生产级分布式训练实战 - 第2节
单机多卡训练基础：从CUDA到DDP的进化之路
大家好，我是你们的老朋友。今天咱们聊聊单机多卡训练的基础。说实话，这块内容我讲了不下几十次，但每次都有新感悟。你想想看，从单卡到多卡，看似只是多了几张卡，但背后的思维模式完全不一样了。
我个人习惯把单机多卡训练分成三个层次：CUDA编程模型是地基，DataParallel是新手村，DistributedDataParallel才是真正的战场。咱们一层层往上走。
CUDA编程模型回顾
先别急着跳过去。我知道很多人觉得CUDA是搞底层的人学的，咱们做训练的不用管。嗯，这个想法我也有过，直到有一次排查一个诡异的显存泄漏问题，折腾了两天，最后发现是CUDA stream没同步。从那以后，我老老实实把CUDA编程模型补了一遍。
核心概念：Host vs Device
说白了，CUDA把计算分成了两个世界：
- Host：CPU端，负责控制逻辑、数据准备
- Device：GPU端，负责大规模并行计算
数据要从Host搬到Device，算完再搬回来。这个搬运过程，就是训练中最容易出性能瓶颈的地方。
Kernel Launch与Stream
CUDA的kernel launch是异步的。什么意思？就是你调用一个核函数，CPU不会等它算完，而是直接往下执行。这本来是好事，能充分利用CPU和GPU的并行性。但问题来了——如果你不小心，就会在数据还没算完的时候去读结果。
内存层次结构
CUDA的内存模型，我建议你记住这张图：
这张图我每次培训都会画。你注意看，越往下容量越小、速度越快。写CUDA代码的核心思路就是：尽量把数据往快了放。
PyTorch DataParallel原理与局限
好，现在咱们进入PyTorch的世界。DataParallel（简称DP）是PyTorch最早提供的多卡方案。用法很简单：
model = nn.DataParallel(model, device_ids=[0,1,2,3])
outputs = model(inputs) # 一行代码搞定多卡是不是很爽？但爽的背后有代价。
DP的工作原理
DP的工作流程，我总结为三步：
- Scatter：把输入batch切分成多个小batch，分发到各GPU
- Replicate：把模型复制到每个GPU上
- Gather & Reduce：各GPU算完梯度后，汇总到主卡（device_ids[0]）
问题就出在第三步。所有梯度都要先汇集到主卡，主卡更新完参数，再广播回去。这意味着：
- 主卡的显存压力巨大（要存所有卡的梯度）
- 通信瓶颈明显（所有流量经过主卡）
- 主卡的计算和通信无法重叠
我在项目中遇到的DP问题
记得有一次，我用4张V100跑一个BERT-large模型。用DP跑，batch size只能开到16，再大主卡就OOM了。后来换成DDP，batch size直接翻到64，训练速度提升了将近3倍。这就是DP的"主卡瓶颈"在作祟。
PyTorch DistributedDataParallel（DDP）原理
DDP才是真正的分布式训练方案。它的核心思想是：各卡独立计算，只在梯度同步时通信。
DDP的核心机制
DDP的工作流程，我习惯用这张图来解释：
看到没？DDP的关键在于AllReduce操作。各卡算完自己的梯度后，通过AllReduce把梯度平均，然后各卡独立更新参数。这样就没有了主卡瓶颈。
DDP的三大优势
| 对比维度 | DataParallel | DistributedDataParallel |
|---|---|---|
| 通信模式 | 主从式（主卡收集再广播） | 对等式（AllReduce） |
| 显存均衡 | 主卡显存占用远高于其他卡 | 各卡显存基本均衡 |
| 扩展性 | 4卡以上性能下降明显 | 可扩展到数百卡 |
| 计算通信重叠 | 不支持 | 支持（反向传播时同步梯度） |
DDP的启动方式：torchrun vs mp.spawn
好，现在咱们聊聊怎么启动DDP。主要有两种方式：torchrun和mp.spawn。
方式一：torchrun（推荐）
torchrun是PyTorch官方推荐的启动方式。用法很简单：
# 训练脚本 train.py
import torch
import torch.distributed as dist
def main():
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
model = MyModel().cuda()
ddp_model = DDP(model, device_ids=[local_rank])
# ... 训练逻辑
if __name__ == "__main__":
main()
# 命令行启动
# torchrun --nproc_per_node=4 train.pytorchrun会自动帮你设置好环境变量：
LOCAL_RANK：当前进程在本机上的编号（0,1,2,3）RANK：全局编号（多机时有用）WORLD_SIZE：总进程数
我个人特别喜欢torchrun的一点是：它帮你处理了失败重启的逻辑。如果某个进程挂了，torchrun会自动重启所有进程。这在长时间训练时特别有用。
方式二：mp.spawn（手动控制）
mp.spawn是更底层的启动方式，适合需要精细控制的场景：
import torch.multiprocessing as mp
def train_worker(local_rank, world_size):
dist.init_process_group(
"nccl",
rank=local_rank,
world_size=world_size
)
torch.cuda.set_device(local_rank)
model = MyModel().cuda()
ddp_model = DDP(model, device_ids=[local_rank])
# ... 训练逻辑
if __name__ == "__main__":
world_size = 4
mp.spawn(train_worker, args=(world_size,), nprocs=world_size)train_worker函数里加上dist.barrier()做同步。有一次我忘了加，结果rank 0都开始算下一个batch了，rank 3还在初始化，直接报错。torchrun默认帮你做了这些，但mp.spawn需要你自己处理。
两种方式的对比
| 特性 | torchrun | mp.spawn |
|---|---|---|
| 易用性 | ⭐⭐⭐⭐⭐ 一行命令启动 | ⭐⭐⭐ 需要手动管理进程 |
| 失败重启 | 自动支持 | 需要自己实现 |
| 多机扩展 | 原生支持（--nnodes） | 需要额外配置 |
| 调试灵活性 | 较低（环境变量自动设置） | 较高（可自定义rank分配） |
好了，单机多卡的基础就聊到这儿。记住一句话：DP是玩具，DDP才是工具。下一章咱们会深入DDP的通信优化，包括梯度压缩、异步通信这些实战技巧。到时候见。