CUDA 语义

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

然而，一旦张量被分配，无论当前选定的设备如何，都可以对其进行操作，结果将始终放在与张量相同的设备上。

默认情况下，跨 GPU 操作是不允许的，但有一些例外，例如 copy_() 以及具有类似复制功能的方法，如 to() 和 cuda()。除非启用了对等内存访问，否则尝试在不同设备上的张量上启动操作将引发错误。

下面是一个小示例，展示了这一点：

在安培（及以后）设备上使用 TensorFloat-32 (TF32)

从 PyTorch 1.7 开始，引入了一个新的标志称为 allow_tf32。这个标志在 PyTorch 1.7 到 PyTorch 1.11 中默认为 True，而在 PyTorch 1.12 及以后版本中默认为 False。这个标志控制 PyTorch 是否允许使用自安培架构以来在 NVIDIA GPU 上可用的 TensorFloat32 (TF32) 张量核心，以内部计算矩阵乘法（包括批量矩阵乘法）和卷积。

TF32 张量核心设计用于通过将输入数据四舍五入到 10 位小数部分，并以 FP32 精度累积结果，从而在 torch.float32 张量上实现更好的矩阵乘法和卷积性能，同时保持 FP32 的动态范围。

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

# 下面的标志控制是否在矩阵乘法中允许使用 TF32。此标志在 PyTorch 1.12 及以后版本中默认为 False。
torch.backends.cuda.matmul.allow_tf32 = True

# 下面的标志控制是否在 cuDNN 中允许使用 TF32。此标志默认为 True。
torch.backends.cudnn.allow_tf32 = True

矩阵乘法的精度可以通过 set_float_32_matmul_precision() 更广泛地设置（不仅限于 CUDA）。请注意，除了矩阵乘法和卷积本身，内部使用矩阵乘法或卷积的函数和 nn 模块也会受到影响。这些包括 nn.Linear、nn.Conv*、cdist、tensordot、affine grid 和 grid sample、adaptive 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()

# 在 TF32 模式下进行矩阵乘法。
torch.backends.cuda.matmul.allow_tf32 = True
ab_tf32 = a @ b  # 在 GA100 上耗时 0.016 秒
error = (ab_tf32 - ab_full).abs().max()  # 0.1747
relative_error = error / mean  # 0.0022

# 禁用 TF32 模式后进行矩阵乘法。
torch.backends.cuda.matmul.allow_tf32 = False
ab_fp32 = a @ b  # 在 GA100 上耗时 0.11 秒
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 GEMM 中的降低精度归约

fp16 GEMM 可能在某些中间步骤中使用降低精度的归约（例如，在 fp16 而不是 fp32 中进行）。这些选择性降低精度的操作可以在某些工作负载（尤其是 k 维度较大时）和 GPU 架构上提高性能，但会牺牲数值精度并增加溢出的风险。

以下是在 V100 上的一些基准测试数据：

如果需要全精度归约，用户可以通过以下方式禁用 fp16 GEMM 中的减少精度的归约：

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

在 C++ 中切换减少精度的归约标志，可以这样做：

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

BFloat16 GEMM 中的减少精度的归约

BFloat16 GEMM 也有类似的标志（如上所述）。请注意，默认情况下此开关对于 BFloat16 是设置为 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()

# 在这里运行一些操作

end_event.record()
torch.cuda.synchronize()  # 等待事件记录完成！
elapsed_time_ms = start_event.elapsed_time(end_event)

例外情况，某些函数（如to() 和 copy_()）接受显式的 non_blocking 参数，这允许调用者在不需要同步时跳过同步步骤。另一个例外是 CUDA 流，如下所述。

CUDA 流

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

每个流中的操作按创建顺序依次执行，但不同流中的操作可以并发执行，且相对顺序不确定，除非使用显式的同步函数（如synchronize() 或 wait_stream()）。例如，以下代码是错误的：

```

plain cuda = torch.device('cuda') s = torch.cuda.Stream() # 创建一个新的流。 A = torch.empty((100, 100), device=cuda).normal(0.0, 1.0) with torch.cuda.stream(s): # sum() 可能在 normal() 完成之前就开始执行！ B = torch.sum(A)

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

plain cuda = torch.device('cuda') s = torch.cuda.Stream() # 创建一个新的流。 A = torch.empty((100, 100), device=cuda).normal(0.0, 1.0) s.waitstream(torch.cuda.defaultstream(cuda)) # 新增！ with torch.cuda.stream(s): B = torch.sum(A) A.recordstream(s) # 新增！

这里有两处新的改动。`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()` 的文档。

请注意，即使在没有读取依赖的情况下，这种同步也是必要的，例如在这个示例中：

plain cuda = torch.device('cuda') s = torch.cuda.Stream() # 创建一个新流。 A = torch.empty((100, 100), device=cuda) s.waitstream(torch.cuda.defaultstream(cuda)) # 依然需要！ with torch.cuda.stream(s): A.normal(0.0, 1.0) A.recordstream(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. 利用梯度进行后续操作

这些操作之间的流语义关系与任何一组操作相同：

plain s = torch.cuda.Stream()

安全，梯度在与 backward() 相同的流环境中使用

with torch.cuda.stream(s): loss.backward() 利用梯度进行后续操作

不安全

with torch.cuda.stream(s): loss.backward() 利用梯度进行后续操作

安全，带有同步机制

with torch.cuda.stream(s): loss.backward() torch.cuda.currentstream().waitstream(s) 利用梯度进行后续操作

安全，初始化初始梯度和调用 backward 在相同的流环境中

with torch.cuda.stream(s): loss.backward(gradient=torch.ones_like(loss))

不安全，初始化初始梯度和调用 backward 在不同的流环境中，且没有同步

initialgrad = torch.oneslike(loss) with torch.cuda.stream(s): loss.backward(gradient=initial_grad)

安全，带有同步机制

initialgrad = torch.oneslike(loss) s.waitstream(torch.cuda.currentstream()) with torch.cuda.stream(s): initialgrad 记录到流 s 中 loss.backward(gradient=initialgrad)

#### BC 注意：在默认流上使用 grads

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

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

只要 `use grads` 在默认流上执行，就是安全的。但在当前的 PyTorch 中，这种模式不再安全。如果 `backward()` 和 `use grads` 处于不同的流上下文中，必须同步这些流：

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

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

## 内存管理

为了更好地了解 CUDA 内存在不同时间点的使用情况，[了解 CUDA 内存使用](../torch_cuda_memory.html#torch-cuda-memory) 描述了用于捕获和可视化内存使用情况的工具。

对于更高级的用户，我们通过 [`memory_stats()`](../generated/torch.cuda.memory_stats.html#torch.cuda.memory_stats) 提供了更全面的内存性能测试。我们还通过 [`memory_snapshot()`](../generated/torch.cuda.memory_snapshot.html#torch.cuda.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 内置的异步分配器](https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-1/)。`cudaMallocAsync` 需要 CUDA 11.4 或更高版本。默认值是 `native`。`backend` 对进程使用的所有设备生效，不能针对单个设备进行设置。

* `max_split_size_mb` 防止原生分配器拆分大于此大小（以 MB 为单位）的块。这可以减少碎片化，并可能使一些临界工作负载在不耗尽内存的情况下完成。性能成本取决于分配模式，范围从‘零’到‘重大’。默认值是无限制，即所有块都可以拆分。[`memory_stats()`](../generated/torch.cuda.memory_stats.html#torch.cuda.memory_stats) 和 [`memory_summary()`](../generated/torch.cuda.memory_summary.html#torch.cuda.memory_summary) 方法对于调优非常有用。此选项应作为最后手段，用于因‘内存不足’而终止且显示大量非活动拆分块的工作负载。`max_split_size_mb` 仅在 `backend` 为 `native` 时有意义。如果 `backend` 为 `cudaMallocAsync`，则 `max_split_size_mb` 将被忽略。

*

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

* ```markdown
  `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>
// 使用 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 << "分配 " << ptr << " " << size << std::endl;
    return ptr;
}

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

这可以通过 torch.cuda.memory.CUDAPluggableAllocator 在 Python 中使用。用户负责提供 .so 文件的路径以及与上述参数匹配的 alloc/free 函数的名称。

import torch

# 加载分配器
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
    'alloc.so', 'my_malloc', 'my_free')
# 切换当前分配器
torch.cuda.memory.change_current_allocator(new_alloc)
# 这将使用新的分配器在设备上分配内存
b = torch.zeros(10, device='cuda')

import torch

# 进行一次初始的内存分配
b = torch.zeros(10, device='cuda')
# 加载分配器
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
    'alloc.so', 'my_malloc', 'my_free')
# 这将导致错误，因为当前的分配器已经实例化了
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()`](../generated/torch.fft.fft.html#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 操作会从多种内核中选择，因此同一个操作可能会根据输入的不同而多次编译。这种编译在每个进程中只会发生一次，但如果使用内核缓存，则只会编译一次。

默认情况下，如果定义了 `XDG_CACHE_HOME`，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()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available)。在以下示例中，`args.device` 生成一个 `torch.device` 对象，该对象可以用于将张量移动到 CPU 或 GPU。

python import argparse import torch

parser = argparse.ArgumentParser(description='PyTorch 示例') parser.addargument('--disable-cuda', action='storetrue', help='禁用 CUDA') args = parser.parseargs() args.device = None if not args.disablecuda and torch.cuda.is_available(): args.device = torch.device('cuda') else: args.device = torch.device('cpu')

注意

在评估给定环境中 CUDA 的可用性（[`is_available()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available)）时，PyTorch 的默认行为是调用 CUDA 运行时 API 方法 [cudaGetDeviceCount](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html#group__CUDART__DEVICE_1g18808e54893cfcaafefeab31a73cc55f)。由于此调用会初始化 CUDA 驱动程序 API（通过 [cuInit](https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__INITIALIZE.html#group__CUDA__INITIALIZE_1g0a2f1517e1bd8502c7194c3a8c134bc3)），如果它尚未初始化，则后续的进程复制在运行过 [`is_available()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available) 后将会因 CUDA 初始化错误而失败。

可以在导入执行 [`is_available()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available) 的 PyTorch 模块之前（或在直接执行它之前）设置环境变量 `PYTORCH_NVML_BASED_CUDA_CHECK=1`，以指示 [`is_available()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available) 使用 NVML 进行评估（[nvmlDeviceGetCount\_v2](https://docs.nvidia.com/deploy/nvml-api/group__nvmlDeviceQueries.html#group__nvmlDeviceQueries_1ga93623b195bff04bbe3490ca33c8a42d)）。如果 NVML 评估成功（即 NVML 发现/初始化未失败），则 [`is_available()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available) 调用不会影响后续的进程复制。

如果 NVML 发现或初始化失败，[`is_available()`](../generated/torch.cuda.is_available.html#torch.cuda.is_available) 将回退到标准的 CUDA 运行时 API 评估，上述的 fork 限制将生效。

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

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

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

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

plain 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` 上下文管理器。

plain print("外部设备编号为 0") # 在设备 0 上（在大多数情况下是默认的） with torch.cuda.device(1): print("内部设备编号为 1") # 在设备 1 上 print("外部设备编号仍为 0") # 在设备 0 上

如果你有一个张量并希望在同一设备上创建一个相同类型的新的张量，可以使用 `torch.Tensor.new_*` 方法（参见 [`torch.Tensor`](../tensors.html#torch.Tensor)）。虽然之前提到的 `torch.*` 工厂函数（[创建操作](../torch.html#tensor-creation-ops)）依赖于当前的 GPU 上下文和你传递的属性参数，但 `torch.Tensor.new_*` 方法会保留张量的设备和其他属性。

这是在创建模块时推荐的做法，特别是在前向传递过程中需要内部创建新张量时。

plain cuda = torch.device('cuda') xcpu = torch.empty(2) xgpu = torch.empty(2, device=cuda) xcpulong = torch.empty(2, dtype=torch.int64)

ycpu = xcpu.newfull([3, 2], fillvalue=0.3) print(y_cpu)

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

ygpu = xgpu.newfull([3, 2], fillvalue=-5) print(y_gpu)

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

ycpulong = xcpulong.newtensor([[1, 2, 3]]) print(ycpu_long)

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

如果你想要创建一个与另一个张量类型和大小相同，并且填充为全1或全0的新张量，可以使用 [`ones_like()`](../generated/torch.ones_like.html#torch.ones_like) 或 [`zeros_like()`](../generated/torch.zeros_like.html#torch.zeros_like) 这两个方便的辅助函数（这些函数还会保留张量的 `torch.device` 和 `torch.dtype`）。

plain xcpu = torch.empty(2, 3) xgpu = torch.empty(2, 3)

ycpu = torch.oneslike(xcpu) ygpu = torch.zeroslike(xgpu)

### 使用固定内存缓冲区

警告

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

从固定（页锁定）内存中复制数据到 GPU 的速度更快。CPU 张量和存储对象提供了一个 [`pin_memory()`](../generated/torch.Tensor.pin_memory.html#torch.Tensor.pin_memory) 方法，该方法返回一个对象的副本，数据被放置在固定内存区域中。

此外，一旦你固定了一个张量或存储，你可以使用异步的 GPU 复制功能。只需在调用 [`to()`](../generated/torch.Tensor.to.html#torch.Tensor.to) 或 [`cuda()`](../generated/torch.Tensor.cuda.html#torch.Tensor.cuda) 时传递一个额外的 `non_blocking=True` 参数。这可以用于将数据传输与计算过程重叠。

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

### 使用 nn.parallel.DistributedDataParallel 而不是多进程或多 GPU 数据并行

大多数涉及批量输入和多个 GPU 的场景应默认使用 [`DistributedDataParallel`](../generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) 来充分利用多个 GPU。

使用 CUDA 模型与 `multiprocessing` 时需要注意一些重要事项；除非严格满足数据处理要求，否则你的程序可能会出现错误或未定义的行为。

推荐使用 [`DistributedDataParallel`](../generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel)，而不是 [`DataParallel`](../generated/torch.nn.DataParallel.html#torch.nn.DataParallel) 进行多 GPU 训练，即使只有一个节点的情况下。

[`DistributedDataParallel`](../generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) 和 [`DataParallel`](../generated/torch.nn.DataParallel.html#torch.nn.DataParallel) 的主要区别在于：[`DistributedDataParallel`](../generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) 使用多进程，每个 GPU 都有一个独立的进程，而 [`DataParallel`](../generated/torch.nn.DataParallel.html#torch.nn.DataParallel) 使用多线程。通过使用多进程，每个 GPU 都有自己的专用进程，这样可以避免 Python 解释器的全局解释器锁（GIL）导致的性能开销。

如果你使用 [`DistributedDataParallel`](../generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel)，你可以使用 `torch.distributed.launch` 工具来启动你的程序，请参阅 [第三方后端](../distributed.html#distributed-launch)。

## CUDA 图

CUDA 图记录了 CUDA 流及其依赖流执行的任务（主要是内核调用及其参数）。有关 CUDA 图的基本概念和底层 CUDA API 的详细信息，请参阅 [CUDA 图入门](https://developer.nvidia.com/blog/cuda-graphs/) 以及 CUDA C 编程指南中的 [图部分](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-graphs)。

PyTorch 支持使用 [流捕获](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#creating-a-graph-using-stream-capture) 技术来构建 CUDA 图，该技术会将 CUDA 流置于 *捕获模式*。在这种模式下，流中的操作会被记录下来而不是立即执行。

在捕获模式下，发送到捕获流的 CUDA 操作不会实际在 GPU 上运行，而是被记录在图中。

捕获完成后，可以 *启动* 图以多次运行 GPU 工作。每次重放都会运行相同的内核和相同的参数。对于指针参数，这意味着使用相同的内存地址，即每次重放时，内核会访问相同的内存位置。通过在每次重放前用新数据（例如，来自新批次的数据）填充输入内存，可以在新数据上重新运行相同的操作。

### 为什么使用 CUDA 图？

重放图牺牲了即时执行的动态灵活性，以换取**大幅减少的 CPU 开销**。图的参数和内核都是固定的，因此图的重放会跳过所有参数设置和内核调度的层次，包括 Python、C++ 和 CUDA 驱动程序开销。在底层，重放通过一次调用 [cudaGraphLaunch](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__GRAPH.html#group__CUDART__GRAPH_1g1accfe1da0c605a577c22d9751a09597) 将整个图的工作提交给 GPU。重放中的内核在 GPU 上执行得更快，但主要的好处是减少了 CPU 开销。

如果你的网络全部或部分是图安全的（通常这意味着静态形状和静态控制流，但请参阅其他[约束条件](#capture-constraints)），并且你怀疑其运行时间至少部分受到 CPU 限制，那么你应该尝试使用 CUDA 图。

### PyTorch API

警告

此 API 仍在测试阶段，未来版本可能会有所变化。

PyTorch 通过一个原始的 [`torch.cuda.CUDAGraph`](../generated/torch.cuda.CUDAGraph.html#torch.cuda.CUDAGraph) 类和两个方便的包装器，[`torch.cuda.graph`](../generated/torch.cuda.graph.html#torch.cuda.graph) 和 [`torch.cuda.make_graphed_callables`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables)，来暴露图。

markdown [torch.cuda.graph`](../generated/torch.cuda.graph.html#torch.cuda.graph) 是一个简单且多功能的上下文管理器，用于捕获其上下文中的 CUDA 工作。在捕获之前，通过运行几次急切模式（eager mode）的迭代来预热要捕获的工作负载。预热必须在一个辅助流（side stream）上进行。因为图在每次重放时都会从相同的内存地址读取和写入，所以在捕获期间必须保持对输入和输出数据张量的长期引用。要在新的输入数据上运行图，将新数据复制到捕获的输入张量中，重放图，然后从捕获的输出张量中读取新的输出。示例：

g = torch.cuda.CUDAGraph()

# 用于捕获的占位符输入
static_input = torch.empty((5,), device="cuda")

# 捕获前预热
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)

# 捕获图
# 为了允许捕获，自动在上下文中将一个辅助流设置为当前流
with torch.cuda.graph(g):
    static_output = static_input * 2

# 用新数据填充图的输入内存以进行计算
static_input.copy_(torch.full((5,), 3, device="cuda"))
g.replay()
# static_output 持有结果
print(static_output)  # 全部为 3 * 2 = 6

# 用更多数据填充图的输入内存以进行计算
static_input.copy_(torch.full((5,), 4, device="cuda"))
g.replay()
print(static_output)  # 全部为 4 * 2 = 8

```

请参阅 [全网络捕获](#whole-network-capture)、[与 torch.cuda.amp 的使用](#graphs-with-amp) 和 [多流使用](#multistream-capture)，了解实际和高级模式。

[`make_graphed_callables`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 更加高级。[`make_graphed_callables`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 接受 Python 函数和 [`torch.nn.Module`](../generated/torch.nn.Module.html#torch.nn.Module)。对于每个传递的函数或模块，它会分别创建前向和后向传递的图。请参阅 [部分网络捕获](#partial-network-capture)。

#### 约束

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

这些约束适用于所有在 [`torch.cuda.graph`](../generated/torch.cuda.graph.html#torch.cuda.graph) 上下文中执行的工作，以及传递给 [`torch.cuda.make_graphed_callables()`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 的任何可调用对象的前向和后向传递中的所有操作。

违反这些规则中的任何一条都可能导致运行时错误：

捕获必须在非默认流上进行。（只有在使用原始的 CUDAGraph.capture_begin 和 CUDAGraph.capture_end 调用时才需要关注这一点。graph 和 make_graphed_callables() 会为你设置一个侧流。）
禁止使用同步 CPU 和 GPU 的操作（如 .item() 调用）。
允许使用 CUDA 随机数生成器操作。如果在图中使用多个 torch.Generator 实例，必须在图捕获之前使用 CUDAGraph.register_generator_state 注册这些生成器。避免在捕获期间使用 Generator.get_state 和 Generator.set_state。相反，应使用 Generator.graphsafe_set_state 和 Generator.graphsafe_get_state 在图上下文中安全地管理生成器状态。这确保了在 CUDA 图中正确运行随机数生成器并管理生成器状态。

违反以下任何一条都可能导致静默的数值错误或未定义行为：

* 在一个进程中，同一时间只能进行一个捕获。

* 在捕获过程中，此进程（无论哪个线程）不得运行任何非捕获的 CUDA 工作。

* CPU 工作不会被捕获。如果捕获的操作中包含 CPU 工作，在重放时这些工作将被忽略。

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

* 禁止使用基于 CPU 或 GPU 数据的动态控制流。

* 禁止使用动态形状。图假设每次重放时，捕获操作序列中的每个张量都具有相同的大小和布局。

* 在捕获过程中使用多个流是允许的，但有一些[限制](#multistream-capture)。

#### 非约束条件

* 一旦捕获，图可以在任意流上重放。

### 整个网络捕获

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

### 部分网络捕获

如果您的网络中某些部分不安全捕获（例如，由于动态控制流、动态形状、CPU 同步或必要的 CPU 端逻辑），您可以以急切模式运行这些不安全的部分，并使用 [`torch.cuda.make_graphed_callables()`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 仅将那些捕获安全的部分图化。

默认情况下，`make_graphed_callables()` 返回的可调用对象支持自动梯度计算，可以直接在训练循环中替代您传递的函数或 `nn.Module` 使用。

`make_graphed_callables()` 内部创建了 [`CUDAGraph`](../generated/torch.cuda.CUDAGraph.html#torch.cuda.CUDAGraph) 对象，运行预热迭代，并根据需要维护静态输入和输出。因此，与使用 `torch.cuda.graph` 不同，您无需手动处理这些操作。

在以下示例中，基于数据的动态控制流意味着整个网络无法一次性被捕获，但 [`make_graphed_callables()`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 使我们能够捕获并运行图安全的部分。

### 与 torch.cuda.amp 的使用

对于典型的优化器，`GradScaler.step` 会同步 CPU 和 GPU，这在捕获期间是禁止的。为了避免错误，可以使用 [部分网络捕获](#partial-network-capture)，或者（如果前向传播、损失计算和反向传播都是安全的）捕获前向传播、损失计算和反向传播，但不捕获优化器的更新步骤。

### 多流使用

捕获模式会自动传播到与捕获流同步的所有流。在捕获过程中，您可以通过向不同的流发出调用来暴露并行性，但总体的流依赖关系图必须在捕获开始后从初始捕获流分支出来，并在捕获结束前重新合并到初始流：

plain with torch.cuda.graph(g): # 在上下文管理器入口处，torch.cuda.current_stream() # 是初始捕获流

# 错误（没有从初始流分支出来或在捕获结束前重新合并到初始流）
with torch.cuda.stream(s):
    cuda_work()

# 正确：
# 从初始流分支出来
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    cuda_work()
# 在捕获结束前重新合并到初始流
torch.cuda.current_stream().wait_stream(s)

注意

为了避免高级用户在 nsight systems 或 nvprof 中查看重放时产生混淆：与急切执行不同，图在捕获过程中将复杂的流 DAG 视为提示，而不是命令。在重放期间，图可能会将独立的操作重新组织到不同的流上或以不同的顺序入队（同时保留您原始 DAG 的总体依赖关系）。

### 与 DistributedDataParallel 一起使用

#### NCCL < 2.9.6

NCCL 版本早于 2.9.6 的情况下不允许捕获集体操作。您必须使用 [部分网络捕捉](#partial-network-capture)，这会将所有归约操作推迟到图外的反向传播阶段。

在将网络包装为 DDP 之前，调用 [`make_graphed_callables()`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 对可图形化的部分进行处理。

#### NCCL >= 2.9.6

NCCL 版本 2.9.6 或更高版本允许在图中进行集体操作。可以采用捕获 [整个反向传播过程](#whole-network-capture) 的方法，但需要三个设置步骤。

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

plain os.environ["NCCLASYNCERRORHANDLING"] = "0" torch.distributed.initprocess_group(…)

2. 在完全反向捕获之前，DDP 必须在侧流环境中构建：

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

3. 在捕获前，您的预热必须至少运行 11 次启用 DDP 的急切迭代。

### 图内存管理

捕获的图每次重放时都会使用相同的虚拟地址。如果 PyTorch 释放了这些内存，后续的重放可能会引发非法内存访问。如果 PyTorch 将这些内存重新分配给新的张量，重放可能会损坏这些张量的值。因此，图使用的虚拟地址必须在多次重放中保持保留。PyTorch 的缓存分配器通过检测捕获过程，并从图专用的内存池中分配内存来实现这一点。私有池会一直存在，直到其对应的 `CUDAGraph` 对象和捕获期间创建的所有张量超出作用范围。

私有池会自动维护。默认情况下，分配器为每次捕获创建一个独立的私有池。如果你捕获多个图，这种保守的方法可以确保图的重放不会相互干扰，但有时会浪费内存。

#### 跨捕获共享内存

为了节省私有池中的内存，[`torch.cuda.graph`](../generated/torch.cuda.graph.html#torch.cuda.graph) 和 [`torch.cuda.make_graphed_callables()`](../generated/torch.cuda.make_graphed_callables.html#torch.cuda.make_graphed_callables) 可以选择让不同的捕获共享同一个私有池。如果知道这些图总是按捕获时的顺序重放，并且永远不会并发重放，则可以安全地让一组图共享一个私有池。

[`torch.cuda.graph`](../generated/torch.cuda.graph.html#torch.cuda.graph) 的 `pool` 参数是一个提示，用于指定使用特定的私有池，可以像下面所示那样在图之间共享内存：

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

(为 g1 和 g2 创建静态输入，并对它们的工作负载进行预热…)

捕获 g1 的计算图

with torch.cuda.graph(g1): staticout1 = g1workload(staticin_1)

捕获 g2，并提示 g2 可能与 g1 共享内存池

with torch.cuda.graph(g2, pool=g1.pool()): staticout2 = g2workload(staticin_2)

将 realdata1 复制到 staticin1

staticin1.copy(realdata_1)

将 realdata2 复制到 staticin2

staticin2.copy(realdata_2)

重放 g1 的计算图

g1.replay()

重放 g2 的计算图

g2.replay() ```

使用 torch.cuda.make_graphed_callables()，如果您希望为多个可调用对象生成图，并且知道它们总是按相同的顺序执行（并且从不并发执行），请将它们按实际工作负载中的运行顺序作为一个元组传递，make_graphed_callables() 将使用共享的私有池来捕获它们的图。

如果在实际工作负载中，您的可调用对象偶尔会以不同的顺序执行，或者它们会并发执行，则不允许将它们作为一个元组传递给单个调用的 make_graphed_callables()。相反，您必须分别对每个可调用对象调用 make_graphed_callables()。