TunableOp

概述

此模块公开了一个 TunableOp 接口。

某些操作（例如 GEMM）可以使用不止一个库或不止一种技术来实现。例如，可以使用 blas 或 blasLt 库为 CUDA 或 ROCm 实现 GEMM。此外，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 毫秒内可以运行的最大迭代次数（以较小者为准），并计算其平均执行时间。将选择所有成功分析的解决方案中最快的一个。如果给定解决方案未达到与默认实现相同的精度，或者如果解决方案返回错误代码，则性能分析可能会失败。

当前可调优操作符

用于 ROCm 的 TunableGemm

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

离线调优

动机

离线调优有几种用例。

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

另一种用例是计算密集型工作负载。

在这种情况下，更具资源效率的做法是先收集工作负载的 GEMM 一次，然后使用不同的调优参数或库重复进行调优。

工作流程

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

2. 基本上有两步：1) 设置环境变量以收集未调优的 GEMM，这将生成 tunableop_untuned0.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")

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

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)

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

请注意，mgpu_tune_gemm_in_file API 的用法与其单 GPU 对应 API (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 功能是否已启用。

返回类型

bool

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

启用 TunableOp 实现的调优。

启用后，如果未找到已调优条目，则运行调优步骤并记录该条目。

torch.cuda.tunable.is_enabled()

返回 TunableOp 功能是否已启用。

torch.cuda.tunable.tuning_is_enabled()

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

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

启用记录未调优的 TunableOp 操作以进行离线调优。

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

torch.cuda.tunable.is_enabled()

返回 TunableOp 功能是否已启用。

torch.cuda.tunable.record_untuned_is_enabled()

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

torch.cuda.tunable.set_max_tuning_duration(duration)

设置调优给定解决方案的最大时间（毫秒）。

如果同时设置了最大调优持续时间和迭代次数，则以两者中较小者为准。

torch.cuda.tunable.is_enabled()

至少会运行 1 次调优迭代。

torch.cuda.tunable.get_max_tuning_duration()

获取调优给定解决方案的最大时间。

torch.cuda.tunable.set_max_tuning_duration(duration)

int

torch.cuda.tunable.set_max_tuning_iterations(iterations)

torch.cuda.tunable.is_enabled()

至少会运行 1 次调优迭代。

设置调优给定解决方案的最大迭代次数。

torch.cuda.tunable.get_max_tuning_iterations()

获取调优给定解决方案的最大迭代次数。

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

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

torch.cuda.tunable.is_enabled()

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

torch.cuda.tunable.get_filename()

获取结果文件名。

torch.cuda.tunable.is_enabled()

str

torch.cuda.tunable.get_results()

返回所有 TunableOp 结果。

torch.cuda.tunable.is_enabled()

tuple[str, str, str, float]

torch.cuda.tunable.get_validators()

返回 TunableOp 验证器。

tuple[str, str]

torch.cuda.tunable.write_file_on_exit(val)

在销毁调优上下文期间，将文件写入磁盘。

如果您的应用程序因正常运行或错误而终止，这有助于将最终结果刷新到磁盘。可以通过手动调用 write_file() 来手动刷新结果。

torch.cuda.tunable.is_enabled()

返回 TunableOp 功能是否已启用。

Manual flushing of your results can be achieved by manually calling write_file().

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

如果您的应用程序因正常运行或错误而终止，这有助于将最终结果刷新到磁盘。可以通过手动调用 write_file() 来手动刷新结果。

torch.cuda.tunable.is_enabled()

返回 TunableOp 功能是否已启用。

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

获取旋转缓冲区大小（单位：千字节）。

torch.cuda.tunable.is_enabled()

至少会运行 1 次调优迭代。