CUDA 语义

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 1.7 开始，有一个名为 allow_tf32 的新标志。此标志在 PyTorch 1.7 到 PyTorch 1.11 中默认为 True，在 PyTorch 1.12 及更高版本中默认为 False。此标志控制是否允许 PyTorch 内部使用 TensorFloat32 (TF32) 张量核心（自 Ampere 以来在 NVIDIA GPU 上可用）来计算 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_float_32_matmul_precision() 更广泛地设置（不仅限于 CUDA）。请注意，除了 matmul 和卷积本身之外，内部使用 matmul 或卷积的函数和 nn 模块也会受到影响。这些包括 nn.Linear、nn.Conv*、cdist、tensordot、仿射网格和网格采样、自适应对数 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 倍，并且相对于双精度，相对误差大约大 2 个数量级。请注意，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 中的降低精度归约

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 中的降低精度归约

BF16 GEMM 存在类似的标志（如上所述）。请注意，如果观察到工作负载中的数值不稳定，则 BF16 的此开关默认设置为 True，您可能希望将其设置为 False。

如果不需要降低精度归约，用户可以使用以下方式禁用 bf16 GEMM 中的降低精度归约

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False

要在 C++ 中切换降低精度归约标志，可以执行

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

异步执行

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

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

您可以通过设置环境变量 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() 调用确保在我们在侧流上开始运行 sum(A) 之前，normal_() 执行已完成。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 注意：在默认流上使用 grads

在早期版本的 PyTorch（1.9 及更早版本）中，autograd 引擎始终将默认流与所有反向操作同步，因此以下模式

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

只要 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_CUDA_ALLOC_CONF 优化内存使用

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

缓存分配器的行为可以通过环境变量 PYTORCH_CUDA_ALLOC_CONF 控制。格式为 PYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>... 可用选项

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

• 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 的幂次方，并更好地利用块。在原生 CUDACachingAllocator 中，大小以 512 的块大小的倍数向上舍入，因此这对于较小的大小效果很好。但是，这对于较大的附近分配可能效率低下，因为每个分配都将转到不同大小的块，并且这些块的重用被最小化。这可能会创建大量未使用的块，并浪费 GPU 内存容量。此选项启用将分配大小四舍五入到最接近的 2 的幂次方。例如，如果我们需要向上舍入大小为 1200 的大小，并且如果除数数量为 4，则大小 1200 介于 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 之间。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 选项是一个布尔标志，用于启用后台线程来处理事件。这避免了与快速分配路径中查询/处理事件相关的任何慢速路径。此功能默认禁用。

注意

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 文件的路径以及与上面指定的签名匹配的 alloc/free 函数的名称。

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)

cuBLAS 工作区

对于 cuBLAS 句柄和 CUDA 流的每种组合，如果该句柄和流组合执行需要工作区的 cuBLAS 内核，则将分配一个 cuBLAS 工作区。为了避免重复分配工作区，除非调用 torch._C._cuda_clearCublasWorkspaces()，否则不会释放这些工作区。每个分配的工作区大小可以通过环境变量 CUBLAS_WORKSPACE_CONFIG 指定，格式为 :[SIZE]:[COUNT]。例如，每个分配的默认工作区大小为 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。

即时编译

当在 CUDA 张量上执行某些操作时，PyTorch 会即时编译这些操作，例如 torch.special.zeta。此编译可能非常耗时（根据您的硬件和软件，最多需要几秒钟），并且对于单个运算符可能会发生多次，因为许多 PyTorch 运算符实际上是从各种内核中选择的，每个内核都必须编译一次，具体取决于它们的输入。此编译每个进程执行一次，或者如果使用内核缓存，则仅执行一次。

默认情况下，如果定义了 XDG_CACHE_HOME，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() 的进程的后续 fork 将失败，并显示 CUDA 初始化错误。

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

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

请注意，上述基于 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 环境变量标志来管理哪些 GPU 可供 PyTorch 使用。如上所述，要手动控制在哪个 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.* 工厂函数 (创建操作) 取决于当前的 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 的构造函数，使 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 实用程序来启动您的程序，请参阅 第三方后端。

CUDA 图

CUDA 图是 CUDA 流及其依赖流执行的工作（主要是内核及其参数）的记录。有关底层 CUDA API 的一般原则和详细信息，请参阅 CUDA 图入门 和 CUDA C 编程指南的 图部分。

PyTorch 支持使用 流捕获 构建 CUDA 图，这会将 CUDA 流置于捕获模式。发布到捕获流的 CUDA 工作实际上不会在 GPU 上运行。相反，工作记录在图中。

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

为什么使用 CUDA 图？

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

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

PyTorch API

警告

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

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

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

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

有关实际和高级模式，请参阅全网络捕获、与 torch.cuda.amp 一起使用 和 与多个流一起使用。

make_graphed_callables 更复杂。make_graphed_callables 接受 Python 函数和 torch.nn.Module。对于每个传递的函数或 Module，它会创建前向传递和后向传递工作的单独图。请参阅部分网络捕获。

约束

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

约束适用于 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 端逻辑），您可以 eager 运行不安全的部分，并使用 torch.cuda.make_graphed_callables() 仅对捕获安全的部分进行图形化。

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

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

在以下示例中，数据相关的动态控制流意味着网络无法端到端捕获，但 make_graphed_callables() 允许我们捕获并运行图安全的 sections 作为图，无论如何。

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 中查看回放的高级用户感到困惑：与 eager execution 不同，图将捕获中重要的流 DAG 解释为提示，而不是命令。在回放期间，图可能会将独立的 ops 重新组织到不同的流上，或者以不同的顺序入队（同时尊重你原始 DAG 的整体依赖关系）。

与 DistributedDataParallel 的用法

NCCL < 2.9.6

早于 2.9.6 的 NCCL 版本不允许捕获 collectives。你必须使用 部分网络捕获，这会将 allreduces 推迟到反向传播的图 sections 之外发生。

在用 DDP 包装网络之前，在可图化的网络 sections 上调用 make_graphed_callables()。

NCCL >= 2.9.6

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

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

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

2. 在完全反向传播捕获之前，DDP 必须在 side-stream 上下文中构建

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

3. 你的预热必须在捕获之前至少运行 11 次启用 DDP 的 eager 迭代。

图内存管理

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

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

跨捕获共享内存

为了节省私有池中存储的内存，torch.cuda.graph 和 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()

使用 torch.cuda.make_graphed_callables()，如果你想对几个 callables 进行图化，并且你知道它们将始终以相同的顺序运行（并且永远不会并发），请将它们作为元组传递，顺序与它们在实际工作负载中运行的顺序相同，并且 make_graphed_callables() 将使用共享的私有池捕获它们的图。

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