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)

安培（及后续）设备上的张量浮点32 (TF32)¶

从PyTorch 1.7开始，有一个新的标志叫做allow_tf32。在PyTorch 1.7到PyTorch 1.11版本中，默认值为True，在PyTorch 1.12及以后版本中，默认值为False。此标志控制PyTorch是否可以使用自Ampere架构以来NVIDIA GPU上可用的TensorFloat32 (TF32) 张量核心，用于内部计算矩阵乘法（包括批量矩阵乘法）和卷积。

TF32 张量核心旨在通过对输入数据进行舍入以具有 10 位尾数并在 FP32 精度下累积结果，实现在 torch.float32 张量上实现更好的矩阵乘法和卷积性能，同时保持 FP32 动态范围。

矩阵乘法和卷积是分别控制的，它们对应的标志可以在以下位置访问：

# 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

矩阵乘法的精度也可以更广泛地设置（不限于CUDA）通过set_float_32_matmul_precision()。请注意，除了矩阵乘法和卷积本身之外，内部使用矩阵乘法或卷积的函数和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倍，并且与双精度相比的相对误差大约大两个数量级。请注意，TF32与单精度的速度确切比率取决于硬件代际，因为诸如内存带宽与计算的比例以及TF32与FP32矩阵乘法吞吐量的比例等特性可能会因代际或模型的不同而有所变化。如果需要完整的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的更多信息，请参见：

FP16 GEMMs中的精度降低¶

fp16 矩阵乘法可能会在某些中间步骤中使用较低精度的缩减（例如，在 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 GEMMs中的降低精度归约：

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

要在C++中切换降低精度的减少标志，可以这样做

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

BF16 GEMMs中的降低精度减少¶

类似地，BFloat16 GEMM 也存在一个相同的标志（如上所述）。请注意，默认情况下此开关设置为True，如果你的工作负载中观察到数值不稳定，可以将其设置为False。

如果用户不希望使用降低精度的归约操作，可以通过以下方式在bf16 GEMMs中禁用降低精度的归约操作：

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False

要在C++中切换降低精度的减少标志，可以这样做

at::globalContext().setAllowBF16ReductionCuBLAS(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 流是属于特定设备的线性执行序列。通常情况下，您不需要显式创建一个：默认情况下，每个设备使用自己的“默认”流。 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)), 的行为是

可选地填充初始梯度（s），
调用反向传播，
使用梯度

具有与任何一组操作相同的流语义关系：

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及更早版本）中，自动梯度引擎总是将默认流与所有反向操作同步，因此以下模式：

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 将被忽略。
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。

注意

一些由 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);
}
}

这可以通过python使用torch.cuda.memory.CUDAPluggableAllocator实现。用户负责提供.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设备，使用LRU缓存来存储cuFFT计划，以加速在相同几何和配置下的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.backends.cuda.cufft_plan_cache 对象，该对象可以是一个 torch.device 对象或设备索引，并访问上述其中一个属性。例如，要设置设备 1 的缓存容量，可以编写 torch.backends.cuda.cufft_plan_cache[1].max_size = 10。

即时编译¶

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

默认情况下，PyTorch 会在 $XDG_CACHE_HOME/torch/kernels 中创建一个内核缓存，如果 XDG_CACHE_HOME 已定义；如果没有定义，则在 $HOME/.cache/torch/kernels 中创建（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初始化错误而失败。

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

如果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 环境标志来管理哪些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)

使用固定内存缓冲区¶

警告

这是一个高级提示。如果你过度使用固定内存，在内存不足时可能会导致严重问题，并且你应该意识到，固定操作通常是昂贵的。

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

另外，一旦你固定了一个张量或存储，你可以使用异步GPU拷贝。只需在调用to() 或 cuda()时传递一个额外的non_blocking=True参数。这可以用来重叠数据传输与计算。

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

使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel¶

大多数涉及批量输入和多GPU的用例应默认使用DistributedDataParallel来利用多个GPU。

使用CUDA模型时存在重要的注意事项； multiprocessing；除非小心处理以满足数据处理要求，否则您的程序可能会出现错误或未定义的行为。

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

4 和 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图？¶

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

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

PyTorch API¶

警告

该 API 处于测试阶段，未来版本可能会有所更改。

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

查看全网捕获, 与 torch.cuda.amp 一起使用, 以及与多个流一起使用以了解现实和高级模式。

make_graphed_callables 更加复杂。 make_graphed_callables 接受 Python 函数和 torch.nn.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 实例并将其作为 generator 参数传递给RNG函数是禁止的。

违反其中任何一项都可能导致无声的数值错误或未定义行为：

在过程中，任何时候只能有一个捕获正在进行。
在捕获进行期间，此进程中（任何线程上）不得运行未被捕获的 CUDA 操作。
CPU 的工作未被捕获。如果被捕获的操作包括 CPU 工作，那么在重放期间会省略这部分工作。
每次重放都会读取和写入相同的（虚拟）内存地址。
禁止动态控制流（基于CPU或GPU数据）。
动态形状不受支持。该图假设捕获的操作序列中的每个张量在每次重放时都具有相同的大小和布局。
在捕获中使用多个流是允许的，但存在限制。

Non-constraints¶

一旦被捕获，该图可以在任何流上重放。

全网络捕获¶

如果你整个网络都可以被捕获，你可以捕获并重放整个迭代：

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()返回的可调用函数是自动求梯度感知的，并且可以在训练循环中作为您传递的函数或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 延迟到图的反向传播部分之外进行。

调用make_graphed_callables()在可绘图网络部分之前使用DDP包装网络。

NCCL ≥ 2.9.6¶

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

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

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

在完整反向传播捕获之前，DDP 必须在一个旁流上下文中构建。
```
with torch.cuda.stream(s):
    model = DistributedDataParallel(model)
```
您的热身运行必须在捕获之前至少进行 11 次启用 DDP 的急切迭代。

图形内存管理¶

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

私有池会自动维护。默认情况下，分配器为每个捕获创建一个单独的私有池。如果你捕获了多个图，这种保守的方法可以确保图的重播永远不会相互破坏各自的值，但有时会无谓地浪费内存。

CUDA语义¶

安培（及后续）设备上的张量浮点32 (TF32)¶

FP16 GEMMs中的精度降低¶

BF16 GEMMs中的降低精度减少¶

异步执行¶

CUDA 流¶

反向传播的流语义¶

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

内存管理¶

优化内存使用 `PYTORCH_CUDA_ALLOC_CONF`¶

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

cuBLAS 工作区¶

cuFFT计划缓存¶

即时编译¶

最佳实践¶

设备无关代码¶

使用固定内存缓冲区¶

使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel¶

CUDA 图形¶

为什么使用CUDA图？¶

PyTorch API¶

约束条件¶

Non-constraints¶

全网络捕获¶

部分网络捕获¶

与torch.cuda amp的用法¶

多流使用¶

使用DistributedDataParallel¶

NCCL < 2.9.6¶

NCCL ≥ 2.9.6¶

图形内存管理¶

文档

教程

资源

CUDA语义¶

安培（及后续）设备上的张量浮点32 (TF32)¶

FP16 GEMMs中的精度降低¶

BF16 GEMMs中的降低精度减少¶

异步执行¶

CUDA 流¶

反向传播的流语义¶

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

内存管理¶

优化内存使用 PYTORCH_CUDA_ALLOC_CONF¶

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

cuBLAS 工作区¶

cuFFT计划缓存¶

即时编译¶

最佳实践¶

设备无关代码¶

使用固定内存缓冲区¶

使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel¶

CUDA 图形¶

为什么使用CUDA图？¶

PyTorch API¶

约束条件¶

Non-constraints¶

全网络捕获¶

部分网络捕获¶

与torch.cuda amp的用法¶

多流使用¶

使用DistributedDataParallel¶

NCCL < 2.9.6¶

NCCL ≥ 2.9.6¶

图形内存管理¶

在捕获之间共享内存¶

文档

教程

资源

优化内存使用 `PYTORCH_CUDA_ALLOC_CONF`¶