训练基础设施与分布式训练

摘要：当单卡放不下模型或训练太慢时，需要分布式训练。本文讲清数据并行、ZeRO（0/1/2/3）、PyTorch FSDP2、张量并行（TP）与流水并行（PP）的分工，给出 DeepSpeed 配置示例、ZeRO 选型决策，以及 OOM / 通信瓶颈的排查。适用对象：需要在 8 卡~多节点集群上训练 7B~70B+ 模型的平台/SRE 工程师。版本：DeepSpeed ≥ 0.14、torch ≥ 2.4（FSDP2）、accelerate ≥ 0.30（[版本相关]）。

适用版本与前提

框架：transformers 4.40+、DeepSpeed 0.14+（ZeRO 各阶段）、PyTorch 2.4+（FSDP2）、accelerate 0.30+
互联：节点内 NVLink（~300+ GB/s），跨节点 InfiniBand（~50 GB/s 量级）——影响并行策略选择
模型示例：Llama-3.1-8B/70B、Mixtral 8x7B（MoE）

核心概念：并行维度的分工

并行方式	切什么	通信	典型用途
数据并行 DP	数据分片，模型全复制	AllReduce 梯度	加速吞吐，基础
ZeRO-1/2/3	优化器态/梯度/参数分片	额外 AllGather	省显存，替代纯 DP
张量并行 TP	层内矩阵按列/行切	高频 AllReduce（节点内快）	单层太大放不下
流水并行 PP	按层切到不同卡	跨 stage 传递激活	超深/跨节点
序列并行 SP	长序列切分	RingAttention	长上下文

ZeRO 三阶段（DeepSpeed 官方定义）

ZeRO-1：只分片优化器状态。每卡显存约降至 1/N（N 为 DP 度）。通信量与 DP 相同。
ZeRO-2：再分片梯度。进一步降显存，通信略增。
ZeRO-3：再分片模型参数，前向/反向时按需 AllGather。显存随 N 线性下降，通信约 1.5× DDP。
ZeRO-Infinity：在 ZeRO-3 基础上把状态卸载到 CPU/NVMe，可训超大模型。

架构与原理：3D 并行拓扑

ZeRO-3 与 FSDP 等价，但别混用

ZeRO-3（FULL_SHARD）在思想上与 PyTorch FSDP 完全等价；FSDP 的 FULL_SHARD≈ZeRO-3、SHARD_GRAD_OP≈ZeRO-2、NO_SHARD≈DDP、HYBRID_SHARD 组内分片组间 DDP。二者不要同时启用——在 accelerate 配置里只选其一。FSDP2（PyTorch 2.4）改为 per-parameter 调度，配合 torch.compile 更激进；DeepSpeed 在 Offload（ZeRO-Infinity / NVMe 卸载）与老代码库生态上仍有优势。

生产实践：怎么选

选型经验

小模型小卡救急：单/少卡放不下时用 ZeRO-3 或 FSDP 救急。
大模型大集群：ZeRO-1/2 + TP + PP 组合拳更优——TP 走节点内 NVLink，ZeRO-3 的 AllGather 跨节点走 IB 更贵。
新项目从 0 起、只用 PyTorch 配方：优先 FSDP（与 transformers/torchtune/trl 集成度高）。
接老代码 / 要 Offload / 对接 NeMo·Megatron：继续用 DeepSpeed。
MoE / 非 Transformer 边角场景：DeepSpeed 历史积累更深。

操作步骤：DeepSpeed ZeRO-2 配置

json

// ds_zero2.json —— 8×A100 训 7B/13B 全量 SFT
{
  "bf16": { "enabled": "auto" },
  "zero_optimization": {
    "stage": 2,
    "contiguous_gradients": true,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": "auto",
    "allgather_bucket_size": "auto"
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": "auto", "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0 }
  }
}

bash

# 用 accelerate 启动（需先 accelerate config 生成加速配置）
accelerate launch --num_processes 8 --mixed_precision bf16 \
  --config_file ds_zero2.json train.py

# 或用 deepspeed 直接启动
deepspeed --num_gpus 8 train.py --deepspeed ds_zero2.json

ZeRO-3 保存权重注意

ZeRO-3 下 state_dict 只含占位（参数被分片）。要导出完整权重需：

json

{ "zero_optimization": { "stage": 3, "stage3_gather_16bit_weights_on_model_save": true } }

再用 deepspeed.save_16bit_model(...)；或用 zero_to_fp32.py（保存 checkpoint 时自动生成）离线合并，无需 GPU。否则只保存分片，无法直接用于推理。

验证

bash

# 1) 启动后确认每张卡都被占用且显存分布均衡
nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv
# [未实测具体输出；应看到 8 卡均有占用]

# 2) 监控通信：若 GPU-util 长期 <50% 且 NCCL 等待高，多半是数据加载或 TP 跨节点
# 3) loss 曲线与单卡 QLoRA 趋势应一致（同数据同超参下）

回滚与清理

大 checkpoint 的回滚

ZeRO-3 分片 checkpoint 占空间大；保留版本化目录，回滚即切换路由指向旧 checkpoint。
删除前确认无推理服务 / 下游微调任务引用；跨节点 NFS 删除大目录注意 IO 风暴，分批删。
用 accelerate / DeepSpeed 的 save_total_limit 限制保留份数，避免存储爆满。

故障排查

OOM（ZeRO-3）：开启 offload_param/offload_optimizer 到 CPU；降 train_micro_batch_size_per_gpu；或改用 ZeRO-2（若参数能放下）。
NCCL 超时 / 卡死：检查 IB 网卡、防火墙、多网卡路由；设 NCCL_SOCKET_IFNAME 与 NCCL_IB_DISABLE 排查；topo 不一致时统一驱动。
负载不均：device_map 或 TP 跨节点导致；TP 应尽量留在单节点 NVLink 内。
精度/收敛异常：确认 bf16 全链路一致，避免 fp16/bf16 混用；开 torch.compile 时先关闭验证基线。

安全与合规

多租户与数据边界

Checkpoint 泄露：全量微调产出的完整权重含全部训练数据痕迹，存储需加密 + ACL，禁止跨租户可读。
网络隔离：分布式训练跨节点走 IB/以太网，需隔离训练 VPC，避免梯度/激活被嗅探（理论风险，取决于部署）。
越权：训练节点常以特权访问对象存储拉权重，凭据用短时效 token，避免长期密钥落地。
合规：集群级训练能耗与数据驻留可能触发合规要求（数据出境、行业监管），上线前核对。

成本与性能（估算，[未实测]）

规模	配置	示例时长	单价假设	估算
7B 全量 SFT	8× A100 80GB	数小时	~$16/h(整机)	$数十
70B 全量	多节点 A100/H100	数十~上百小时	~$数十/h	$数千
70B ZeRO-3+Offload	更少卡更长时	更长	低单价	权衡时间

性能备注

ZeRO-3 通信约 1.5× DDP；FSDP2 在部分场景比 ZeRO-3 快 2–5×/iter（社区基准，[版本相关]，非通用结论）。利用率目标 > 70%，靠 nvidia-smi/DCGM 监控；线性扩展效率（强扩展）低于 80% 说明通信或 IO 瓶颈。Offload 到 CPU/NVMe 能降显存但显著拖慢，仅作「放不下」时的兜底。[时长/单价为估算，非实测报价]

训练基础设施与分布式训练 ​

适用版本与前提 ​

核心概念：并行维度的分工 ​

架构与原理：3D 并行拓扑 ​

生产实践：怎么选 ​

操作步骤：DeepSpeed ZeRO-2 配置 ​

验证 ​

回滚与清理 ​

故障排查 ​

安全与合规 ​

成本与性能（估算，[未实测]） ​

参考资料 ​