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 使用从 Ampere 开始在 NVIDIA GPU 上可用的 TensorFloat32 (TF32) 张量核心来在内部计算矩阵乘法（矩阵乘法和批处理矩阵乘法）和卷积。
TF32 张量核心旨在通过将输入数据舍入为具有 10 位尾数并在保持 FP32 动态范围的情况下以 FP32 精度累积结果，从而在 torch.float32 张量上的矩阵乘法和卷积中实现更好的性能。
矩阵乘法和卷积分别控制，可以通过以下方式访问其对应的标志
# 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
还可以通过 set_float_32_matmul_precision() 更广泛地设置矩阵乘法的精度（不仅限于 CUDA）。请注意，除了矩阵乘法和卷积本身之外，内部使用矩阵乘法或卷积的函数和 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 矩阵乘法吞吐量的比率可能会因代而异，或者模型而异。如果需要完整的 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 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 中的降精度缩减¶
BFloat16 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）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() 调用确保在我们在侧流上开始运行 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 = 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引擎始终将默认流与所有反向操作同步，因此以下模式
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);
}
}
这可以通过 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 张量上执行 torch.special.zeta 等操作时，PyTorch 会对其进行即时编译。此编译过程可能非常耗时（取决于您的硬件和软件，最多可能需要几秒钟），并且对于单个运算符可能会发生多次，因为许多 PyTorch 运算符实际上是从各种内核中选择的，每个内核必须编译一次，具体取决于它们的输入。此编译每个进程发生一次，或者如果使用内核缓存，则仅发生一次。
默认情况下，如果定义了 XDG_CACHE_HOME，PyTorch 会在 $XDG_CACHE_HOME/torch/kernels 中创建内核缓存；如果未定义，则在 $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 初始化错误而失败。
您可以在导入执行 is_available() 的 PyTorch 模块之前（或在直接执行它之前）在环境中设置 PYTORCH_NVML_BASED_CUDA_CHECK=1，以便指示 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 环境标志来管理哪些 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 返回放置在固定内存中的批次。
使用 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 图？¶
重放图形会牺牲典型急切执行的动态灵活性，以换取**大大减少的 CPU 开销**。图形的参数和内核是固定的，因此图形重放会跳过所有层的参数设置和内核调度，包括 Python、C++ 和 CUDA 驱动程序开销。在底层，重放通过单个调用将整个图形的工作提交给 GPU，以使用 cudaGraphLaunch。重放中的内核在 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实例时，必须在图形捕获之前使用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() 返回的可调用对象是自动梯度感知的，并且可以在训练循环中用作函数或 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¶
早于 2.9.6 的 NCCL 版本不允许捕获集体操作。您必须使用部分网络捕获，它会将 allreduce 推迟到反向的图形化部分之外进行。
在使用 DDP 包装网络*之前*，在可绘制图形的网络部分上调用 make_graphed_callables()。
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 的 Eager 模式迭代。
图内存管理¶
捕获的图每次重放时都会作用于相同的虚拟地址。如果 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()
使用 torch.cuda.make_graphed_callables() 时，如果您想绘制多个可调用对象的图形，并且您知道它们将始终以相同的顺序运行（并且永远不会并发运行），请按照它们在实时工作负载中运行的相同顺序将它们作为元组传递，并且 make_graphed_callables() 将使用共享的私有池捕获它们的图形。
如果在实时工作负载中，您的可调用对象的运行顺序偶尔会发生变化，或者它们将并发运行，则不允许将它们作为元组传递给 make_graphed_callables() 的单个调用。相反，您必须为每个可调用对象单独调用 make_graphed_callables()。
