HIP (ROCm) 语义#
创建于: 2021年5月12日 | 最后更新于: 2025年8月8日
ROCm™ 是 AMD 的开源软件平台,用于 GPU 加速的高性能计算和机器学习。HIP 是 ROCm 的 C++ 方言,旨在简化 CUDA 应用程序向可移植 C++ 代码的转换。HIP 用于转换现有的 CUDA 应用程序(如 PyTorch)为可移植 C++ 代码,以及用于需要 AMD 和 NVIDIA 之间可移植性的新项目。
HIP 接口复用 CUDA 接口#
PyTorch for HIP 故意复用了现有的 torch.cuda 接口。这有助于加速现有 PyTorch 代码和模型的移植,因为所需的代码更改非常少(如果有的话)。
来自 CUDA 语义 的示例对于 HIP 来说效果完全相同。
cuda = torch.device('cuda') # Default HIP device
cuda0 = torch.device('cuda:0') # 'rocm' or 'hip' are not valid, use 'cuda'
cuda2 = torch.device('cuda:2') # GPU 2 (these are 0-indexed)
x = torch.tensor([1., 2.], device=cuda0)
# x.device is device(type='cuda', index=0)
y = torch.tensor([1., 2.]).cuda()
# y.device is device(type='cuda', index=0)
with torch.cuda.device(1):
# allocates a tensor on GPU 1
a = torch.tensor([1., 2.], device=cuda)
# transfers a tensor from CPU to GPU 1
b = torch.tensor([1., 2.]).cuda()
# a.device and b.device are device(type='cuda', index=1)
# You can also use ``Tensor.to`` to transfer a tensor:
b2 = torch.tensor([1., 2.]).to(device=cuda)
# b.device and b2.device are device(type='cuda', index=1)
c = a + b
# c.device is device(type='cuda', index=1)
z = x + y
# z.device is device(type='cuda', index=0)
# even within a context, you can specify the device
# (or give a GPU index to the .cuda call)
d = torch.randn(2, device=cuda2)
e = torch.randn(2).to(cuda2)
f = torch.randn(2).cuda(cuda2)
# d.device, e.device, and f.device are all device(type='cuda', index=2)
检查 HIP#
无论您使用的是 PyTorch for CUDA 还是 HIP,调用 is_available() 的结果将是相同的。如果您使用的是已构建并支持 GPU 的 PyTorch,它将返回 True。如果您必须检查您使用的是哪个版本的 PyTorch,请参考下面的示例。
if torch.cuda.is_available() and torch.version.hip:
# do something specific for HIP
elif torch.cuda.is_available() and torch.version.cuda:
# do something specific for CUDA
ROCm 上的 TensorFloat-32 (TF32)#
ROCm 不支持 TF32。
内存管理#
PyTorch 使用缓存内存分配器来加速内存分配。这允许快速的内存去分配,而无需设备同步。但是,分配器管理的未使用内存仍会显示为已被使用的状态(在 rocm-smi 中)。您可以使用 memory_allocated() 和 max_memory_allocated() 来监控张量占用的内存,并使用 memory_reserved() 和 max_memory_reserved() 来监控缓存分配器管理的内存总量。调用 empty_cache() 会释放 PyTorch 中所有 **未使用的** 缓存内存,以便其他 GPU 应用程序可以使用。但是,张量占用的 GPU 内存不会被释放,因此无法增加 PyTorch 可用的 GPU 内存量。
对于更高级的用户,我们通过 memory_stats() 提供更全面的内存基准测试。我们还提供通过 memory_snapshot() 捕获内存分配器状态完整快照的功能,这可以帮助您理解您的代码产生的底层分配模式。
要调试内存错误,请在您的环境中设置 PYTORCH_NO_HIP_MEMORY_CACHING=1 来禁用缓存。为了方便移植,也接受 PYTORCH_NO_CUDA_MEMORY_CACHING=1。
hipBLAS 工作空间#
对于每个 hipBLAS 句柄和 HIP 流的组合,如果该句柄和流组合执行了一个需要工作空间的 hipBLAS 内核,则会分配一个 hipBLAS 工作空间。为了避免重复分配工作空间,这些工作空间不会被释放,除非调用 torch._C._cuda_clearCublasWorkspaces();请注意,这是 CUDA 和 HIP 的同一个函数。每个分配的工作空间大小可以通过环境变量 HIPBLAS_WORKSPACE_CONFIG 指定,格式为 :[SIZE]:[COUNT]。例如,环境变量 HIPBLAS_WORKSPACE_CONFIG=:4096:2:16:8 指定了总大小为 2 * 4096 + 8 * 16 KiB,即 8 MiB。默认工作空间大小为 32 MiB;MI300 及更新版本默认为 128 MiB。要强制 hipBLAS 避免使用工作空间,请设置 HIPBLAS_WORKSPACE_CONFIG=:0:0。为了方便起见,也接受 CUBLAS_WORKSPACE_CONFIG。
hipFFT/rocFFT 计划缓存#
不支持设置 hipFFT/rocFFT 计划的缓存大小。
torch.distributed 后端#
目前,ROCm 上只支持 torch.distributed 的 “nccl” 和 “gloo” 后端。
C++ 中的 CUDA API 到 HIP API 映射#
请参考: https://rocm.docs.amd.com/projects/HIP/en/latest/reference/api_syntax.html
注意:CUDA_VERSION 宏、cudaRuntimeGetVersion 和 cudaDriverGetVersion API 在语义上不映射到与 HIP_VERSION 宏、hipRuntimeGetVersion 和 hipDriverGetVersion API 相同的值。在进行版本检查时,请勿互换使用它们。
例如:与其使用
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 来隐式排除 ROCm/HIP,
不如使用以下方法来避免进入 ROCm/HIP 的代码路径:
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 && !defined(USE_ROCM)
或者,如果您希望进入 ROCm/HIP 的代码路径:
#if (defined(CUDA_VERSION) && CUDA_VERSION >= 11000) || defined(USE_ROCM)
或者,如果您希望仅对特定的 HIP 版本进入 ROCm/HIP 的代码路径:
#if (defined(CUDA_VERSION) && CUDA_VERSION >= 11000) || (defined(USE_ROCM) && ROCM_VERSION >= 40300)
参考 CUDA 语义文档#
对于此处未列出的任何部分,请参考 CUDA 语义文档:CUDA 语义。
启用内核断言#
ROCm 支持内核断言,但由于性能开销而禁用。可以通过从源代码重新编译 PyTorch 来启用。
请将以下行作为参数添加到 cmake 命令中:
-DROCM_FORCE_ENABLE_GPU_ASSERTS:BOOL=ON
启用/禁用 ROCm 可组合内核#
启用 SDPA 和 GEMM 的可组合内核 (CK) 是一个两部分的过程。首先,用户必须在构建 pytorch 时将相应的环境变量设置为 '1'。
SDPA: USE_ROCM_CK_SDPA=1
GEMMs: USE_ROCM_CK_GEMM=1
其次,用户必须通过相应的 python 调用显式请求使用 CK 作为后端库。
SDPA: setROCmFAPreferredBackend('<choice>')
GEMMs: setBlasPreferredBackend('<choice>')
要在这两种情况下启用 CK,只需将 'ck' 传递给这些函数即可。
要将后端设置为 CK,用户 **必须** 使用正确的环境变量进行构建。否则,PyTorch 将打印警告并使用“默认”后端。对于 GEMMs,这将路由到 hipblas;对于 SDPA,它将路由到 aotriton。