TunableOp#

创建于: 2024年6月3日 | 最后更新于: 2025年6月13日

概述#

本模块暴露了 TunableOp 接口。

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

单独启用 TunableOp 和调优#

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

文件输入和输出#

第一次调用任何 TunableOp 时，将通过尝试从给定文件读取结果来准备调优后操作的内部数据库。默认文件名是 '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,Gemm_Hipblaslt_1219,1.262
GemmTunableOp_float_NT,nt_4096_4096_64,Gemm_Rocblas_1216,0.033

注意“Validator”行。如果您更改了库版本、ROCm 版本或 PyTorch 版本，TunableOp 将检测到这一点并拒绝该调优文件，因为先前的调优可能受到其他软件更改的影响。

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

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

关于调优行为、预热和缓存效应的说明#

调优一个运算符包括遍历已注册的实现列表并分析每个实现。分析是通过将单个实现在一个循环中运行多次并取平均执行时间来建立的。在调优之前还有一个可选的预热阶段，可以帮助硬件达到稳定的电源状态。在调优工作负载期间，各种硬件缓存更有可能产生命中，而不是在不调优时。有用于刷新指令缓存和轮换输入张量的选项，这可能有助于更真实地分析调优后的运算符，就像运算符在更大的工作负载中运行一样，而不是在一个紧凑、重复的循环中。

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

当前可调优运算符#

TunableGemm for ROCm#

目前仅实现了一个 TunableGemm for ROCm。请注意，PyTorch 的 CUDA 构建在使用 TunableOp 时也能正常工作，但仅适用于 CUDA 构建的解决方案是 'Default' 实现，即通过 TunableOp 调用原始的 cuBLAS 默认实现。任何对 at::cuda::blas::gemm() 或 ::bgemm() 的调用都将在启用时通过 TunableOp 进行路由。调用 gemm() for a given set of input arguments (transa, transb, m, n, k) 将尝试使用 rocblas 和 hipblaslt 中最快的可用实现。

离线调优#

动机#

离线调优有几种用例。

一种用例涉及内存利用率很高的工作负载，此时常规调优可能导致内存不足。

另一种用例是针对计算密集型工作负载。在这种情况下，一次收集工作负载的 GEMM，然后使用不同的调优参数或库进行重复调优，会更有效率。

工作流#

基本上有两种步骤：1) 设置环境变量以收集未调优的 GEMM，这将生成 tunableop_untuned0.csv

PYTORCH_TUNABLEOP_ENABLED=1
PYTORCH_TUNABLEOP_TUNING=0
PYTORCH_TUNABLEOP_RECORD_UNTUNED=1
...

运行一个 Python 脚本，该脚本读取 tunableop_untuned0.csv 并生成 tunableop_results0.csv，如下所示

import torch.cuda.tunable as tunable
import os

os.putenv('PYTORCH_TUNABLEOP_ENABLED', '1')
os.putenv('PYTORCH_TUNABLEOP_TUNING', '1')
os.putenv('PYTORCH_TUNABLEOP_RECORD_UNTUNED', '0')
tunable.tune_gemm_in_file("tunableop_untuned0.csv")

还可以将多个未调优文件组合起来，并将 GEMM 分发到单个节点内的多个 GPU 上进行调优。在第一步中，首先收集 GEMM 并消除重复的 GEMM。接下来，将 GEMM 分发到不同的 GPU 上进行调优。在所有 GEMM 都调优完成后，来自所有 GPU 的结果将收集到一个文件中，该文件的基本文件名会附加 _full0（例如 tunableop_results_full0.csv）。最后，这个包含收集结果的新文件将被复制 N 次，为每个 GPU 复制一份，方便用户在 N 个 GPU 上使用调优后的配置运行工作负载。

if __name__ == "__main__":
    num_gpus = 8 # number of GPUs that will be used during the tuning process
    tunable.mgpu_tune_gemm_in_file("tunableop_untuned?.csv", num_gpus)

请注意，mgpu_tune_gemm_in_file API 的使用方式与其单 GPU 版本（tune_gemm_in_file）不同。由于使用了 concurrent futures 模块，调用该 API 的 Python 脚本主体必须包装在 main() 中，如所示。传递给 mgpu_tune_gemm_in_file 的参数必须包含一个通配符表达式（? 或 *），以生成包含要处理的 GEMM 的未调优文件列表。num_gpus 必须介于 1 和可用 GPU 总数之间。

调优上下文#

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

环境变量接口#

环境变量在首次读取时被缓存。您无法以编程方式使用环境变量接口，因为设置一旦生效就固定不变。请改用 C++ 或 Python API。

API 参考#

torch.cuda.tunable.enable(val=True)[源代码]#

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

torch.cuda.tunable.is_enabled()[源代码]#

返回 TunableOp 功能是否已启用。

返回类型: 布尔值

torch.cuda.tunable.tuning_enable(val=True)[源代码]#

启用 TunableOp 实现的调优。

启用后，如果找不到调优的条目，将执行调优步骤并记录该条目。

torch.cuda.tunable.tuning_is_enabled()[源代码]#

返回 TunableOp 实现是否可以进行调优。

返回类型: 布尔值

torch.cuda.tunable.record_untuned_enable(val=True)[源代码]#

启用 TunableOp 操作的未调优记录，用于离线调优。

启用后，如果找不到调优的条目，将将其写入未调优文件。

torch.cuda.tunable.record_untuned_is_enabled()[源代码]#

返回 TunableOp 操作是否已为离线调优进行记录。

返回类型: 布尔值

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()。

返回类型: 布尔值

torch.cuda.tunable.read_file(filename=None)[源代码]#

从 TunableOp CSV 文件读取结果。

如果未提供 filename，则会调用 get_filename()。

返回类型: 布尔值

torch.cuda.tunable.tune_gemm_in_file(filename)[源代码]#

在文件中调优 GEMM。

torch.cuda.tunable.mgpu_tune_gemm_in_file(filename_pattern, num_gpus)[源代码]#

处理一个或多个文件，并将工作分发到一台或多台 GPU 上。

torch.cuda.tunable.set_rotating_buffer_size(buffer_size)[源代码]#

如果缓冲大小大于零，则将其设置为此值（以 MB 为单位）。

如果小于零，则查询 L2 缓存大小。如果等于零，表示禁用旋转缓冲区。

torch.cuda.tunable.get_rotating_buffer_size()[源代码]#

获取旋转缓冲区大小（以 KB 为单位）。

返回类型: int