训练框架

状态：入门理解 · 更新于 2026-09-14

训练框架要解决三件事：把模型计算拆到设备上，把数据和梯度高效地传起来，以及在故障、显存和吞吐约束下稳定运行。理解 API 之前，先看一次训练迭代在系统中的完整路径。

一次训练迭代

DataLoader 取 batch
    ↓
数据切分 / 搬运到设备
    ↓
Forward → Loss
    ↓
Backward 计算梯度
    ↓
梯度同步或累积
    ↓
Optimizer 更新参数
    ↓
保存 checkpoint / 记录指标

分布式训练中，各阶段可能跨设备、跨进程甚至跨节点。性能问题既可能来自 Kernel，也可能来自数据读取、Python 调度、通信等待或 checkpoint。

数据并行

数据并行让每个副本持有完整模型、处理不同数据切片，并在反向后同步梯度。典型做法是 AllReduce：每个副本得到相同的聚合梯度，再执行一致的参数更新。

优点是改造相对直接；代价是每一步都要承担梯度通信，模型副本还会重复占用参数和优化器状态显存。

张量并行与流水线并行

策略	拆分对象	主要通信 / 气泡	适用问题
数据并行	Batch	梯度同步	模型能放进单卡、数据量大
张量并行	单层权重或激活	层内 AllReduce / AllGather	单层或模型太大，需要多卡协作
流水线并行	模型层	Stage 间传激活，存在 Pipeline Bubble	层数多、设备数较多
组合并行	多种维度同时拆分	通信拓扑更复杂	大模型训练

并行策略不是越多越好。需要同时考虑 Batch、网络拓扑、通信带宽、显存和容错；跨节点 TP 的通信成本往往比单机互联更敏感。

PyTorch 分布式的最小概念

• Process Group：参与通信的一组进程；
• Rank / World Size：进程身份与总进程数；
• Backend：如 NCCL，负责 GPU 集合通信；
• Collective：AllReduce、AllGather、ReduceScatter、Broadcast 等集体操作；
• DistributedDataParallel（DDP）：在反向过程中进行梯度同步的数据并行封装。

启动方式、环境变量和进程绑定必须与集群编排保持一致。排查卡住问题时，先确认每个 Rank 是否都进入同一个 Collective，再检查网络和 NCCL 日志。

显存与吞吐

训练显存不只有参数：还包括梯度、优化器状态、激活和通信缓冲。混合精度、Activation Checkpointing、梯度累积、ZeRO/FSDP 等方法分别在计算、通信和显存之间做权衡。

吞吐也不能只看 samples/s。应同时记录 step time、tokens/s、MFU/利用率、通信占比、显存峰值和重算开销，确保比较时 Batch、序列长度和精度一致。

故障排查顺序

1. 先确认数据、模型和单卡版本能稳定跑通。
2. 对比单卡与多卡的 step time，拆分计算、通信和输入耗时。
3. 检查 Rank、设备绑定、进程数和 NCCL 拓扑。
4. 观察 OOM 是参数/优化器、激活，还是临时通信缓冲导致。
5. 验证 checkpoint 可恢复，并记录失败前的 step 与日志。

信息来源

• PyTorch Distributed Overview
• PyTorch DistributedDataParallel
• DeepSpeed Documentation
• PyTorch FSDP

延伸阅读

• 推理框架
• 训练到推理的衔接
• 训练框架八股

更新日志

• 2026-09-14：补充训练迭代链路、并行策略与排查框架。