数据截至 (上游 commit bd2a0fc7c314)
04 · 训练循环与分布式
这一章讲什么: 以
FullFinetuneRecipeDistributed(recipes/full_finetune_distributed.py:53)为主线,精读一份分布式训练脚本的全部关键动作:并行度怎么算、模型怎么在 meta device 上组装再灌权重、训练循环里那几处「乘 token 数再除回来」到底在修正什么、以及三件套显存技术(AC / offloading / optimizer-in-backward)的实现位置。
1. 它要解决的小问题
同一份训练逻辑要覆盖的局面很宽:单卡 bf16、8 卡 FSDP、多机 + 张量并行 + 上下文并行,外加梯度累积、activation checkpointing、断点续训。传统 做法是把这些藏进 Trainer;torchtune 的选择是全部写在 recipe 脚本里,代价是脚本长(1169 行),收益是每个决策点都能指到行号。
2. setup:模型就位的固定顺序
_setup_model(recipes/full_finetune_distributed.py:584)里动作的顺序是刻意安排的,改动顺序就会踩坑(比如先灌权重再分片会浪费几倍内存):
① meta device 实例化(只占结构,不占显存) :611-612
② torch.compile 包一层 :614-615
③ TP:除小头数 + parallelize_module :636-651
④ activation checkpointing(全量或选择性) :659-670
⑤ FSDP 逐层 fully_shard :673-692
⑥ 落到真设备,补建 RoPE 缓存 :694-698
⑦ 完整 state dict → 分片灌入 :702-708
⑧ 挂上 offloading / CP / train 上下文 :711-724
⑨ 校验没有参数遗留在 meta device + barrier :727, :739
三个关键步骤各自展开:
① 为什么用 meta device。 with training.set_default_dtype(self._dtype), torch.device("meta"): model = config.instantiate(cfg_model)(recipes/full_finetune_distributed.py:611-612)。meta device 上的张量只有 shape/dtype 元信息,不分配存储——8B 模型在每张卡上实例化也只占零显存,等 FSDP 分片后才为本分片分配真存储。docstring 里点明动机:「To minimize GPU peak memory」(:598-602)。
⑤ FSDP 的粒度是「层」。 shard_model(torchtune/training/_distributed.py:661)自底向上遍历 named_modules,凡是名字形如 *.layers.{i} 的子模块(判定函数 get_shard_conditions,:612)各包一层 fully_shard,最后对根模型再包一层但设 reshard_after_forward=False(:704-708)。效果:每层参数用完即重新分片(FULL_SHARD 语义),根模块常驻聚合,避免 embedding 这类散件被反复 gather。
⑥⑦ RoPE 与灌权重的衔接。 RoPE 缓存不进 state dict(第 2 章),所以分片完成后先统一 rope_init()(recipes/full_finetune_distributed.py:694-698),再由 load_from_full_model_state_dict(torchtune/training/_distributed.py:336)把完整权重按各参数的 sharding 布局 distribute_tensor 成分片,load_state_dict(..., assign=True) 灌入。
3. ParallelDims:一张 mesh 装下所有并行
ParallelDims(torchtune/training/_distributed.py:62)是个 dataclass,字段即四个并行度:dp_replicate / dp_shard / tp / cp。
- 推断与校验。
dp_shard=-1表示「剩下的都给 DP 分片」:world_size // (dp_replicate * tp * cp)(:78-82);四者乘积必须等于 world_size,否则断言失败(:84-87)。 - 建 mesh 时把子 mesh 预先 flatten 好。
build_mesh(:92)只对 >1 的维建init_device_mesh,然后把常用组合预先铺平:dp(数据加载用)、dp_shard_cp(FSDP 分片用)、dp_cp(:113-129)。recipe 里self.world_mesh["dp_shard_cp"]直接取用(recipes/full_finetune_distributed.py:686-692)。 - 序列长度的整除约束也在这里算。
min_seq_len_divisor(:165-175):开 CP 要求序列长能被2 * tp * cp整除,否则tp;dataloader 的 collate 会按它 pad(recipes/full_finetune_distributed.py:838)。
TP 侧的额外准备是 prepare_mha_for_tp(torchtune/training/_distributed.py:717):把每个 MultiHeadAttention 的 num_heads/num_kv_heads/embed_dim 就地处以 TP 度数(:756-770),因为 TP 之后每张卡只算一部分头——模块属性与切完的张量保持一致。loss 若在 vocab 维被 TP 切分,LinearCrossEntropyLoss.patch_tp_plan 会把输出投影改成 ColwiseParallel 并要求 loss_parallel 上下文(torchtune/modules/loss/cross_entropy_loss.py:72-85)。
4. 训练循环精读:token 数归一化
主循环在 train()(recipes/full_finetune_distributed.py:949)。大部分是每个训练脚本都有的骨架,真正值得细讲的是变长序列下 loss 怎么平均。问题:loss 函数返回的是「本 batch 每 token 的平均」;梯度累积 + 各 batch 有效 token 数不同时,直接平均会把「短 batch」和「长 batch」等权,数学上错了。
torchtune 的解法分四步(recipes/full_finetune_distributed.py:1004-1050):
# 示意,非源码(对应 :1009-1050)
current_num_tokens = (batch["labels"] != loss_fn.ignore_index).sum()
current_loss = loss_step(batch) * current_num_tokens # 平均 → 求和
current_loss.backward() # 梯度累积的是「和」
# ... 累积够 N 步后:
torch.distributed.all_reduce(num_tokens) # 全体 rank 的总 token 数
grad_scaler(model.parameters(), world_size / num_tokens) # 和 → 全局平均
optimizer.step()
为什么乘的是 world_size / num_tokens:每个 rank 的 backward 产出「本地 token 之和」的梯度,FSDP 的梯度归约把各 rank 梯度按 world_size 平均(inferred:基于 FSDP 默认的均值归约语义与 :1034 的注释「Manually scale the gradients from unnormalized loss by total # of tokens」);再乘 world_size / num_tokens 就恰好得到「全体 token 的平均损失」的梯度。scale_grads_(torchtune/training/_grad_scaler.py:43)本身只是 in-place 的 mul_ 封装,并被 torch.compile 包装(recipes/full_finetune_distributed.py:361-365)。
同一机制顺带让 log 出来的 loss 与优化的目标严格一致(running_loss 也按 token 数除,:1069),并且验证循环用同样的 token 加权平均 + all_reduce(recipes/full_finetune_distributed.py:873-915)。
循环里的其他要点:
| 机制 | 位置 | 一句话 |
|---|---|---|
| 梯度累积门槛 | :1027 | (batch_count + 1) % gradient_accumulation_steps == 0 才 step |
| 梯度裁剪 + DTensor | :1041-1048 | 裁剪后的 grad_norm 若是 DTensor 要 full_tensor() 收回 |
| 断点续训 | :846-847、:485-489 | StatefulDataLoader/StatefulDistributedSampler 恢复 sampler 进度;epoch 边界续训时先 list(self._dataloader) 对齐 |
| profiler | :968、:1128 | 默认 DummyProfiler 空操作,配置后才真开 |
| 收尾 | :1141-1144 | 最后一存 full_tensors=True(纯模型、HF 格式,见第 3 章) |
5. 显存三件套
三件技术都能从 config 开关,且两两有明 确的互斥/依赖关系(recipe 的 __init__ 里集中校验,recipes/full_finetune_distributed.py:222-263)。
5.1 activation checkpointing(两版并存)
- 旧版全量 AC:
enable_activation_checkpointing=True时,set_activation_checkpointing以{TransformerSelfAttentionLayer}为 wrap policy,每层前向只存输入、反向时重算(torchtune/training/memory.py:26;recipe 调用在recipes/full_finetune_distributed.py:667-670)。 - 新版选择性 AC:
ac_mode/ac_option控制,apply_selective_activation_checkpointing(torchtune/training/activations.py:65)逐层包checkpoint_wrapper,可以只包每 N 层或只对特定 op 做。代码注释明说两版并存是「for testing and BC purposes」,测试完会清理(recipes/full_finetune_distributed.py:653-658)。
5.2 activation offloading
OffloadActivations(torchtune/training/_activation_offloading.py:24)是 saved_tensors_hooks 的子类:前向时把超过 1KB 的激活异步搬到 CPU(pinned memory + 独立 CUDA stream 与计算重叠),反向时预取回。几个工程细节:
- 输出头例外:给
model.output挂 forward hook,进入 NoOp 上下文— —lm_head 的激活搬过去马上要用回来,纯亏,还会干扰 chunked CE 的启发式(get_act_offloading_ctx_manager,:388,hook 逻辑在:412-431附近)。 - 依赖关系:必须先开 AC 才允许开 offloading(
recipes/full_finetune_distributed.py:260-263)。 - TP 下建议关 stream:
activation_offloading_use_streams与 TP 同开时 recipe 会 warn 不稳定(recipes/full_finetune_distributed.py:244-254)。
5.3 optimizer-in-backward
最激进的一件:不留全量梯度。做法是给每个参数配一个独立 optimizer(recipes/full_finetune_distributed.py:750-754),再给每个参数挂 register_post_accumulate_grad_hook——该参数的梯度一累加完,就地对它 step() + zero_grad()(torchtune/training/memory.py:222-241 的 optim_step)。
代价写死在 __init__ 的校验里:与梯度累积、梯度裁剪互斥(recipes/full_finetune_distributed.py:222-232),因为这两个操作都需要「看到全部参数的完整梯度」。lr scheduler 也要走 OptimizerInBackwardWrapper 的 set_lr_scheduler 包一层(torchtune/training/memory.py:200-220;recipe 侧 recipes/full_finetune_distributed.py:542-544)。
6. LinearCrossEntropyLoss:把 lm_head 搬进 loss
最后一公里的显存优化。vocab 12.8 万时,[b, s, vocab] 的 logits 张量比模型的一整层还大。LinearCrossEntropyLoss(torchtune/modules/loss/cross_entropy_loss.py:19)的拆法:
- 模型不产 logits。
set_model_output把model.skip_output_layer置 True 并把model.output拿过来当自己的投影层(:67-70);模型 forward 只返回 hidden states(TransformerDecoder.unembed,torchtune/modules/transformer.py:678-688)。 - 先过滤再投影。
mask_inputs先把 label 为ignore_index的位置(padding、被 mask 的输入轮次)从 hidden states 里index_select掉(:117-144)——通常能丢掉一大半 token。 - 分块算 CE。 剩下的 hidden 按
num_output_chunks切块,逐块投影 +F.cross_entropy(reduction="sum")累加(compute_cross_entropy:87-114;forward:147-186),最后除以总有效 token 数。
recipe 侧的对接只有两行:loss 是 SFTLoss 实例就调 set_model_output(recipes/full_finetune_distributed.py:432-433),_loss_step 里对非 SFTLoss 的第三方 loss 才 reshape logits(:858-863)。
7. 上下文并行(CP)
开 CP(context_parallel_dim > 1)后,get_context_parallel_manager(torchtune/training/_distributed.py:811)返回一个上下 文管理器,把 batch 的每个 tensor 沿序列维(dim=1)切给 CP 组,注意力内部靠 ring/allgather 轮换 KV。两个防御性限制写在入口处:
- 输入里出现 flex 的 BlockMask 直接报错——CP 尚不支持 sample packing 的 block mask(
:855-858)。 - 非
TransformerDecoder(多模态)直接报错(:847-849)。
训练时它和 SDPA backend 限定叠加:CP 上下文里把 SDPA 后端限制在 flash/efficient/cuDNN(_get_sdpa_context,:784-808),recipe 主循环用 self.train_context(self.context_parallel_manager(...)) 把两层上下文叠起来(recipes/full_finetune_distributed.py:1004-1006;get_train_context 在 torchtune/training/_distributed.py:889)。
8. 关键细节 / 坑
- compile 的粒度是可配的。
compile: True时 model/loss/optimizer_step/scale_grads 分别编译,其中 optimizer_step 默认不开(recipes/full_finetune_distributed.py:343-355);optimizer_in_bwd与编译 optimizer step 互斥(:399-403)。 - MoE + compile 需要
capture_scalar_outputs(:356-358)。 - 多机启动走 torchrun 语义,
tune run --nnodes ... --rdzv_endpoint ...;未指定 rendezvous 时 CLI 强制 standalone 模式,方便同机多任务(torchtune/_cli/run.py:91-96)。 - tokens/s/GPU 的分母是
non_data_parallel_size(tp×cp),即吞吐按 DP 组归一(recipes/full_finetune_distributed.py:1090-1093)。 - 单机版 recipe 是另一个文件。
recipes/full_finetune_single_device.py没有 process group、mesh 与分片的全部代码,optimizer-in-bwd 是其低显存主力——单卡调试先读它,再读本篇主线,坡度更平。
9. 代码地图
| 主题 | 文件路径 | 符号名 |
|---|---|---|
| 分布式 recipe 主线 | recipes/full_finetune_distributed.py | FullFinetuneRecipeDistributed.setup、_setup_model、train、_loss_step |
| 单机 recipe 对照 | recipes/full_finetune_single_device.py | FullFinetuneRecipeSingleDevice |
| 并行度与 mesh | torchtune/training/_distributed.py | ParallelDims、ParallelDims.build_mesh、min_seq_len_divisor |
| FSDP 分片 | torchtune/training/_distributed.py | shard_model、get_shard_conditions |
| 完整权重 → 分片 | torchtune/training/_distributed.py | load_from_full_model_state_dict |
| TP 准备 | torchtune/training/_distributed.py | prepare_mha_for_tp |
| CP 上下文 | torchtune/training/_distributed.py | get_context_parallel_manager、_get_sdpa_context、get_train_context |
| 梯度缩放 | torchtune/training/_grad_scaler.py | scale_grads_ |
| 全量 AC | torchtune/training/memory.py | set_activation_checkpointing |
| 选择性 AC | torchtune/training/activations.py | apply_selective_activation_checkpointing、checkpoint_wrapper |
| activation offloading | torchtune/training/_activation_offloading.py | OffloadActivations、get_act_offloading_ctx_manager |
| optimizer-in-backward | torchtune/training/memory.py | register_optim_in_bwd_hooks、create_optim_in_bwd_wrapper、OptimizerInBackwardWrapper |
| chunked CE | torchtune/modules/loss/cross_entropy_loss.py | LinearCrossEntropyLoss、mask_inputs、compute_cross_entropy、patch_tp_plan |
| loss 接口 | torchtune/modules/loss/loss_types.py | SFTLoss |