是什么
为什么
当单卡显存或算力无法承载模型规模时,工程师与算法研究员必须通过分布式训练突破硬件瓶颈
怎么做
通过框架配置并行策略(如数据并行/张量并行),启动多进程通信并同步梯度
什么时候
模型参数量超过单卡容量或需缩短训练周期时使用;单卡可跑且时间可接受时不应用
FOCUS
先记住这些
不掌握就无法正确启动训练
数据并行适合小模型大数据,张量并行切分大矩阵,流水线并行按层拆分;选错会导致通信爆炸或显存浪费
决定实际加速比而非理论值
梯度同步频率、通信后端选择(NCCL/Gloo)与网络拓扑直接影响吞吐;高频AllReduce会拖慢迭代
最高频OOM根因
激活值峰值、优化器状态与参数分布不均会导致部分GPU提前耗尽显存;需启用梯度检查点或重计算
长周期训练不可跳过
节点故障或网络抖动必发生;必须配置定期保存检查点与自动重启逻辑,否则前功尽弃
PROBLEM / POSITION / INTERFACE
先弄清它为什么存在,以及谁在使用
-
01 解决的问题
模型参数量或激活值超出单GPU显存,或单卡训练耗时过长
盲目多卡导致通信开销爆炸、OOM、梯度不同步或训练中断
- 成功标准
- 线性加速比接近理论值,显存均衡,训练稳定可恢复
-
02 AI 生态位
连接模型架构与底层硬件的中间层,向上承接训练任务,向下调度GPU/NPU集群
- 上游
- 依赖模型结构定义、数据集划分与硬件拓扑
- 下游
- 为预训练、微调等任务提供可扩展的训练能力
-
03 人的生态位
工程师负责策略选择、拓扑配置与异常排查
- 适合使用
- 模型显存需求>单卡容量,或需将训练时间压缩至可接受范围
- 不必使用
- 模型可单卡运行、数据量小、或团队缺乏集群运维能力时
-
04 独特价值
唯一能突破单设备物理限制、实现千亿参数模型训练的工程路径
通信与计算重叠难优化,拓扑感知配置复杂,故障恢复成本高
- 复杂度判断
- 复杂度不是功能数量,而是控制与验证成本。
INTERFACE FLOW
谁在操作,信息怎样流动
-
WHO OPERATES
HUMAN选择并行策略、配置集群拓扑、监控指标与干预故障
SYSTEM执行进程间通信、梯度同步、检查点保存与设备调度
-
INPUT
训练数据批次、模型定义、并行策略配置
-
CONTROL
并行类型、通信后端、批大小、同步频率、容错策略
-
OUTPUT
更新后的模型权重、训练日志、检查点文件;不产生外部服务调用
CODE / PYTHONPyTorch 分布式最小启动
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def setup():
dist.init_process_group(backend="nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
def train():
setup()
model = MyModel().cuda()
model = DDP(model, device_ids=[int(os.environ["LOCAL_RANK"])])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
for batch in dataloader:
optimizer.zero_grad()
loss = model(batch)
loss.backward()
optimizer.step()
dist.destroy_process_group()
if __name__ == "__main__":
train()最短闭环:初始化NCCL后端 → 包装DDP → 正常训练循环 → 销毁组。需通过torchrun启动。
JSON / RESPONSE典型训练日志与状态(示意)
{
"rank": 0,
"world_size": 4,
"backend": "nccl",
"status": "running",
"gpu_memory_used_gb": 18.2,
"gpu_memory_total_gb": 24.0,
"step": 1250,
"loss": 0.342,
"throughput_samples_per_sec": 1280,
"checkpoint_saved_at": "2024-06-15T10:30:00Z"
}PRACTICE
启动前必查清单
✓确认所有节点GPU驱动与CUDA版本一致
✓验证NCCL通信可用(nccl-tests 带宽测试)
✓设置正确的 LOCAL_RANK 与 WORLD_SIZE 环境变量
✓启用梯度缩放(AMP)避免混合精度溢出
✓配置检查点保存路径与频率(至少每2小时一次)
FAQ
常见问题
01数据并行和张量并行该怎么选?+
模型能放进单卡选数据并行;显存不够但层内矩阵大可加张量并行;两者可组合。
选错策略会导致通信爆炸或显存浪费,直接决定训练能否跑通。02最小可用启动方式是什么?+
用 torchrun 启动脚本,设置 NCCL 后端,用 DDP 包装模型,保持训练循环不变。
避免手动管理进程与通信,快速验证多卡可用性。03训练中途节点挂了怎么办?+
依赖定期保存的检查点重启;配置自动重试逻辑,从最近step恢复优化器状态。
长周期训练必遇故障,无容错机制将导致全部算力浪费。04什么时候不需要分布式训练?+
模型可单卡运行、数据量小、或训练时间可接受时,优先单卡简化调试。
分布式引入通信与运维成本,非必要不增加系统复杂度。