TunableOp#
创建于: 2024年6月3日 | 最后更新于: 2025年6月13日
概述#
此模块公开了 TunableOp 接口。
某些操作，例如 GEMMs，可以使用多个库或多个技术来实现。例如，GEMM 可以为 CUDA 或 ROCm 实现，使用 blas 或 blasLt 库。此外，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
注意“验证器”行。如果您更改了库版本、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 次迭代，或者运行 30ms 内可以运行的尽可能多的迭代，取两者中的较小值，并计算其平均执行时间。在所有成功进行性能分析的解决方案中，将选择最快的解决方案。如果给定的解决方案未达到与默认实现相同的精度或解决方案返回了错误代码，性能分析可能会失败。
当前可调优的操作符#
ROCm 的 TunableGemm#
目前只实现了一个 ROCm 的 TunableGemm。请注意，PyTorch 的 CUDA 构建在使用 TunableOp 时可以正常工作，但对 CUDA 构建可用的唯一解决方案是“Default”实现，即通过 TunableOp 调用的原始 cuBLAS 默认实现。任何对 at::cuda::blas::gemm() 或 ::bgemm() 的调用在启用时都将通过 TunableOp 进行路由。调用具有给定输入参数集（transa、transb、m、n、k）的 gemm() 将尝试使用 rocblas 和 hipblaslt 中可用的最快实现。
离线调优#
动机#
离线调优有几种用例。
一种用例涉及高内存利用率的工作负载，此时常规调优可能导致内存不足。
另一种用例是计算密集型工作负载。在这种情况下，一次收集工作负载的 GEMMs，然后用不同的调优参数或库重复调优，会更节省资源。
工作流程#
基本上有两个步骤：1) 设置环境变量以收集未调优的 GEMM，这将生成 tunableop_untuned0.csv
export PYTORCH_TUNABLEOP_ENABLED=1
export PYTORCH_TUNABLEOP_TUNING=0
export 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")
也可以同时处理多个未调优文件，并将 GEMMs 分布到单个节点中的多个 GPU 上进行调优。第一步，首先收集 GEMMs 并消除重复的 GEMMs。接下来，将 GEMMs 分发到不同的 GPU 上进行调优。在所有 GEMMs 都调优完成后，来自所有 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）不同。调用该 API 的 Python 脚本的主体必须如所示用 main() 包裹起来，因为使用了 concurrent.futures 模块。传递给 mgpu_tune_gemm_in_file 的参数必须包含一个通配符表达式（? 或 *）来生成要处理的 GEMMs 的未调优文件列表。num_gpus 必须介于 1 和可用 GPU 总数之间。
API 参考#
- torch.cuda.tunable.tuning_enable(val=True)[source]#
启用 TunableOp 实现的调优。
启用后，如果找不到已调优的条目，将运行调优步骤并记录该条目。
- torch.cuda.tunable.record_untuned_enable(val=True)[source]#
为离线调优启用 TunableOp 操作的未调优记录。
启用后，如果找不到已调优的条目，则将其写入未调优文件。
- torch.cuda.tunable.set_max_tuning_duration(duration)[source]#
设置调优给定解决方案的最大时间（毫秒）。
如果同时设置了最大调优时间和迭代次数，将遵守两者中的较小值。至少会始终运行 1 次调优迭代。
- torch.cuda.tunable.set_max_tuning_iterations(iterations)[source]#
设置调优给定解决方案的最大迭代次数。
如果同时设置了最大调优时间和迭代次数，将遵守两者中的较小值。至少会始终运行 1 次调优迭代。
- torch.cuda.tunable.set_filename(filename, insert_device_ordinal=False)[source]#
设置用于输入/输出调优结果的文件名。
如果
insert_device_ordinal为True，则当前设备序数将自动添加到给定的文件名中。这可以在每个 GPU 一个进程的场景中使用，以确保所有进程写入单独的文件。
- torch.cuda.tunable.write_file_on_exit(val)[source]#
在上下文销毁时，将文件写入磁盘。
这有助于在应用程序因正常操作或错误而终止时将结果最终刷新到磁盘。可以通过手动调用
write_file()来手动刷新结果。
- torch.cuda.tunable.write_file(filename=None)[source]#
将结果写入 CSV 文件。
如果未给出
filename，则调用get_filename()。- 返回类型
- torch.cuda.tunable.read_file(filename=None)[source]#
从 TunableOp CSV 文件读取结果。
如果未给出
filename，则调用get_filename()。- 返回类型
- torch.cuda.tunable.mgpu_tune_gemm_in_file(filename_pattern, num_gpus)[source]#
处理一个或多个文件并将工作分配到一或多个 GPU 上。
