可调算子

注意

这是一个原型功能，这意味着它处于早期阶段，用于反馈和测试，其组件可能会发生变化。

概述

此模块公开了可调算子接口。

某些操作，例如 GEMM，可以使用多个库或多个技术来实现。例如，可以为 CUDA 或 ROCm 使用 blas 或 blasLt 库来实现 GEMM。此外，ROCm 的 rocblas 和 hipblaslt 库允许用户查询所有可能的算法，然后选择其中一个。如何知道哪个实现最快，应该选择哪个？这就是可调算子提供的功能。

单独启用可调算子并进行调优

可调算子功能与启用调优阶段本身是分开启用的。启用可调算子意味着 PyTorch 将用其可调实现替换任何标准运算符。对可调算子的任何调用都首先检查它是否已经针对给定的运算符输入进行了调优。如果是，它将立即调用调优后的运算；即使在启用调优设置时，也不会进行进一步的调优。相反，如果没有找到调优结果，并且启用了调优，可调算子将为给定的输入集对该运算符的每个已注册实现进行基准测试，并选择最快的实现。

文件输入和输出

首次调用任何可调算子时，将通过尝试从给定文件读取结果来准备已调优运算的内部数据库。默认文件名是 'tunableop_results.csv'。为了支持在多个进程中使用多个 GPU 进行调优，GPU 设备序号会自动插入文件名中，以避免多个进程覆盖同一个文件。

如果启用了调优，并且在工作负载过程中发现新的调优，它还会将所有调优写入同一个文件名，包括启动时读取的调优以及运行时找到的新调优。例如，这可用于通过重用同一个文件跨多个工作负载构建调优文件。应用程序终止时，输出文件将自动创建。此行为可以通过 C++ 和 Python API 控制，但不能通过环境变量控制。

假设您指定了文件名，您最终将得到一个 CSV 文件，其内容如下

Validator,PT_VERSION,2.2.0
Validator,ROCM_VERSION,6.0.0.0-12969-1544e39
Validator,HIPBLASLT_VERSION,0.6.0-a9c5cc7
Validator,ROCBLAS_VERSION,4.0.0-72e57364-dirty
GemmTunableOp_float_NT,nt_25088_4096_64,1219,1.262
GemmTunableOp_float_NT,nt_4096_4096_64,1216,0.033

注意“验证器”行。如果您更改了库版本、ROCm 版本或 PyTorch 版本，可调算子会检测到这一点，并拒绝调优文件，因为之前的调优可能受到其他软件更改的影响。

其余行是您执行过程中遇到的每个可调算子的调优解决方案。每行包含 4 个用逗号分隔的字段：运算符名称、运算符参数、解决方案名称和平均执行时间。执行时间是可选字段。CSV 文件可以编辑，但要谨慎。例如，可以将解决方案名称（字段 3）更改为“默认”，它将回退到原始的未调优 PyTorch 实现。或者，在 ROCm 的 hipBLAS 或 hipBLASLt 库的情况下，如果您知道特定的解决方案索引，您可以通过替换值来覆盖可调算子选择的解决方案。运算符名称和参数（字段 1 和 2）是内部命名的，不应修改。在 GemmTunableOp 的情况下，字段 1 指示数据类型以及输入是转置 (T) 还是否 (N)，字段 2 指示 M、N、K 输入形状。

有一个选项可以启用详细输出，但仅建议用于调试目的。这将生成大量诊断消息，但可能有助于查看是否正在使用可调算子。否则，可调算子将完全静默，除了文件输出之外，除非在其使用过程中存在警告或错误。详细选项仅通过设置环境变量 PYTORCH_TUNABLEOP_VEROBSE=1 可用。

关于调优行为的说明

调优运算符包括遍历已注册实现列表，并分析每个实现。通过在循环中多次运行单个实现并获取平均执行时间来建立分析结果。

默认情况下，给定运算符的每个可能解决方案将运行 100 次迭代或可以在 30 毫秒内运行的迭代次数，以较小者为准，并计算其平均执行结果。在所有已成功分析的解决方案中，将选择最快的解决方案。如果给定解决方案无法达到与默认实现相同的精度，或者如果解决方案返回错误代码，则分析可能会失败。

当前可调运算符

适用于 ROCm 的 TunableGemm

目前仅实现了适用于 ROCm 的 TunableGemm。请注意，PyTorch 的 CUDA 构建在使用可调算子时将正常工作，但 CUDA 构建可用的唯一解决方案是“默认”实现，即原始 cuBLAS 默认值，现在通过可调算子调用。对 at::cuda::blas::gemm() 或 ::bgemm() 的任何调用都将在启用时通过可调算子路由。调用针对给定的输入参数集 (transa, transb, m, n, k) 的 gemm() 将尝试在 rocblas 和 hipblaslt 中使用最快的可用实现。

调优上下文

TunableOp 的行为目前可以通过环境变量、at::cuda::tunable::getTuningContext() 的 C++ 接口或封装 C++ TuningContext 的 torch.cuda.tunable python 接口进行操作。环境变量优先于您使用 C++ 或 Python API 操作的任何设置。

API 参考

torch.cuda.tunable.enable(val=True)

这是所有 TunableOp 实现的总开关。

torch.cuda.tunable.is_enabled()

返回 TunableOp 功能是否启用。

返回类型

bool

torch.cuda.tunable.tuning_enable(val=True)

启用 TunableOp 实现的调整。

启用后，如果找不到调整的条目，则运行调整步骤并记录条目。

torch.cuda.tunable.tuning_is_enabled()

返回 TunableOp 实现是否可以调整。

返回类型

bool

torch.cuda.tunable.set_max_tuning_duration(duration)

设置花费在调整给定解决方案上的最大时间（以毫秒为单位）。

如果同时设置了最大调整持续时间和迭代次数，则将使用较小的值。始终至少运行 1 次调整迭代。

torch.cuda.tunable.get_max_tuning_duration()

获取花费在调整给定解决方案上的最大时间。

返回类型

int

torch.cuda.tunable.set_max_tuning_iterations(iterations)

设置花费在调整给定解决方案上的最大迭代次数。

如果同时设置了最大调整持续时间和迭代次数，则将使用较小的值。始终至少运行 1 次调整迭代。

torch.cuda.tunable.get_max_tuning_iterations()

获取花费在调整给定解决方案上的最大迭代次数。

返回类型

int

torch.cuda.tunable.set_filename(filename, insert_device_ordinal=False)

设置用于调整结果的输入/输出的文件名。

如果 insert_device_ordinal 为 True，则当前设备序号将自动添加到给定的文件名中。这可以在每个 GPU 一个进程的场景中使用，以确保所有进程都写入不同的文件。

torch.cuda.tunable.get_filename()

获取结果文件名。

返回类型

str

torch.cuda.tunable.get_results()

返回所有 TunableOp 结果。

返回类型

Tuple[str, str, str, float]

torch.cuda.tunable.get_validators()

返回 TunableOp 验证器。

返回类型

Tuple[str, str]

torch.cuda.tunable.write_file_on_exit(val)

在 Tuning Context 销毁期间，将文件写入磁盘。

这在您的应用程序因正常操作或错误而终止时，作为对磁盘结果的最终刷新非常有用。可以通过手动调用 write_file() 来手动刷新您的结果。

torch.cuda.tunable.write_file(filename=None)

将结果写入 CSV 文件。

如果未给出 filename，则调用 get_filename()。

返回类型

bool

torch.cuda.tunable.read_file(filename=None)

从 TunableOp CSV 文件中读取结果。

如果未给出 filename，则调用 get_filename()。

返回类型

bool