1. 多GPU训练到底在解决什么问题
单卡跑模型，最直接的感受就是“等”。等一个epoch跑完，等loss降下去，等验证集指标刷新。尤其是当模型参数量往上走、输入序列变长、batch size被迫调小的时候，单卡的算力天花板很快就摸到了。多GPU训练要解决的核心问题就一个：把原本串行堆在一张卡上的计算量和显存压力，拆到多张卡上并行处理，让整体吞吐提上去，同时保证模型收敛行为跟单卡尽量一致。
我最初接触多GPU训练时，脑子里想的很简单：数据分几份，每张卡跑一份，最后把梯度加起来不就完了。实际动手才发现，这里面牵扯的东西比想象中多。数据怎么切、梯度怎么同步、BN层统计量怎么算、学习率要不要跟着放大、保存模型时听谁的，每一个环节都有坑。这篇内容就是围绕“多GPU训练实现”这个主题，把从单卡到多卡的迁移过程拆开讲清楚，适合已经能跑通单卡训练、想进一步压榨硬件效率的读者参考。
核心关键词会贯穿全文：数据并行、梯度同步、DistributedDataParallel、DataParallel、all_reduce、batch size缩放、学习率调整、显存优化。这些不是孤立的概念，它们串起来就是一条完整的多GPU训练链路。
2. 多GPU并行的几种路线与选型逻辑
2.1 数据并行、模型并行与流水线并行的区别
多GPU训练不是只有一种玩法。按照切分维度不同，大致可以分成三类。
数据并行是最常见的一种。每张卡上放一份完整的模型副本，把一个大batch切成若干小batch分给各卡，各卡独立前向和反向，算完梯度后做一次全局同步，再用同步后的梯度更新各自模型。因为每张卡模型一样，更新后参数也保持一致。它的优势是实现相对简单，适合模型能塞进单卡显存的场景。
模型并行是把模型本身切开，不同层放到不同卡上。比如前几层在卡0，中间层在卡1，最后几层在卡2。前向时数据像流水一样在各卡之间传递。它适合单层就特别大、单卡根本放不下的模型，但实现复杂度高，卡间通信频繁，利用率往往不如数据并行。
流水线并行可以看作模型并行的进阶版，把模型按层切成多个阶段，每个阶段放在不同卡上，同时用micro-batch把流水线填满，减少卡间等待。它兼顾了显存和吞吐，但调度逻辑复杂，对框架支持要求高。
实际项目里，纯数据并行覆盖了大多数场景。模型大到单卡放不下时，才会考虑模型并行或流水线并行，或者两者与数据并行组合使用。
2.2 为什么优先选数据并行
选型时我一般先问两个问题：模型能不能塞进单卡？单卡跑一个step要多久？
如果模型能塞进单卡，只是嫌慢，那数据并行就是首选。它的代码改动量最小，从单卡迁移到多卡通常只需要改数据加载器、包一层并行模块、调整学习率这几步。而且主流框架对数据并行的支持最成熟，梯度同步、通信后端、混合精度这些都有现成方案。
如果模型塞不进单卡，那才需要动模型并行或流水线并行的脑筋。这两条路对代码侵入性大，调试成本高，不是万不得已不建议一上来就搞。
提示：不要为了用多卡而用多卡。如果单卡batch size已经很大、GPU利用率也接近满载，多卡带来的收益可能被通信开销吃掉。先确认瓶颈确实在算力，再上多卡。
2.3 DataParallel与DistributedDataParallel的取舍
数据并行在实现层面有两个典型代表：DataParallel（简称DP）和DistributedDataParallel（简称DDP）。
DP是单进程多线程模式，主卡负责把batch切分、分发到各卡，收集各卡输出，再统一算loss和反向，最后把梯度广播回去。它的优点是改起来快，一行包装就能跑。缺点是主卡负载重、通信效率低、Python GIL限制多线程并行，而且只支持单机多卡。实测下来，DP在卡数增多时加速比会明显掉下来。
DDP是多进程模式，每张卡一个独立进程，各自持有完整模型副本，通过通信后端做梯度all_reduce。它没有主卡瓶颈，支持单机多卡和多机多卡，扩展性好得多。代价是启动方式变了，要用torchrun或multiprocessing.spawn，数据采样器要加DistributedSampler，日志和模型保存要只在主进程做。
我的建议很直接：只要不是临时跑个demo，一律用DDP。DP那点上手便利，在性能和扩展性面前不值一提。
| 对比项 | DataParallel | DistributedDataParallel |
|---|---|---|
| 进程模型 | 单进程多线程 | 多进程 |
| 通信效率 | 较低，主卡瓶颈 | 较高，ring all_reduce |
| 扩展性 | 仅单机多卡 | 单机/多机均可 |
| 代码改动 | 小 | 中等 |
| 推荐场景 | 临时验证 | 正式训练 |
3. 从单卡到多卡的核心改造点
3.1 启动方式与进程组初始化
DDP的第一步是让每张卡知道自己是谁、总共几张卡、跟谁通信。这靠环境变量和进程组初始化完成。
用torchrun启动时，框架会自动注入LOCAL_RANK、RANK、WORLD_SIZE这些环境变量。代码里需要读取它们并初始化进程组：
import os import torch import torch.distributed as dist def setup_distributed(): dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) return local_rankbackend选nccl，这是NVIDIA GPU上做集合通信的标准后端，all_reduce、broadcast这些操作都走它。如果是在CPU上调试，可以换成gloo，但正式训练一定用nccl。
启动命令大致长这样：
torchrun --nproc_per_node=4 train.py --batch_size 64 --lr 0.01nproc_per_node等于本机GPU数量。如果只指定部分卡，可以用CUDA_VISIBLE_DEVICES限制可见设备，再让nproc_per_node匹配可见卡数。
注意：init_process_group必须在任何CUDA操作之前调用，set_device也要尽早执行。顺序错了会出现进程组初始化失败或者所有进程都跑到卡0上的情况。
3.2 模型包装与梯度同步机制
模型要包一层DDP：
model = MyModel().to(local_rank) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False )find_unused_parameters这个参数值得单独说。默认False时，DDP假设所有参数都会参与反向传播。如果模型里有分支结构导致某些参数在某些step不参与计算，不设True会报错。但设True会带来额外的图遍历开销，拖慢训练。所以正确做法是先确认模型有没有未使用参数，没有就保持False。
梯度同步发生在反向传播过程中。DDP在构造时会给每个参数注册hook，反向算到某层梯度时，hook触发all_reduce，把各卡该参数的梯度求平均。这样每张卡上的梯度就是全局一致的，更新后参数自然保持一致。
all_reduce的底层实现通常是ring all_reduce，每张卡只跟相邻卡通信，带宽利用充分，卡数越多优势越明显。
3.3 分布式采样器与数据切分
单卡时DataLoader的shuffle是全局打乱。多卡时如果还用普通sampler，每张卡会拿到完整数据集，等于重复训练。正确做法是用DistributedSampler：
from torch.utils.data.distributed import DistributedSampler train_sampler = DistributedSampler( train_dataset, num_replicas=world_size, rank=rank, shuffle=True ) train_loader = DataLoader( train_dataset, batch_size=batch_size_per_gpu, sampler=train_sampler, num_workers=4, pin_memory=True )DistributedSampler会把数据集按rank切分，每张卡只看到自己那份。每个epoch开始前要调用train_sampler.set_epoch(epoch)，否则每轮切分方式一样，shuffle效果打折。
这里有个容易忽略的点：batch_size_per_gpu是每卡batch size，全局batch size等于它乘以卡数。后面调学习率时要按全局batch size来算。
3.4 日志、保存与评估的主进程控制
多进程下如果每个进程都写日志、存模型，会互相覆盖，日志也会重复。标准做法是用rank判断：
is_main = (rank == 0) if is_main: logger.info(f"epoch {epoch} loss {loss.item()}") if is_main: torch.save(model.module.state_dict(), "ckpt.pt")注意保存时要取model.module，因为DDP包装后原始模型在module属性里。评估阶段也要小心，如果每张卡都跑完整验证集，不仅浪费还算重复。常见做法是只在主进程评估，或者各卡评估自己那份再汇总。
4. 超参数调整与显存优化实操
4.1 学习率与batch size的缩放关系
从单卡换到多卡，全局batch size变成原来的N倍（N为卡数）。学习率要不要跟着变，取决于优化器和训练策略。
经验规则是线性缩放：全局batch size扩大N倍，学习率也扩大N倍。这个规则在小batch size范围内比较稳。但batch size很大时，线性缩放会导致初期不稳定，这时用平方根缩放更保守：学习率乘以sqrt(N)。
实际操作中我会先按线性缩放设一个值，跑几百step看loss曲线。如果初期loss震荡明显，就降到平方根缩放的水平。另外warmup阶段可以适当拉长，给大batch一个缓冲。
提示：学习率调整后，weight decay、梯度裁剪阈值这些也要重新审视。大batch下梯度噪声更小，梯度裁剪的阈值可以适当放宽。
4.2 梯度累积与显存占用的平衡
多卡不等于显存就够用。每张卡上的模型副本、激活值、梯度、优化器状态都要占显存。如果单卡batch size调大后OOM，有两个方向：一是用梯度累积模拟大batch，二是上混合精度和梯度检查点。
梯度累积的做法是多次前向反向后再step一次：
for i, batch in enumerate(loader): loss = model(batch) / accum_steps loss.backward() if (i + 1) % accum_steps == 0: optimizer.step() optimizer.zero_grad()这样等效batch size等于per_gpu_batch乘以卡数再乘以accum_steps。代价是训练速度变慢，因为step次数没减少，只是每次step前多攒了几轮梯度。
混合精度用torch.cuda.amp，前向用autocast，反向用GradScaler。实测能省30%到50%显存，速度也有提升。梯度检查点则是用计算换显存，把中间激活值丢掉，反向时重算，适合层数特别深的模型。
4.3 通信开销的隐藏与重叠
多卡训练里通信是绕不开的成本。梯度all_reduce的量等于模型参数量，模型越大通信越久。如果通信和计算不能重叠，卡就会空等。
DDP默认会把梯度按bucket分组，反向传播算出一个bucket的梯度就立刻触发该bucket的all_reduce，实现通信与反向计算的重叠。bucket大小可以通过bucket_cap_mb调整，默认25MB。网络带宽紧张时可以调大，让通信更连续；带宽充裕时可以调小，让重叠更细。
另一个手段是增大batch size，让计算时间变长，相对稀释通信占比。但batch size不能无限大，还要考虑收敛。
| 优化手段 | 显存收益 | 速度影响 | 适用场景 |
|---|---|---|---|
| 混合精度 | 高 | 提速 | 大多数GPU训练 |
| 梯度累积 | 高 | 变慢 | 显存不足但想要大batch |
| 梯度检查点 | 高 | 变慢 | 深层模型 |
| bucket调优 | 无 | 提速 | 通信瓶颈明显 |
5. 常见报错与排查思路
5.1 进程组初始化失败
报错信息里出现“Address already in use”或者“Connection refused”，多半是端口冲突或MASTER_ADDR、MASTER_PORT设置有问题。torchrun会自动选端口，手动用spawn时要指定一个空闲端口。多机训练时MASTER_ADDR要设成主节点可达的地址。
还有一种情况是NCCL初始化超时，通常跟网卡、驱动或者卡间拓扑有关。可以先跑一个最小的all_reduce测试脚本，确认通信链路本身没问题，再排查训练代码。
5.2 梯度不同步或loss异常
如果发现各卡loss差异很大，或者训练完全不收敛，先检查DistributedSampler有没有正确set_epoch，再确认模型是不是在DDP包装前就to(device)了。顺序应该是先to(device)再包DDP。
另一个常见原因是find_unused_parameters设置不当。模型有分支时设False会报“Expected to have finished reduction”，设True又可能因为某些参数确实没参与而静默不同步。稳妥做法是打印出未使用参数列表，确认是否符合预期。
5.3 显存OOM的定位方法
OOM不一定是batch size太大。可以用torch.cuda.memory_summary()看显存分布，确认是模型参数、激活值还是优化器状态占了大头。如果是激活值占大头，优先上混合精度或梯度检查点；如果是优化器状态占大头，考虑换用Adafactor这类省显存的优化器。
还有一个隐蔽的坑：验证阶段如果忘了torch.no_grad()，会构建计算图，显存暴涨。训练循环里也要注意及时del中间变量、清空缓存。
| 报错现象 | 可能原因 | 排查方向 |
|---|---|---|
| 进程组初始化超时 | 端口/地址/网卡 | 检查环境变量与通信链路 |
| loss各卡不一致 | 采样器或同步问题 | 检查set_epoch与DDP包装顺序 |
| OOM | batch过大或图未释放 | memory_summary定位 |
| 训练变慢 | 通信瓶颈 | 调bucket、查带宽 |
6. 我踩过的几个坑和实操心得
第一个坑是忘了set_epoch。有次训练loss一直不降，查了半天以为是学习率问题，最后发现DistributedSampler每轮切分一样，模型看到的样本顺序固定，变相减少了数据多样性。加上set_epoch后立刻正常。
第二个坑是保存模型时存了DDP包装后的state_dict。加载时key多了一层module前缀，跟单卡模型不兼容。正确做法是存model.module.state_dict()，加载时用原始模型结构。
第三个坑是验证集重复计算。早期图省事让每张卡都跑完整验证集，结果评估时间翻了几倍，指标还因为重复统计出现偏差。后来改成只在主进程评估，或者各卡算自己那份再all_reduce汇总。
第四个坑是学习率没跟着batch size调。从单卡batch 32换到4卡batch 128，学习率还是原来的值，收敛慢了一大截。按线性缩放调到4倍后，收敛速度才回到预期。
最后分享一个调试技巧：先用单卡跑通，确认loss曲线正常，再切多卡。多卡出问题时，把卡数降到2，用最小数据集跑几个step，逐步加卡、加数据，定位问题比一上来就4卡8卡高效得多。多GPU训练本身不复杂，复杂的是那些藏在细节里的同步和一致性问题，把每个环节都验证一遍，基本就稳了。
