AI 在多GPU环境中：理解主机与设备范式

Source: Towards Data Science

• 第 1 部分：Understanding the Host and Device Paradigm — this article
• 第 2 部分：Point‑to‑Point and Collective Operations — coming soon
• 第 3 部分：How GPUs Communicate — coming soon
• 第 4 部分：Gradient Accumulation & Distributed Data Parallelism (DDP) — coming soon
• 第 5 部分：ZeRO — coming soon
• 第 6 部分：Tensor Parallelism — coming soon

Introduction

本指南解释了 CPU 与离散显卡（GPU）协同工作的基础概念。它提供了一个高层次的概览，帮助你建立关于主机‑设备范式的思维模型，重点关注 NVIDIA GPU，因为它们是 AI 工作负载中最常用的显卡。

注意： 集成显卡（例如 Apple Silicon 芯片中的显卡）拥有不同的架构，本篇文章不涉及。

大局观：主机 vs. 设备

最重要的概念是 Host（主机） 与 Device（设备） 之间的关系。

• 你的程序总是从 CPU 开始运行。
• 当你希望 GPU 执行任务（例如，矩阵相乘）时，CPU 会将指令和数据发送到 GPU。

理解这种主机‑设备交互是高效 GPU 编程的基础。

CPU‑GPU 交互

• CPU 发起指令 – 当你的脚本（在 CPU 上运行）执行到面向 GPU 的代码行时（例如 tensor.to('cuda')），它会为 GPU 创建一条指令。
• 指令被排队 – CPU 不会等待；它会把指令放入 GPU 的一个特殊待办列表，称为 CUDA 流（在下一节会详细说明）。
• 异步执行 – CPU 在 GPU 处理排队操作的同时继续执行后续代码行。此 异步执行 对于高性能至关重要，使 CPU 能在 GPU 进行计算时准备下一批数据或执行其他任务。

Source: …

CUDA 流

CUDA 流 是 GPU 操作的有序队列。提交到同一流的操作会 按顺序 执行，一个接一个。不同流中的操作则可以 并发 运行，使 GPU 能够同时处理多个独立的工作负载。

默认情况下，所有 PyTorch GPU 操作都会被加入 当前活动流（通常是 PyTorch 自动创建的默认流）。这简单且可预测：每个操作都会等前一个操作完成后才开始。大多数代码中你不会注意到这一点，但在有可以重叠的工作时，它会导致性能损失。

多流：并发

使用多流的经典场景是 计算与数据传输的重叠。当 GPU 正在处理批次 N 时，你可以同时把批次 N + 1 从 CPU RAM 复制到 GPU VRAM：

Stream 0 (compute): [process batch 0]────[process batch 1]───
Stream 1 (transfer): ────[copy batch 1]────[copy batch 2]───

该流水线之所以可行，是因为计算和数据传输使用 GPU 内部的不同硬件单元，从而实现真正的并行。

在 PyTorch 中，你可以使用上下文管理器创建流并将工作调度到相应的流上：

import torch

compute_stream  = torch.cuda.Stream()
transfer_stream = torch.cuda.Stream()

# -------------------------------------------------
# Transfer work (runs on `transfer_stream`)
# -------------------------------------------------
with torch.cuda.stream(transfer_stream):
    # `non_blocking=True` 让拷贝以异步方式进行
    next_batch = next_batch_cpu.to('cuda', non_blocking=True)

# -------------------------------------------------
# Compute work (runs on `compute_stream`)
# -------------------------------------------------
with torch.cuda.stream(compute_stream):
    # 这段代码会与上面的拷贝并发执行
    output = model(current_batch)

提示： .to() 上的 non_blocking=True 标志至关重要。若不加此标志，即使你希望异步执行，拷贝仍会阻塞 CPU 线程。

流之间的同步

由于流是相互独立的，你必须显式地指示何时一个流依赖于另一个流。

• 全局同步（粗粒度工具）：

  torch.cuda.synchronize()   # 等待设备上 *所有* 流完成

• 细粒度同步 使用 CUDA 事件（精细工具）。
  事件在流中标记一个特定点；另一个流可以在不阻塞 CPU 线程的情况下等待该事件。

  event = torch.cuda.Event()
  
  # -------------------------------------------------
  # Transfer stream: copy data and record an event
  # -------------------------------------------------
  with torch.cuda.stream(transfer_stream):
      next_batch = next_batch_cpu.to('cuda', non_blocking=True)
      event.record()                     # 标记：传输已完成
  
  # -------------------------------------------------
  # Compute stream: wait for the transfer before using the data
  # -------------------------------------------------
  with torch.cuda.stream(compute_stream):
      compute_stream.wait_event(event)    # 只在 GPU 上阻塞此流
      output = model(next_batch)

使用事件比 stream.synchronize() 更高效，因为只有依赖的流会在 GPU 端停下来；CPU 线程仍然可以继续排队其他工作。

何时需要手动管理流？

对于典型的 PyTorch 训练循环，你很少需要直接操作流。然而，许多高级工具——例如 DataLoader(pin_memory=True) 和自定义的预取管线——在内部都利用了该机制。理解流的工作原理可以帮助你：

• 明白这些设置背后的原因。
• 在出现细微性能瓶颈时进行诊断。
• 构建能够重叠计算、数据移动，甚至 kernel 启动的高级管线，以实现最大吞吐量。

Source: …

PyTorch 张量

PyTorch 是一个强大的框架，它抽象了许多细节，但这种抽象有时会掩盖底层实际发生的事情。

当你创建一个 PyTorch 张量时，它由两部分组成：

1. 元数据 – 形状、数据类型、设备等。
2. 数值数据 – 实际存储在设备上的数值。

t = torch.randn(100, 100, device=device)

• 元数据 位于主机的 RAM 中。
• 数据 位于设备的内存（例如 GPU VRAM）中。

为什么这种区分很重要

• print(t.shape) – CPU 可以立即检索形状，因为元数据已经在 RAM 中。
• print(t) – 为了显示张量的内容，PyTorch 必须将数据从 VRAM 传输到主机，这会产生一次设备到主机的拷贝，对于大张量来说代价很高。

主机‑设备同步

从 CPU 访问 GPU 数据会触发 主机‑设备同步，这是常见的性能瓶颈。只要 CPU 需要 GPU 尚未写回主存的结果，就会发生这种情况。

为什么重要

print(gpu_tensor)

gpu_tensor 仍在 GPU 上计算中。CPU 必须等到 GPU 完成计算 并且 将数据从显存复制到内存后才能打印其值。此时 CPU 阻塞（即等待），导致整个程序停滞。

高效的张量创建

要点

每一次同步都会迫使主机和设备相互等待，显著降低吞吐量。优秀的 GPU 编程应 尽量减少这些同步点，让 CPU 与 GPU 都保持忙碌。

“你希望你的 GPU 发出 brrrrr 的声音。”

图片作者：本人（使用 ChatGPT 生成）

扩展规模：分布式计算与 Rank

训练大型模型——尤其是大语言模型（LLMs）——往往超出单个 GPU 的计算能力。要利用多 GPU，需要 分布式计算，随之而来的概念是 rank（进程编号）。

什么是 Rank？

• Rank 是一个独立的 CPU 进程，分配有：
  1. 一个 GPU 设备（例如 cuda:0、cuda:1 …）
  2. 唯一的整数 ID（0, 1, 2, …）

当你在两块 GPU 上启动训练脚本时，会创建两个进程：

每个进程运行自己的 Python 脚本实例。在单机（单节点）上，这些进程共享同一 CPU，但 不共享内存或状态。它们独立运行，仅通过你配置的分布式后端（例如 NCCL、Gloo、MPI）进行协同。

为什么 Rank 很重要

即使每个 rank 执行相同的代码，你仍然可以利用 rank ID 来：

• 划分数据：每个 rank 处理数据集的不同切片。
• 分配角色：rank 0 通常负责日志记录、检查点保存或验证，而其他 rank 只专注于训练。
• 控制设备分配：torch.cuda.set_device(rank) 确保每个进程使用自己的 GPU。

最小示例（PyTorch）

import os
import torch
import torch.distributed as dist

def main():
    # 1️⃣  初始化进程组
    dist.init_process_group(
        backend="nccl",          # NCCL 在 GPU 上表现最佳
        init_method="env://",    # 使用环境变量进行初始化
        world_size=int(os.environ["WORLD_SIZE"]),
        rank=int(os.environ["RANK"]),
    )

    # 2️⃣  为当前 rank 设置设备
    rank = dist.get_rank()
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    # 3️⃣  示例：每个 rank 获取自己的数据切片
    dataset = torch.arange(0, 100)          # 虚拟数据集
    per_rank = len(dataset) // dist.get_world_size()
    start = rank * per_rank
    end = start + per_rank
    local_data = dataset[start:end].to(device)

    # 4️⃣  你的训练循环写在这里 …
    print(f"Rank {rank} handling data {start}:{end} on {device}")

    # 5️⃣  清理
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

启动方式（两 GPU 节点）：

torchrun --nproc_per_node=2 \
    --master_addr=127.0.0.1 \
    --master_port=29500 \
    your_script.py

• torchrun 会自动为每个进程设置 RANK 和 WORLD_SIZE 环境变量。
• 每个进程获得自己的 GPU（rank 0 对应 cuda:0，rank 1 对应 cuda:1），并可以处理各自的数据切片。

小结

• Rank = 独立进程 + 唯一 ID + 专用 GPU。
• 利用 rank ID 来划分工作、管理日志，并保持每块 GPU 高效运行。
• 下一篇博客将进一步探讨数据并行、梯度同步以及实用的规模化训练技巧。

结论

恭喜你阅读到文章的末尾！你已经学到了：

• 主机/设备关系
• 异步执行
• CUDA 流 以及它们如何实现并发 GPU 工作
• 主机‑设备同步

接下来： 在即将发布的博客文章中，我们将深入探讨 点对点 和 集合 操作，这些操作允许多个 GPU 协调复杂的工作流，例如分布式神经网络训练。