CUDA 语义

创建日期：2017 年 1 月 16 日 | 最后更新日期：2026 年 1 月 15 日

torch.cuda 用于设置和运行 CUDA 操作。它跟踪当前选定的 GPU，并且您分配的所有 CUDA 张量默认将在该设备上创建。可以使用 torch.cuda.device 上下文管理器更改选定的设备。

但是，一旦分配了张量，无论选定的设备如何，您都可以对其执行操作，并且结果始终放置在与张量相同的设备上。

默认情况下不允许跨 GPU 操作，但 copy_() 以及其他具有类似复制功能的函数（例如 to() 和 cuda()）除外。除非您启用点对点内存访问，否则任何尝试在跨不同设备的张量上启动操作都将引发错误。

下面您可以看到一个展示此功能的示例

cuda = torch.device('cuda')     # Default CUDA device
cuda0 = torch.device('cuda:0')
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)

Ampere（及更高版本）设备上的 TensorFloat-32 (TF32)

在 Pytorch 2.9 之后，我们提供了一组新的 API，以便以更细粒度的方式控制 TF32 行为，并建议使用新的 API 以获得更好的控制。我们可以为每个后端和每个运算符设置 float32 精度。我们还可以覆盖特定运算符的全局设置。

torch.backends.fp32_precision = "ieee"
torch.backends.cuda.matmul.fp32_precision = "ieee"
torch.backends.cudnn.fp32_precision = "ieee"
torch.backends.cudnn.conv.fp32_precision = "tf32"
torch.backends.cudnn.rnn.fp32_precision = "tf32"

对于 cuda/cudnn，可以将 fp32_precision 设置为 ieee 或 tf32。ieee fp32_precision 表示我们将使用 FP32 作为内部计算精度。tf32 fp32_precision 表示我们将允许使用 TF32 作为内部计算精度。

如果 fp32_precision 设置为 ieee，我们可以覆盖特定运算符的通用设置。

torch.backends.cudnn.fp32_precision = "tf32"
torch.backends.cudnn.conv.fp32_precision = "ieee"
torch.backends.cudnn.rnn.fp32_precision = "ieee"

如果 fp32_precision 设置为 ieee，我们可以覆盖特定后端的通用设置。

torch.backends.fp32_precision = "tf32"
torch.backends.cudnn.fp32_precision = "ieee"
torch.backends.cudnn.conv.fp32_precision = "ieee"
torch.backends.cudnn.rnn.fp32_precision = "ieee"

对于上述两种情况，torch.backends.cudnn.conv.fp32_precision 和 torch.backends.cudnn.rnn.fp32_precision 均被覆盖为 ieee。

我们建议使用新的设置以获得更好的控制。并且我们不支持同时使用新旧设置。

警告

以下旧设置与 allow_tf32 相关，将被弃用。我们建议使用上述新设置以获得更好的控制。并且我们不支持同时使用新旧设置。

从 PyTorch 1.7 开始，有一个名为 allow_tf32 的新标志。此标志在 PyTorch 1.7 到 PyTorch 1.11 中默认为 True，在 PyTorch 1.12 及更高版本中默认为 False。此标志控制 PyTorch 是否允许在内部使用 NVIDIA GPU 上提供的 TensorFloat32 (TF32) 张量核心来计算 matmul（矩阵乘法和批量矩阵乘法）和卷积。

TF32 张量核心旨在通过将输入数据四舍五入到具有 10 位尾数，并以 FP32 精度累积结果（保持 FP32 动态范围），从而在 torch.float32 张量上的 matmul 和卷积上实现更好的性能。

matmul 和卷积分别受控，并且可以访问其相应的标志：

# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
# in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True

# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True

matmul 的精度也可以通过 set_float32_matmul_precision() 更广泛地设置（不限于 CUDA）。请注意，除了 matmul 和卷积本身之外，内部使用 matmul 或卷积的函数和 nn 模块也会受到影响。这些包括 nn.Linear、nn.Conv*、cdist、tensordot、仿射网格和网格采样、自适应 log softmax、GRU 和 LSTM。

为了了解精度和速度，请参阅下面的示例代码和基准数据（在 A100 上）

a_full = torch.randn(10240, 10240, dtype=torch.double, device='cuda')
b_full = torch.randn(10240, 10240, dtype=torch.double, device='cuda')
ab_full = a_full @ b_full
mean = ab_full.abs().mean()  # 80.7277

a = a_full.float()
b = b_full.float()

# Do matmul at TF32 mode.
torch.backends.cuda.matmul.allow_tf32 = True
ab_tf32 = a @ b  # takes 0.016s on GA100
error = (ab_tf32 - ab_full).abs().max()  # 0.1747
relative_error = error / mean  # 0.0022

# Do matmul with TF32 disabled.
torch.backends.cuda.matmul.allow_tf32 = False
ab_fp32 = a @ b  # takes 0.11s on GA100
error = (ab_fp32 - ab_full).abs().max()  # 0.0031
relative_error = error / mean  # 0.000039

从上面的示例可以看出，启用 TF32 后，A100 上的速度大约快 7 倍，与双精度相比，相对误差大约大两个数量级。请注意，TF32 与单精度速度的精确比率取决于硬件代、内存带宽与计算的比率以及 TF32 与 FP32 matmul 吞吐量的比率等属性，这些属性可能因代或模型而异。如果需要完全 FP32 精度，用户可以通过以下方式禁用 TF32：

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

要在 C++ 中关闭 TF32 标志，您可以执行以下操作：

at::globalContext().setAllowTF32CuBLAS(false);
at::globalContext().setAllowTF32CuDNN(false);

有关 TF32 的更多信息，请参阅：

• TensorFloat-32

• CUDA 11

• Ampere 架构

FP16 GEMM 中降低精度的缩减

（与旨在用于具有比 FP32 累积更高的吞吐量的硬件的完全 FP16 累积不同，请参阅 完全 FP16 累积）

fp16 GEMM 可能会以一些中间降低精度的缩减（例如，在 fp16 而不是 fp32 中）完成。这些精度选择性缩减可以允许在某些工作负载（特别是具有大 k 维度的工作负载）和 GPU 架构上获得更高的性能，但代价是数值精度和溢出的可能性。

V100 上的示例基准数据：

[--------------------------- bench_gemm_transformer --------------------------]
      [  m ,  k  ,  n  ]    |  allow_fp16_reduc=True  |  allow_fp16_reduc=False
1 threads: --------------------------------------------------------------------
      [4096, 4048, 4096]    |           1634.6        |           1639.8
      [4096, 4056, 4096]    |           1670.8        |           1661.9
      [4096, 4080, 4096]    |           1664.2        |           1658.3
      [4096, 4096, 4096]    |           1639.4        |           1651.0
      [4096, 4104, 4096]    |           1677.4        |           1674.9
      [4096, 4128, 4096]    |           1655.7        |           1646.0
      [4096, 4144, 4096]    |           1796.8        |           2519.6
      [4096, 5096, 4096]    |           2094.6        |           3190.0
      [4096, 5104, 4096]    |           2144.0        |           2663.5
      [4096, 5112, 4096]    |           2149.1        |           2766.9
      [4096, 5120, 4096]    |           2142.8        |           2631.0
      [4096, 9728, 4096]    |           3875.1        |           5779.8
      [4096, 16384, 4096]   |           6182.9        |           9656.5
(times in microseconds).

如果需要完全精度缩减，用户可以使用以下方法禁用 fp16 GEMM 中的降低精度缩减：

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

要在 C++ 中切换降低精度缩减标志，可以执行以下操作：

at::globalContext().setAllowFP16ReductionCuBLAS(false);

BF16 GEMM 中降低精度的缩减

存在一个与上述类似的标志，用于 BFloat16 GEMM。请注意，如果默认情况下此开关设置为 True，如果您在工作负载中观察到数值不稳定，您可能希望将其设置为 False。

如果不需要降低精度缩减，用户可以使用以下方法禁用 bf16 GEMM 中的降低精度缩减：

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False

要在 C++ 中切换降低精度缩减标志，可以执行以下操作：

at::globalContext().setAllowBF16ReductionCuBLAS(true);

FP16 GEMM 中的完全 FP16 累积

某些 GPU 在以 FP16 进行 _所有_ FP16 GEMM 累积时具有更高的性能，代价是数值精度和溢出的可能性增加。请注意，此设置仅对计算能力为 7.0（Volta）或更高版本的 GPU 有效。

可以通过以下方式启用此行为：

torch.backends.cuda.matmul.allow_fp16_accumulation = True

要在 C++ 中切换降低精度缩减标志，可以执行以下操作：

at::globalContext().setAllowFP16AccumulationCuBLAS(true);

异步执行

默认情况下，GPU 操作是异步的。当您调用使用 GPU 的函数时，操作会排队到特定设备，但不一定立即执行。这允许我们并行执行更多计算，包括在 CPU 或其他 GPU 上的操作。

通常，异步计算的效果对调用者不可见，因为 (1) 每个设备按创建的顺序执行流中的操作，并且 (2) PyTorch 会自动执行必要地同步，以便在 CPU 和 GPU 之间或两个 GPU 之间复制数据。因此，计算将像每个操作都同步执行一样进行。

可以通过设置环境变量 CUDA_LAUNCH_BLOCKING=1 来强制同步计算。当 GPU 上发生错误时，这会很有用。（使用异步执行，直到实际执行操作后才会报告此类错误，因此堆栈跟踪不会显示请求发生的位置。）

异步计算的一个后果是，没有同步的计时测量是不准确的。要获得精确的测量结果，应该调用 torch.cuda.synchronize() 进行测量之前，或者使用 torch.cuda.Event 记录时间，如下所示：

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()

# Run some things here

end_event.record()
torch.cuda.synchronize()  # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)

作为例外，一些函数（例如 to() 和 copy_()）都允许使用显式的 non_blocking 参数，该参数允许调用者在不需要时绕过同步。另一个例外是 CUDA 流，如下所述。

CUDA 流

一个 CUDA 流 是属于特定设备的执行的线性序列。通常，您不需要显式创建一个：默认情况下，每个设备使用自己的“默认”流。

流中的操作按创建的顺序序列化，但来自不同流的操作可以以任何相对顺序并发执行，除非使用显式同步函数（例如 synchronize() 或 wait_stream()）。例如，以下代码是不正确的：

cuda = torch.device('cuda')
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
with torch.cuda.stream(s):
    # sum() may start execution before normal_() finishes!
    B = torch.sum(A)

当“当前流”是默认流时，PyTorch 会自动执行必要的同步，以便在数据移动时进行同步，如上所述。但是，当使用非默认流时，确保正确同步是用户的责任。此示例的修复版本是：

cuda = torch.device('cuda')
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
s.wait_stream(torch.cuda.default_stream(cuda))  # NEW!
with torch.cuda.stream(s):
    B = torch.sum(A)
A.record_stream(s)  # NEW!

有两个新添加项。torch.cuda.Stream.wait_stream() 调用确保 normal_() 执行完成后，我们才开始在侧流上运行 sum(A)。torch.Tensor.record_stream()（有关更多详细信息，请参阅）确保我们在 sum(A) 完成之前不会释放 A 的内存。您也可以稍后使用 torch.cuda.default_stream(cuda).wait_stream(s) 在流上手动等待（请注意，立即等待是没有意义的，因为这将阻止流执行与默认流上的其他工作并行运行）。有关更多详细信息，请参阅 torch.Tensor.record_stream() 的文档，了解何时使用一个或另一个。

请注意，即使没有读取依赖关系，例如在本例中所示，也需要进行此同步：

cuda = torch.device('cuda')
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda)
s.wait_stream(torch.cuda.default_stream(cuda))  # STILL REQUIRED!
with torch.cuda.stream(s):
    A.normal_(0.0, 1.0)
    A.record_stream(s)

尽管 s 上的计算不会读取 A 的内容，并且没有 A 的其他用途，但仍然需要同步，因为 A 可能对应于 CUDA 缓存分配器重新分配的内存，其中包含旧（已释放）内存中的待处理操作。

反向传播的流语义

每个反向 CUDA 操作都在用于其对应前向操作的相同流上运行。如果你的前向传递在不同的流上并行运行独立的运算，这将有助于反向传递利用相同的并行性。

反向调用相对于周围运算的流语义与任何其他调用相同。反向传递会插入内部同步以确保这一点，即使反向运算如前一段所述在多个流上运行。更具体地说，当调用 autograd.backward、autograd.grad 或 tensor.backward 时，并且可选地将 CUDA 张量作为初始梯度提供（例如，autograd.backward(..., grad_tensors=initial_grads)、autograd.grad(..., grad_outputs=initial_grads) 或 tensor.backward(..., gradient=initial_grad)），那么

1. 可选地填充初始梯度，

2. 调用反向传递，以及

3. 使用梯度

具有与任何运算组相同的流语义关系

s = torch.cuda.Stream()

# Safe, grads are used in the same stream context as backward()
with torch.cuda.stream(s):
    loss.backward()
    use grads

# Unsafe
with torch.cuda.stream(s):
    loss.backward()
use grads

# Safe, with synchronization
with torch.cuda.stream(s):
    loss.backward()
torch.cuda.current_stream().wait_stream(s)
use grads

# Safe, populating initial grad and invoking backward are in the same stream context
with torch.cuda.stream(s):
    loss.backward(gradient=torch.ones_like(loss))

# Unsafe, populating initial_grad and invoking backward are in different stream contexts,
# without synchronization
initial_grad = torch.ones_like(loss)
with torch.cuda.stream(s):
    loss.backward(gradient=initial_grad)

# Safe, with synchronization
initial_grad = torch.ones_like(loss)
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    initial_grad.record_stream(s)
    loss.backward(gradient=initial_grad)

BC 注意：在默认流上使用梯度

在 PyTorch 的早期版本（1.9 及更早版本）中，autograd 引擎始终将默认流与所有反向运算同步，因此只要 use grads 在默认流上发生，以下模式

with torch.cuda.stream(s):
    loss.backward()
use grads

都是安全的。在当前的 PyTorch 中，该模式不再安全。如果 backward() 和 use grads 在不同的流上下文中，你必须同步这些流

with torch.cuda.stream(s):
    loss.backward()
torch.cuda.current_stream().wait_stream(s)
use grads

即使 use grads 在默认流上也是如此。

内存管理

PyTorch 使用缓存内存分配器来加速内存分配。这允许快速内存释放，而无需设备同步。但是，分配器管理的未用内存仍然会显示为已用内存，如在 nvidia-smi 中所示。你可以使用 memory_allocated() 和 max_memory_allocated() 来监控张量占用的内存，并使用 memory_reserved() 和 max_memory_reserved() 来监控缓存分配器管理的总内存量。调用 empty_cache() 会释放 PyTorch 中所有**未用**的缓存内存，以便其他 GPU 应用程序可以使用这些内存。但是，张量占用的 GPU 内存不会被释放，因此无法增加 PyTorch 可用的 GPU 内存量。

为了更好地了解 CUDA 内存随时间的使用情况，了解 CUDA 内存使用情况 描述了用于捕获和可视化内存使用情况跟踪的工具。

对于更高级的用户，我们通过 memory_stats() 提供更全面的内存基准测试。我们还提供通过 memory_snapshot() 捕获内存分配器状态的完整快照的功能，这可以帮助你了解代码产生的底层分配模式。

使用 PYTORCH_ALLOC_CONF 优化内存使用

使用缓存分配器可能会干扰内存检查工具，例如 cuda-memcheck。要使用 cuda-memcheck 调试内存错误，请在你的环境中设置 PYTORCH_NO_CUDA_MEMORY_CACHING=1 以禁用缓存。

缓存分配器的行为可以通过环境变量 PYTORCH_ALLOC_CONF 进行控制。PYTORCH_CUDA_ALLOC_CONF 是它的别名，仅为了向后兼容而提供。格式为 PYTORCH_ALLOC_CONF=<option>:<value>,<option2>:<value2>... 可用选项

• backend 允许选择底层的分配器实现。当前，有效选项是 native，它使用 PyTorch 的本机实现，以及 cudaMallocAsync，它使用 CUDA 内置的异步分配器。cudaMallocAsync 需要 CUDA 11.4 或更高版本。默认值为 native。backend 适用于进程使用的所有设备，并且不能按设备指定。

• large_segment_size_mb 本机分配器使用小块和大块来管理分配的内存。此设置用于配置大块的大小。默认值为 20 MB。

• max_split_size_mb 防止本机分配器拆分大于此大小（以 MB 为单位）的块。这可以减少碎片，并且可能允许一些临界工作负载完成而不会耗尽内存。性能成本范围从“零”到“巨大”，具体取决于分配模式。默认值为无限制，即可以拆分所有块。memory_stats() 和 memory_summary() 方法对于调整很有用。应该将此选项作为最后的手段用于由于“内存不足”而中止并显示大量非活动拆分块的工作负载。max_split_size_mb 仅在 backend:native 有意义时才有意义。对于 backend:cudaMallocAsync，max_split_size_mb 将被忽略。

• roundup_power2_divisions 有助于将请求的分配大小四舍五入到最接近的 2 的幂次方除法，并更好地利用块。在 native CUDACachingAllocator 中，大小以 512 的块大小为倍数进行四舍五入，因此这对于较小的大小来说很好。但是，这对于大型附近分配来说效率可能较低，因为每个分配都会进入不同大小的块，并且这些块的重用被最小化。这可能会创建许多未使用的块并浪费 GPU 内存容量。此选项启用将分配大小四舍五入到最接近的 2 的幂次方除法。例如，如果我们需要四舍五入大小为 1200，并且除数数量为 4，则该大小位于 1024 和 2048 之间，如果我们在这两者之间进行 4 次除法，则值为 1024、1280、1536 和 1792。因此，分配大小 1200 将四舍五入到 1280 作为最接近的上限的 2 的幂次方除法。指定单个值以应用于所有分配大小，或指定键值对数组以针对每个 2 的幂次方区间单独设置 2 的幂次方除法。例如，要为所有小于 256MB 的分配设置 1 次除法，为 256MB 到 512MB 之间的分配设置 2 次除法，为 512MB 到 1GB 之间的分配设置 4 次除法，为任何更大的分配设置 8 次除法，请将旋钮值设置为：[256:1,512:2,1024:4,>:8]。roundup_power2_divisions 仅在 backend:native 有意义时才有意义。对于 backend:cudaMallocAsync，roundup_power2_divisions 将被忽略。

• max_non_split_rounding_mb 将允许不拆分块以更好地重用，例如，

  一个 1024MB 缓存块可以被用于一个 512MB 分配请求。在默认情况下，我们只允许最多 20MB 的不拆分块四舍五入，因此一个 512MB 块只能用 512-532 MB 大小的块提供。如果我们将此选项的值设置为 1024，它将允许 512-1536 MB 大小的块用于一个 512MB 块，从而提高较大块的重用。这也有助于减少避免昂贵的 cudaMalloc 调用中的停顿。

• garbage_collection_threshold 有助于主动回收未使用的 GPU 内存，以避免触发昂贵的同步和回收所有操作（release_cached_blocks），这可能对延迟敏感的 GPU 应用程序（例如服务器）不利。在设置此阈值（例如 0.8）后，如果 GPU 内存容量使用量超过阈值（即 GPU 应用程序分配的总内存的 80%），分配器将开始回收 GPU 内存块。该算法首先释放旧的和未使用的块，以避免释放正在积极重用的块。阈值应大于 0.0 且小于 1.0。默认值为 1.0。

  garbage_collection_threshold 仅在 backend:native 有意义时才有意义。对于 backend:cudaMallocAsync，garbage_collection_threshold 将被忽略。

• expandable_segments (实验性，默认：False) 如果设置为 True，此设置指示分配器创建可以随后扩展的 CUDA 分配，以便更好地处理作业频繁更改分配大小的情况，例如批处理大小不断变化。通常，对于大型（>2MB）分配，分配器会调用 cudaMalloc 获取与用户请求大小相同的分配。将来，这些分配的部分可以重用于其他请求（如果它们是空闲的）。当程序发出许多精确相同大小或大小为该大小倍数的请求时，这会很好地工作。许多深度学习模型遵循这种行为。但是，一个常见的例外是当批处理大小从一个迭代到下一个迭代略有变化时，例如在批量推理中。当程序最初以批处理大小 N 运行时，它将为该大小进行适当的分配。如果将来它以大小 N - 1 运行，现有的分配仍然足够大。但是，如果它以大小 N + 1 运行，那么它将不得不进行新的分配，这些分配略大一些。并非所有张量的大小都相同。有些可能是 (N + 1)*A，而另一些可能是 (N + 1)*A*B，其中 A 和 B 是模型中的一些非批处理维度。由于分配器在现有分配足够大时会重用现有分配，因此一些 (N + 1)*A 分配实际上可以放入已经存在的 N*B*A 分段中，尽管不完美。随着模型的运行，它将部分填充所有这些分段，在这些分段的末尾留下不可用的空闲内存切片。分配器在某个时候需要 cudaMalloc 一个新的 (N + 1)*A*B 分段。如果没有足够的内存，现在无法恢复现有分段末尾的空闲内存切片。对于 50 多层的模型，这种模式可能会重复 50 多次，从而创建许多细小的碎片。

  expandable_segments 允许分配器最初创建一个分段，然后在需要更多内存时以后扩展其大小。它不是为每个分配创建一个分段，而是尝试为每个流创建一个随着需要增长的分段。现在，当 N + 1 案例运行时，分配将很好地平铺到单个大型分段中，直到它被填满。然后请求更多内存并附加到分段的末尾。此过程不会创建那么多不可用的内存碎片，因此更有可能找到该内存。

• pinned_use_cuda_host_register 选项是一个布尔标志，用于确定是否使用 CUDA API 的 cudaHostRegister 函数来分配固定内存，而不是默认的 cudaHostAlloc。当设置为 True 时，内存使用常规 malloc 分配，然后在调用 cudaHostRegister 之前将页面映射到内存。对页面进行预映射有助于减少 cudaHostRegister 执行期间的锁定时间。

• pinned_num_register_threads 选项仅在 pinned_use_cuda_host_register 设置为 True 时有效。默认情况下，使用一个线程来映射页面。此选项允许使用更多线程来并行化页面映射操作，以减少固定内存的总体分配时间。根据基准测试结果，此选项的一个良好值是 8。

• pinned_use_background_threads 选项是一个布尔标志，用于启用后台线程来处理事件。这避免了与在快速分配路径中查询/处理事件相关的任何缓慢路径。此功能默认禁用。

• pinned_reserve_segment_size_mb 选项是为固定内存分段预留的大小（MB）。这会预先分配大量固定内存，然后用于分配小尺寸请求。这有助于减少设备库调用的次数。

• graph_capture_record_stream_reuse (实验性，默认：False) 如果设置为 True，CUDA 缓存分配器将尝试通过使用图拓扑（而不是 CUDA 事件）来确定何时可以安全地重用释放的块，从而在 CUDA 图捕获期间回收设备内存。这可以减少长时间捕获的峰值内存，这些捕获在多个流中释放和重新分配缓冲区，尤其是在捕获 DAG 频繁到达连接前沿时。

• per_process_memory_fraction 选项限制了可以在所有 CUDA 设备上分配的内存量，达到可用内存的指定比例。这是一个介于 0 和 1 之间的值。尝试分配更多内存将引发内存不足错误。

注意

由 CUDA 内存管理 API 报告的一些统计信息特定于 backend:native，并且不适用于 backend:cudaMallocAsync。有关详细信息，请参阅每个函数的文档字符串。

使用自定义内存分配器进行 CUDA

可以将分配器定义为 C/C++ 中的简单函数，并将其编译为共享库。下面的代码显示了一个基本的分配器，它只是跟踪所有内存操作。

#include <sys/types.h>
#include <cuda_runtime_api.h>
#include <iostream>
// Compile with g++ alloc.cc -o alloc.so -I/usr/local/cuda/include -shared -fPIC
extern "C" {
void* my_malloc(ssize_t size, int device, cudaStream_t stream) {
   void *ptr;
   cudaMalloc(&ptr, size);
   std::cout<<"alloc "<<ptr<<size<<std::endl;
   return ptr;
}

void my_free(void* ptr, ssize_t size, int device, cudaStream_t stream) {
   std::cout<<"free "<<ptr<< " "<<stream<<std::endl;
   cudaFree(ptr);
}
}

这可以通过 torch.cuda.memory.CUDAPluggableAllocator 在 python 中使用。用户负责提供指向 .so 文件的路径以及与上述签名匹配的分配/释放函数的名称。

import torch

# Load the allocator
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
    'alloc.so', 'my_malloc', 'my_free')
# Swap the current allocator
torch.cuda.memory.change_current_allocator(new_alloc)
# This will allocate memory in the device using the new allocator
b = torch.zeros(10, device='cuda')

import torch

# Do an initial memory allocator
b = torch.zeros(10, device='cuda')
# Load the allocator
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
    'alloc.so', 'my_malloc', 'my_free')
# This will error since the current allocator was already instantiated
torch.cuda.memory.change_current_allocator(new_alloc)

在同一程序中使用不同的 CUDA 系统分配器

根据您的用例，change_current_allocator() 可能不是您想要使用的，因为它会交换整个程序的 CUDA 分配器（类似于 PYTORCH_ALLOC_CONF=backend:cudaMallocAsync）。例如，如果交换的分配器没有缓存机制，您将失去 PyTorch 的 CUDACachingAllocator 的所有好处。相反，您可以使用 torch.cuda.MemPool 有选择地标记 PyTorch 代码区域以使用自定义分配器。这将允许您在同一个 PyTorch 程序中使用多个 CUDA 系统分配器，同时保留 CUDACachingAllocator 的大部分好处（例如缓存）。使用 torch.cuda.MemPool，您可以利用自定义分配器来启用几个功能，例如

• 使用 ncclMemAlloc 分配器分配 all-reduce 的输出缓冲区可以启用 NVLink Switch Reductions (NVLS)。这可以减少 GPU 资源（SM 和复制引擎）上重叠计算和通信内核之间的争用，尤其是在张量并行工作负载上。

• 对于 Grace CPU 基于的系统，使用 cuMemCreate 并指定 CU_MEM_LOCATION_TYPE_HOST_NUMA 分配 all-gather 的主机输出缓冲区，可以启用基于 Extended GPU Memory (EGM) 的内存传输，从源 GPU 到目标 CPU。这会加速 all-gather，因为传输发生在 NVLink 上，否则会发生在带宽受限的网络接口卡 (NIC) 链路上传输。这种加速的 all-gather 反过来可以加速模型检查点。

• 如果您正在构建一个模型并且一开始不想考虑内存密集型模块的最佳内存放置（例如嵌入表），或者您有一个对性能不敏感且不适合 GPU 的模块，那么您可以只使用 cudaMallocManaged 带有首选 CPU 位置来分配该模块，并首先使您的模型正常工作。

注意

虽然 cudaMallocManaged 提供了使用 CUDA Unified Virtual Memory (UVM) 的便捷自动内存管理，但它不建议用于 DL 工作负载。对于适合 GPU 内存的 DL 工作负载，显式放置始终优于 UVM，因为没有页面错误并且访问模式保持可预测。当 GPU 内存饱和时，UVM 必须执行代价高昂的双重传输，在将新页面引入之前将页面逐出到 CPU。

下面的代码显示了 ncclMemAlloc 包装在 torch.cuda.memory.CUDAPluggableAllocator 中。

import os

import torch
import torch.distributed as dist
from torch.cuda.memory import CUDAPluggableAllocator
from torch.distributed.distributed_c10d import _get_default_group
from torch.utils import cpp_extension


# create allocator
nccl_allocator_source = """
#include <nccl.h>
#include <iostream>
extern "C" {

void* nccl_alloc_plug(size_t size, int device, void* stream) {
  std::cout << "Using ncclMemAlloc" << std::endl;
  void* ptr;
  ncclResult_t err = ncclMemAlloc(&ptr, size);
  return ptr;

}

void nccl_free_plug(void* ptr, size_t size, int device, void* stream) {
  std::cout << "Using ncclMemFree" << std::endl;
  ncclResult_t err = ncclMemFree(ptr);
}

}
"""
nccl_allocator_libname = "nccl_allocator"
nccl_allocator = torch.utils.cpp_extension.load_inline(
    name=nccl_allocator_libname,
    cpp_sources=nccl_allocator_source,
    with_cuda=True,
    extra_ldflags=["-lnccl"],
    verbose=True,
    is_python_module=False,
    build_directory="./",
)

allocator = CUDAPluggableAllocator(
    f"./{nccl_allocator_libname}.so", "nccl_alloc_plug", "nccl_free_plug"
).allocator()

# setup distributed
rank = int(os.getenv("RANK"))
local_rank = int(os.getenv("LOCAL_RANK"))
world_size = int(os.getenv("WORLD_SIZE"))
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")
device = torch.device(f"cuda:{local_rank}")
default_pg = _get_default_group()
backend = default_pg._get_backend(device)

# Note: for convenience, ProcessGroupNCCL backend provides
# the ncclMemAlloc allocator as backend.mem_allocator
allocator = backend.mem_allocator

现在，您可以将此分配器传递给 torch.cuda.MemPool 来定义一个新的内存池

pool = torch.cuda.MemPool(allocator)

然后可以使用 torch.cuda.use_mem_pool 上下文管理器将张量分配到该池中

with torch.cuda.use_mem_pool(pool):
    # tensor gets allocated with ncclMemAlloc passed in the pool
    tensor = torch.arange(1024 * 1024 * 2, device=device)
    print(f"tensor ptr on rank {rank} is {hex(tensor.data_ptr())}")

# register user buffers using ncclCommRegister (called under the hood)
backend.register_mem_pool(pool)

# Collective uses Zero Copy NVLS
dist.all_reduce(tensor[0:4])
torch.cuda.synchronize()
print(tensor[0:4])

请注意在上面的示例中使用了 register_mem_pool。这是 NVLS 缩减的额外步骤，其中用户缓冲区需要使用 NCCL 注册。用户可以使用类似的 deregister_mem_pool 调用来取消注册缓冲区。

要回收内存，用户首先需要确保没有任何内容正在使用该池。当没有张量持有对池的引用时，在删除池时内部将调用 empty_cache()，从而将所有内存返回给系统。

del tensor, del pool

用户可以选择在 MemPool 创建期间指定一个 use_on_oom 布尔值（默认值为 False）。如果为 true，则 CUDACachingAllocator 将能够将此池中的内存作为最后的手段，而不是 OOM。

pool = torch.cuda.MemPool(allocator, use_on_oom=True)
with torch.cuda.use_mem_pool(pool):
    a = torch.randn(40 * 1024 * 1024, dtype=torch.uint8, device="cuda")
del a

# at the memory limit, this will succeed by using pool's memory in order to avoid the oom
b = torch.randn(40 * 1024 * 1024, dtype=torch.uint8, device="cuda")

以下 torch.cuda.MemPool.use_count() 和 torch.cuda.MemPool.snapshot() API 可用于调试目的

pool = torch.cuda.MemPool(allocator)

# pool's use count should be 1 at this point as MemPool object
# holds a reference
assert pool.use_count() == 1

nelem_1mb = 1024 * 1024 // 4

with torch.cuda.use_mem_pool(pool):
    out_0 = torch.randn(nelem_1mb, device="cuda")

    # pool's use count should be 2 at this point as use_mem_pool
    # holds a reference
    assert pool.use_count() == 2

# pool's use count should be back to 1 at this point as use_mem_pool
# released its reference
assert pool.use_count() == 1

with torch.cuda.use_mem_pool(pool):
    # pool should have 1 segment since we made a small allocation (1 MB)
    # above and so the CUDACachingAllocator packed it into a 2 MB buffer
    assert len(pool.snapshot()) == 1

    out_1 = torch.randn(nelem_1mb, device="cuda")

    # pool should still have 1 segment since we made another small allocation
    # (1 MB) that got packed into the existing 2 MB buffer
    assert len(pool.snapshot()) == 1

    out_2 = torch.randn(nelem_1mb, device="cuda")

    # pool now should have 2 segments since the CUDACachingAllocator had
    # to make a new 2 MB buffer to accommodate out_2
    assert len(pool.snapshot()) == 2

注意

• torch.cuda.MemPool 持有对池的引用。当您使用 torch.cuda.use_mem_pool 上下文管理器时，它也会获取对池的另一个引用。在退出上下文管理器时，它将释放其引用。之后，理想情况下，应该只有张量持有对池的引用。一旦张量释放了它们的引用，池的使用计数将为 1，反映了只有 torch.cuda.MemPool 对象持有引用。只有那时，当使用 del 调用池的析构函数时，池中持有的内存才能返回给系统。

• torch.cuda.MemPool 目前不支持 CUDACachingAllocator 的 expandable_segments 模式。

• NCCL 对缓冲区有特定要求，以使其与 NVLS 规约兼容。这些要求在动态工作负载中可能会被破坏，例如，由 CUDACachingAllocator 发送到 NCCL 的缓冲区可能会被分割，因此无法正确对齐。在这种情况下，NCCL 可以使用回退算法代替 NVLS。

• 诸如 ncclMemAlloc 之类的分配器由于对齐要求（CU_MULTICAST_GRANULARITY_RECOMMENDED、CU_MULTICAST_GRANULARITY_MINIMUM）可能会使用比请求更多的内存，并可能导致您的工作负载耗尽内存。

H100/H200 GPU 上使用自定义内存分配器调整 NVLink 性能

在极少数情况下，H100/H200 GPU 上的 NVLink 性能可能会受到数据物理内存布局的影响，从而为开发人员提供调整应用程序以实现最佳吞吐量的机会。

一个例子是当通信内核发出不平衡的 NVLink 读取/写入操作时，数据物理内存布局如何影响性能。在下面的图中，我们可以看到每个 warp 以一致的步长模式访问每个单波中的内存地址。我们可以通过调整工作负载中的步长大小来获得更平衡的负载，或者我们可以实现一个自定义 CUDA 分配器。

_______________________________  _______________________________      _______________________________
| Warp 0 Reading | No-reading |  | Warp 1 Reading | No-reading |  ...  Warp N Reading | No-reading |
_______________________________  _______________________________      _______________________________
<----------------------------->
        Stride size

这样的分配器可以在保持内核的连续虚拟内存地址的同时，策略性地安排到物理内存地址的映射（例如，通过洗牌）。这种技术允许开发人员探索不同的物理访问模式以找到最有效的一种，从而在不修改内核逻辑的情况下解锁更高的性能。正如之前提到的，可以使用 PyTorch 的自定义分配器支持来实现这种分配器的实用实现，其中 malloc 和 free 函数是

// assuming a system with 8 GPUs
struct CustomAllocInfo {
  void** devPtr;  // This will be the usable virtual memory address
  CUdeviceptr dptr;
  size_t totalSize;  // Total size of the allocated memory
  size_t padded_size;
  int device_id;
  std::vector<CUmemGenericAllocationHandle> handles;  // Handles to physical memory allocations
};

// loop over pages
cudaError_t customCudaMalloc(CustomAllocInfo* info) {
    if (!info) return cudaErrorInvalidValue;

    CUdeviceptr dptr;

    // Handles to redundant physical memory allocations which help truncate stride pattern in physical memory
    std::vector<CUmemGenericAllocationHandle> handles_redundant;

    size_t granularity = 0;
    CUmemAllocationProp prop = {};

    int currentDev = info->device_id;
    size_t totalSize = info->totalSize;

    prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
    prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
    prop.location.id = currentDev;
    cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM);
    size_t padded_size = ROUND_UP(totalSize, granularity);

    info->padded_size = padded_size;

    // loop over pages
    size_t iter_granularity = granularity * 64; // 64 * granularity with shift_size = 2 works
    uint32_t iteration_count = (totalSize + iter_granularity - 1) / iter_granularity;

    cuMemAddressReserve(&dptr, padded_size, 0ULL, 0ULL, 0ULL);

    const int shift_size = 2;
    for (size_t i = 0; i < iteration_count; i+=shift_size) {

        CUmemGenericAllocationHandle allocHandle[shift_size];
        for (int shift = 0; (shift < shift_size)&&(i+shift < iteration_count); shift++){
            CHECK_CUDA(cuMemCreate(&allocHandle[shift], iter_granularity, &prop, 0));
            info->handles.push_back(allocHandle[shift]);
        }

        for (int shift = 0; (shift < shift_size)&&(i+shift < iteration_count); shift++){

            // mapping makes the shift (shift -> (shift+1)%shift_size  )
            CHECK_CUDA(cuMemMap(dptr + (i+shift) * iter_granularity, iter_granularity, 0, allocHandle[(shift+1)%shift_size], 0));

            setupMultiGPUAccess(dptr + (i+shift) * iter_granularity, iter_granularity, {0, 1, 2, 3, 4, 5, 6, 7}); // Enable access for all 8 GPUs
        }

        // std::cout << "Here we allocate one redundant page (2MB)..." << std::endl;
        // this is an extra optimization on top of the swizzling. It helps "break"
        // the physical access pattern even more. It can be left out if workload is already
        // performing at SOL with just swizzling.
        CUmemGenericAllocationHandle allocHandle_redundant;
        CHECK_CUDA(cuMemCreate(&allocHandle_redundant, granularity, &prop, 0));
        handles_redundant.push_back(allocHandle_redundant);
    }

    *info->devPtr = (void*)dptr;
    info->dptr = dptr;

    // Release each redundant allocation
    for (auto handle : handles_redundant) {
        // std::cout << "Here we release one redundant page (2MB)..." << std::endl;
        CHECK_CUDA(cuMemRelease(handle));
    }

    return cudaSuccess;
}

void customCudaFree(CustomAllocInfo* info) {
    if (!info) return;

    // CHECK_CUDA(cudaSetDevice(info->device_id));

    CHECK_CUDA(cuMemUnmap(info->dptr, info->padded_size));

    // Unmap and release each allocation
    for (auto handle : info->handles) {
        CHECK_CUDA(cuMemRelease(handle));
    }

    // Unreserve the virtual address space
    // CHECK_CUDA(cuMemAddressFree((CUdeviceptr)*info->devPtr, info->padded_size));
    CHECK_CUDA(cuMemAddressFree(info->dptr, info->padded_size));
}

cuBLAS 工作区

对于 cuBLAS 处理程序和 CUDA 流的每种组合，如果该处理程序和流组合执行需要工作区的 cuBLAS 内核，则将分配一个 cuBLAS 工作区。为了避免重复分配工作区，除非调用 torch._C._cuda_clearCublasWorkspaces()，否则这些工作区不会被释放。可以通过格式为 :[SIZE]:[COUNT] 的环境变量 CUBLAS_WORKSPACE_CONFIG 指定每个分配的工作区大小。例如，默认的工作区大小为每个分配是 CUBLAS_WORKSPACE_CONFIG=:4096:2:16:8，它指定总大小为 2 * 4096 + 8 * 16 KiB。要强制 cuBLAS 避免使用工作区，请设置 CUBLAS_WORKSPACE_CONFIG=:0:0。

cuFFT 计划缓存

对于每个 CUDA 设备，都使用 cuFFT 计划的 LRU 缓存来加速重复运行 CUDA 张量上相同几何形状和相同配置的 FFT 方法（例如，torch.fft.fft()）。由于某些 cuFFT 计划可能会分配 GPU 内存，因此这些缓存具有最大容量。

您可以使用以下 API 来控制和查询当前设备的缓存属性

• torch.backends.cuda.cufft_plan_cache.max_size 提供缓存的容量（CUDA 10 及更高版本上的默认值为 4096，旧版本 CUDA 上的默认值为 1023）。直接设置此值会修改容量。

• torch.backends.cuda.cufft_plan_cache.size 提供当前驻留在缓存中的计划数。

• torch.backends.cuda.cufft_plan_cache.clear() 清除缓存。

要控制和查询非默认设备的计划缓存，您可以使用 torch.device 对象或设备索引对 torch.backends.cuda.cufft_plan_cache 对象进行索引，并访问上述属性之一。例如，要设置设备 1 的缓存容量，可以编写 torch.backends.cuda.cufft_plan_cache[1].max_size = 10。

即时编译

PyTorch 会对某些操作（如 torch.special.zeta）在 CUDA 张量上执行时进行即时编译。此编译可能耗时（取决于您的硬件和软件，最多可能需要几秒钟），并且可能为单个运算符发生多次，因为许多 PyTorch 运算符实际上会从各种内核中选择，每个内核必须根据其输入编译一次。

默认情况下，PyTorch 在 $XDG_CACHE_HOME/torch/kernels 中创建一个内核缓存（如果定义了 XDG_CACHE_HOME）或在 $HOME/.cache/torch/kernels 中（如果未定义）（Windows 除外，Windows 上尚不支持内核缓存）。可以使用两个环境变量直接控制缓存行为。如果 USE_PYTORCH_KERNEL_CACHE 设置为 0，则不会使用缓存，如果设置了 PYTORCH_KERNEL_CACHE_PATH，则该路径将用作内核缓存而不是默认位置。

最佳实践

设备无关代码

由于 PyTorch 的结构，您可能需要显式编写设备无关（CPU 或 GPU）代码；一个例子是在创建循环神经网络的初始隐藏状态时创建一个新的张量。

第一步是确定是否应该使用 GPU。一种常见模式是使用 Python 的 argparse 模块读取用户参数，并使用一个标志来禁用 CUDA，并结合 is_available()。在下文中，args.device 产生一个 torch.device 对象，可用于将张量移动到 CPU 或 CUDA。

import argparse
import torch

parser = argparse.ArgumentParser(description='PyTorch Example')
parser.add_argument('--disable-cuda', action='store_true',
                    help='Disable CUDA')
args = parser.parse_args()
args.device = None
if not args.disable_cuda and torch.cuda.is_available():
    args.device = torch.device('cuda')
else:
    args.device = torch.device('cpu')

注意

在评估给定环境中的 CUDA 可用性（is_available()）时，PyTorch 的默认行为是调用 CUDA 运行时 API 方法 cudaGetDeviceCount。由于此调用反过来初始化 CUDA 驱动程序 API（通过 cuInit），如果尚未初始化，后续对已运行 is_available() 的进程进行分叉将失败，并出现 CUDA 初始化错误。

可以在您的环境中设置 PYTORCH_NVML_BASED_CUDA_CHECK=1，然后再导入执行 is_available() 的 PyTorch 模块（或直接执行它），以便指示 is_available() 尝试基于 NVML 的评估（nvmlDeviceGetCount_v2）。如果基于 NVML 的评估成功（即 NVML 发现/初始化未失败），is_available() 调用将不会破坏后续的分叉。

如果 NVML 发现/初始化失败，is_available() 将回退到标准的 CUDA 运行时 API 评估，并且上述分叉约束将适用。

请注意，上述基于 NVML 的 CUDA 可用性评估提供的保证比默认 CUDA 运行时 API 方法弱（后者要求 CUDA 初始化成功）。在某些情况下，基于 NVML 的检查可能会成功，而后续的 CUDA 初始化可能会失败。

现在我们有了 args.device，我们可以使用它在所需设备上创建张量。

x = torch.empty((8, 42), device=args.device)
net = Network().to(device=args.device)

这可以在许多情况下用于生成设备无关的代码。以下是在使用数据加载器时的一个示例

cuda0 = torch.device('cuda:0')  # CUDA GPU 0
for i, x in enumerate(train_loader):
    x = x.to(cuda0)

当在具有多个 GPU 的系统上工作时，可以使用 CUDA_VISIBLE_DEVICES 环境变量来管理 PyTorch 可用的 GPU。如上所述，要手动控制在哪个 GPU 上创建张量，最佳实践是使用 torch.cuda.device 上下文管理器。

print("Outside device is 0")  # On device 0 (default in most scenarios)
with torch.cuda.device(1):
    print("Inside device is 1")  # On device 1
print("Outside device is still 0")  # On device 0

如果您有一个张量，并且想要创建一个相同类型且位于相同设备上的新张量，那么您可以使用一个 torch.Tensor.new_* 方法（参见 torch.Tensor）。虽然前面提到的 torch.* 工厂函数（Creation Ops）依赖于当前的 GPU 上下文以及您传递的属性参数，但 torch.Tensor.new_* 方法会保留张量的设备和其他属性。

当在模块中创建需要在前向传递期间内部创建新张量时，这是推荐的做法。

cuda = torch.device('cuda')
x_cpu = torch.empty(2)
x_gpu = torch.empty(2, device=cuda)
x_cpu_long = torch.empty(2, dtype=torch.int64)

y_cpu = x_cpu.new_full([3, 2], fill_value=0.3)
print(y_cpu)

    tensor([[ 0.3000,  0.3000],
            [ 0.3000,  0.3000],
            [ 0.3000,  0.3000]])

y_gpu = x_gpu.new_full([3, 2], fill_value=-5)
print(y_gpu)

    tensor([[-5.0000, -5.0000],
            [-5.0000, -5.0000],
            [-5.0000, -5.0000]], device='cuda:0')

y_cpu_long = x_cpu_long.new_tensor([[1, 2, 3]])
print(y_cpu_long)

    tensor([[ 1,  2,  3]])

如果您想创建一个与另一个张量相同类型和大小的张量，并用 1 或 0 填充它，ones_like() 或 zeros_like() 作为方便的辅助函数提供（它们也保留 torch.device 和 torch.dtype 的张量）。

x_cpu = torch.empty(2, 3)
x_gpu = torch.empty(2, 3)

y_cpu = torch.ones_like(x_cpu)
y_gpu = torch.zeros_like(x_gpu)

使用固定内存缓冲区

警告

这是一个高级技巧。如果您过度使用固定内存，当 RAM 不足时可能会导致严重问题，并且您应该意识到固定内存通常是一项代价高昂的操作。

从固定（页面锁定）内存发起的从主机到 GPU 的复制速度要快得多。CPU 张量和存储公开了一个 pin_memory() 方法，该方法返回对象的副本，并将数据放入固定区域。

此外，一旦您固定了一个张量或存储，您就可以使用异步 GPU 复制。只需将额外的 non_blocking=True 参数传递给 to() 或 cuda() 调用。这可用于重叠数据传输和计算。

您可以通过在构造函数中传递 pin_memory=True，使 DataLoader 返回放置在固定内存中的批次。

使用 nn.parallel.DistributedDataParallel 代替 multiprocessing 或 nn.DataParallel

大多数涉及批处理输入和多个 GPU 的用例都应默认使用 DistributedDataParallel 以利用多个 GPU。

使用 CUDA 模型与 multiprocessing 存在一些重要的注意事项；除非小心处理以满足数据处理要求，否则您的程序很可能行为不正确或未定义。

建议使用 DistributedDataParallel，而不是 DataParallel 进行多 GPU 训练，即使只有一个节点。

DistributedDataParallel 和 DataParallel 之间的区别是：DistributedDataParallel 为每个 GPU 创建一个进程，而 DataParallel 使用多线程。通过使用多进程，每个 GPU 都有其专用进程，这避免了由 Python 解释器的 GIL 引起的性能开销。

如果您使用 DistributedDataParallel，您可以使用 torch.distributed.launch 实用程序来启动您的程序，请参阅 Launch utility。

CUDA 图形

CUDA 图形是 CUDA 流及其依赖流执行的工作（主要是内核及其参数）的记录。有关一般原理和底层 CUDA API 的详细信息，请参阅 Getting Started with CUDA Graphs 和 CUDA C 编程指南的 Graphs section。

PyTorch 支持使用 stream capture 构建 CUDA 图形，这会将 CUDA 流置于 capture mode。向捕获流发出的 CUDA 工作实际上不会在 GPU 上运行。相反，工作被记录在图形中。

捕获后，可以 launch 图形来运行 GPU 工作多次。每次重放都运行相同的内核，使用相同的参数。对于指针参数，这意味着使用相同的内存地址。通过在每次重放之前用新数据（例如，来自新的批次）填充输入内存，您可以重新运行相同的工作，处理新的数据。

为什么使用 CUDA 图形？

重放图形牺牲了典型渴望执行的动态灵活性，以换取 大大降低的 CPU 开销。图形的参数和内核是固定的，因此图形重放跳过了所有参数设置和内核调度层，包括 Python、C++ 和 CUDA 驱动程序开销。在底层，重放使用单个调用 cudaGraphLaunch 将整个图形的工作提交给 GPU。图形中的内核也在 GPU 上执行得稍快一些，但消除 CPU 开销是主要好处。

如果您的网络全部或部分是图形安全的（通常这意味着静态形状和静态控制流，但请参阅其他 constraints），并且您怀疑其运行时至少在某种程度上受到 CPU 限制，则应尝试 CUDA 图形。

PyTorch API

警告

此 API 处于 Beta 版，未来版本中可能会更改。

PyTorch 通过原始 torch.cuda.CUDAGraph 类和两个便利包装器 torch.cuda.graph 和 torch.cuda.make_graphed_callables 公开图形。

torch.cuda.graph 是一个简单且通用的上下文管理器，它在其上下文中捕获 CUDA 工作。在捕获之前，请通过运行几个渴望迭代来预热要捕获的工作负载。预热必须发生在侧流上。由于图形每次重放都从相同的内存地址读取和写入，因此您必须在捕获期间维护对包含输入和输出数据的张量的长期引用。要在新的输入数据上运行图形，请将新的数据复制到捕获的输入张量，重放图形，然后从捕获的输出张量读取新的输出。示例

g = torch.cuda.CUDAGraph()

# Placeholder input used for capture
static_input = torch.empty((5,), device="cuda")

# Warmup before capture
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        static_output = static_input * 2
torch.cuda.current_stream().wait_stream(s)

# Captures the graph
# To allow capture, automatically sets a side stream as the current stream in the context
with torch.cuda.graph(g):
    static_output = static_input * 2

# Fills the graph's input memory with new data to compute on
static_input.copy_(torch.full((5,), 3, device="cuda"))
g.replay()
# static_output holds the results
print(static_output)  # full of 3 * 2 = 6

# Fills the graph's input memory with more data to compute on
static_input.copy_(torch.full((5,), 4, device="cuda"))
g.replay()
print(static_output)  # full of 4 * 2 = 8

请参阅 Whole-network capture、Usage with torch.cuda.amp 和 Usage with multiple streams 以获取真实和高级模式。

make_graphed_callables 更加复杂。make_graphed_callables 接受 Python 函数和 torch.nn.Module。对于传递的每个函数或模块，它都会创建前向传递和反向传递工作的单独图形。请参阅 Partial-network capture。

约束

如果一组操作不违反以下任何约束，则该操作是 capturable。

约束适用于 torch.cuda.graph 上下文中的所有工作，以及您传递给 torch.cuda.make_graphed_callables() 的任何可调用对象的正向和反向传递中的所有工作。

违反任何一项都可能导致运行时错误

• 捕获必须发生在非默认流上。（这仅在您使用原始 CUDAGraph.capture_begin 和 CUDAGraph.capture_end 调用时才是一个问题。graph 和 make_graphed_callables() 会为您设置一个侧流。）

• 禁止将 CPU 与 GPU 同步的操作（例如，.item() 调用）。

• CUDA RNG 操作是允许的，并且在使用图中的多个 torch.Generator 实例时，必须在使用 CUDAGraph.register_generator_state 注册它们，然后再进行图捕获。避免在捕获期间使用 Generator.get_state 和 Generator.set_state；相反，请使用 Generator.graphsafe_set_state 和 Generator.graphsafe_get_state 在图上下文中安全地管理生成器状态。这确保了 CUDA 图中正确的 RNG 操作和生成器管理。

违反任何这些规定很可能导致静默的数值错误或未定义行为

• 在进程内，一次只能进行一次捕获。

• 在捕获进行时，此进程（在任何线程上）中不能运行任何未捕获的 CUDA 工作。

• CPU 工作不会被捕获。如果捕获的操作包括 CPU 工作，则在重放期间该工作将被省略。

• 每次重放都从相同的（虚拟）内存地址读取和写入。

• 禁止基于 CPU 或 GPU 数据进行动态控制流。

• 禁止动态形状。该图假定捕获的操作序列中的每个张量在每次重放中都具有相同的大小和布局。

• 允许在捕获中使用多个流，但存在 限制。

非约束

• 捕获后，该图可以在任何流上重放。

全网络捕获

如果您的整个网络都可以捕获，您可以捕获并重放整个迭代

N, D_in, H, D_out = 640, 4096, 2048, 1024
model = torch.nn.Sequential(torch.nn.Linear(D_in, H),
                            torch.nn.Dropout(p=0.2),
                            torch.nn.Linear(H, D_out),
                            torch.nn.Dropout(p=0.1)).cuda()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Placeholders used for capture
static_input = torch.randn(N, D_in, device='cuda')
static_target = torch.randn(N, D_out, device='cuda')

# warmup
# Uses static_input and static_target here for convenience,
# but in a real setting, because the warmup includes optimizer.step()
# you must use a few batches of real data.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for i in range(3):
        optimizer.zero_grad(set_to_none=True)
        y_pred = model(static_input)
        loss = loss_fn(y_pred, static_target)
        loss.backward()
        optimizer.step()
torch.cuda.current_stream().wait_stream(s)

# capture
g = torch.cuda.CUDAGraph()
# Sets grads to None before capture, so backward() will create
# .grad attributes with allocations from the graph's private pool
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
    static_y_pred = model(static_input)
    static_loss = loss_fn(static_y_pred, static_target)
    static_loss.backward()
    optimizer.step()

real_inputs = [torch.rand_like(static_input) for _ in range(10)]
real_targets = [torch.rand_like(static_target) for _ in range(10)]

for data, target in zip(real_inputs, real_targets):
    # Fills the graph's input memory with new data to compute on
    static_input.copy_(data)
    static_target.copy_(target)
    # replay() includes forward, backward, and step.
    # You don't even need to call optimizer.zero_grad() between iterations
    # because the captured backward refills static .grad tensors in place.
    g.replay()
    # Params have been updated. static_y_pred, static_loss, and .grad
    # attributes hold values from computing on this iteration's data.

部分网络捕获

如果您的网络的部分内容不安全可捕获（例如，由于动态控制流、动态形状、CPU 同步或必要的 CPU 端逻辑），您可以以渴望模式运行不安全的部分，并使用 torch.cuda.make_graphed_callables() 仅对捕获安全的部分进行图形化。

默认情况下，由 make_graphed_callables() 返回的可调用对象是 autograd 感知的，并且可以用作您传递的函数或 nn.Module 的直接替代品，用于训练循环。

make_graphed_callables() 内部创建 CUDAGraph 对象，运行预热迭代，并根据需要维护静态输入和输出。因此（与 torch.cuda.graph 不同），您无需手动处理这些。

在以下示例中，数据依赖的动态控制流意味着网络不能端到端地捕获，但 make_graphed_callables() 允许我们将图安全的区域作为图进行捕获和运行，无论如何

N, D_in, H, D_out = 640, 4096, 2048, 1024

module1 = torch.nn.Linear(D_in, H).cuda()
module2 = torch.nn.Linear(H, D_out).cuda()
module3 = torch.nn.Linear(H, D_out).cuda()

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(chain(module1.parameters(),
                                  module2.parameters(),
                                  module3.parameters()),
                            lr=0.1)

# Sample inputs used for capture
# requires_grad state of sample inputs must match
# requires_grad state of real inputs each callable will see.
x = torch.randn(N, D_in, device='cuda')
h = torch.randn(N, H, device='cuda', requires_grad=True)

module1 = torch.cuda.make_graphed_callables(module1, (x,))
module2 = torch.cuda.make_graphed_callables(module2, (h,))
module3 = torch.cuda.make_graphed_callables(module3, (h,))

real_inputs = [torch.rand_like(x) for _ in range(10)]
real_targets = [torch.randn(N, D_out, device="cuda") for _ in range(10)]

for data, target in zip(real_inputs, real_targets):
    optimizer.zero_grad(set_to_none=True)

    tmp = module1(data)  # forward ops run as a graph

    if tmp.sum().item() > 0:
        tmp = module2(tmp)  # forward ops run as a graph
    else:
        tmp = module3(tmp)  # forward ops run as a graph

    loss = loss_fn(tmp, target)
    # module2's or module3's (whichever was chosen) backward ops,
    # as well as module1's backward ops, run as graphs
    loss.backward()
    optimizer.step()

与 torch.cuda.amp 的用法

对于典型的优化器，GradScaler.step 会在 CPU 和 GPU 之间同步，这在捕获期间是被禁止的。为了避免错误，请使用 部分网络捕获，或者（如果前向传播、损失和反向传播是捕获安全的），捕获前向传播、损失和反向传播，但不捕获优化器步骤

# warmup
# In a real setting, use a few batches of real data.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for i in range(3):
        optimizer.zero_grad(set_to_none=True)
        with torch.cuda.amp.autocast():
            y_pred = model(static_input)
            loss = loss_fn(y_pred, static_target)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
torch.cuda.current_stream().wait_stream(s)

# capture
g = torch.cuda.CUDAGraph()
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
    with torch.cuda.amp.autocast():
        static_y_pred = model(static_input)
        static_loss = loss_fn(static_y_pred, static_target)
    scaler.scale(static_loss).backward()
    # don't capture scaler.step(optimizer) or scaler.update()

real_inputs = [torch.rand_like(static_input) for _ in range(10)]
real_targets = [torch.rand_like(static_target) for _ in range(10)]

for data, target in zip(real_inputs, real_targets):
    static_input.copy_(data)
    static_target.copy_(target)
    g.replay()
    # Runs scaler.step and scaler.update eagerly
    scaler.step(optimizer)
    scaler.update()

与多个流的用法

捕获模式会自动传播到与捕获流同步的任何流。在捕获过程中，您可以通过向不同的流发出调用来暴露并行性，但整体流依赖 DAG 必须在捕获开始后从初始捕获流分支出来，并在捕获结束前重新连接到初始流

with torch.cuda.graph(g):
    # at context manager entrance, torch.cuda.current_stream()
    # is the initial capturing stream

    # INCORRECT (does not branch out from or rejoin initial stream)
    with torch.cuda.stream(s):
        cuda_work()

    # CORRECT:
    # branches out from initial stream
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        cuda_work()
    # rejoins initial stream before capture ends
    torch.cuda.current_stream().wait_stream(s)

注意

为了避免让熟悉 nsight systems 或 nvprof 的高级用户混淆：与渴望执行不同，该图将捕获中的非平凡流 DAG 解释为提示，而不是命令。在重放期间，该图可以将独立的运算重新组织到不同的流上，或以不同的顺序排队（同时尊重您原始 DAG 的整体依赖关系）。

与 DistributedDataParallel 的用法

NCCL < 2.9.6

NCCL 版本早于 2.9.6 不允许捕获集体运算。您必须使用 部分网络捕获，这会将 allreduces 推迟到反向传播的图形化部分之外。

在用 DDP 包装网络之前，对可图形化的网络部分调用 make_graphed_callables()。

NCCL >= 2.9.6

NCCL 版本 2.9.6 或更高版本允许在图中进行集体运算。捕获 整个反向传播 的方法是一种可行的选择，但需要三个设置步骤。

1. 禁用 DDP 的内部异步错误处理

   os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
   torch.distributed.init_process_group(...)
   

2. 在完整反向传播捕获之前，DDP 必须在侧流上下文中构造

   with torch.cuda.stream(s):
       model = DistributedDataParallel(model)
   

3. 您的预热必须在捕获之前运行至少 11 个启用了 DDP 的渴望迭代。

图内存管理

捕获的图每次重放都作用于相同的虚拟地址。如果 PyTorch 释放了内存，后续重放可能会遇到非法内存访问。如果 PyTorch 将内存重新分配给新的张量，重放可能会破坏这些张量看到的值。因此，图使用的虚拟地址必须在重放期间为图保留。PyTorch 缓存分配器通过检测何时进行捕获并从图私有内存池满足捕获的分配来做到这一点。私有池在 CUDAGraph 对象和在捕获期间创建的所有张量超出范围时保持活动状态。

私有池是自动维护的。默认情况下，分配器为每个捕获创建一个单独的私有池。如果您捕获多个图，这种保守的方法可确保图重放不会相互破坏值，但有时会不必要地浪费内存。

跨捕获共享内存

为了节约私有池中缓存的内存，torch.cuda.graph 和 torch.cuda.make_graphed_callables() 可选地允许不同的捕获共享相同的私有池。如果知道它们总是以捕获的顺序重放，并且从不并发重放，则一组图共享私有池是安全的。

torch.cuda.graph 的 pool 参数是使用特定私有池的提示，可以用作如图所示那样在图之间共享内存。

g1 = torch.cuda.CUDAGraph()
g2 = torch.cuda.CUDAGraph()

# (create static inputs for g1 and g2, run warmups of their workloads...)

# Captures g1
with torch.cuda.graph(g1):
    static_out_1 = g1_workload(static_in_1)

# Captures g2, hinting that g2 may share a memory pool with g1
with torch.cuda.graph(g2, pool=g1.pool()):
    static_out_2 = g2_workload(static_in_2)

static_in_1.copy_(real_data_1)
static_in_2.copy_(real_data_2)
g1.replay()
g2.replay()

在不相互依赖于彼此输出的情况下，也可以在单独的图之间共享内存池，前提是它们从不并发运行。请注意，当它们共享池时，重放一个图可能会破坏另一个图的输出，除非事先对输出调用 clone()。这种模式经常用于在运行时接受可变批量大小的推理服务器。vLLM 是一个值得注意的例子；请参阅 此处 和 此处。

使用 torch.cuda.make_graphed_callables()，如果你想对几个可调用对象进行图形化，并且知道它们总是以相同的顺序运行（并且从不并发），那么将它们作为元组传递，顺序与它们在实际工作负载中运行的顺序相同，make_graphed_callables() 将使用共享的私有池捕获它们的图形。

如果在实际工作负载中，你的可调用对象运行的顺序偶尔会发生变化，或者它们会并发运行，那么将它们作为元组传递给单个 make_graphed_callables() 调用是不允许的。相反，你必须为每个可调用对象单独调用 make_graphed_callables()。