06 · 训练框架：FSDP 与 BF16 优化器

本篇在总分总中是"分"的训练侧深度拆解。DSpark 训练的难点不在算法（已在 03 DSpark 建模 讲清），而在工程：如何用 8 卡 GPU 把 5 层 draft 模型训到收敛、如何在 bf16 下稳定优化、如何应对 hfai 抢占式调度、如何恢复训练。本篇拆解 BaseTrainer 10 步初始化、训练循环、FSDP 分片、BF16 master 权重、checkpoint 保存/恢复。

---

总览段（总）

BaseTrainer（[base_trainer.py:148](file:///workspace/deepspec/trainer/base_trainer.py#L148)）是 DSpark / DFlash / Eagle3 三算法的训练骨架。它抽象了"初始化 → forward → backward → optimizer step → checkpoint"的全部公共流程，子类只实现 _build_draft_model 与 run_batch 两个钩子。

图说明： 初始化 10 步严格顺序依赖：每步用前一步产物。训练循环是经典的"梯度累积 + 同步优化"——should_sync = (next_micro_step + 1) % gradient_accumulation_steps == 0，同步步做 clip+step+checkpoint+suspend 检查。SuspendController 是 hfai 抢占式调度的关键——后台线程轮询 hfai.client.receive_suspend_command()，收到信号就保存 checkpoint 并挂起（[hfai_suspend.py:18-71](file:///workspace/deepspec/utils/hfai_suspend.py#L18-71)）。

关键文件清单：

文件	角色
[deepspec/trainer/base_trainer.py](file:///workspace/deepspec/trainer/base_trainer.py)	BaseTrainer 抽象
[deepspec/trainer/dspark_trainer.py](file:///workspace/deepspec/trainer/dspark_trainer.py)	Qwen3/Gemma4 DSpark trainer
[deepspec/trainer/eagle3_trainer.py](file:///workspace/deepspec/trainer/eagle3_trainer.py)	Qwen3/Gemma4 Eagle3 trainer
[deepspec/trainer/ckpt_manager.py](file:///workspace/deepspec/trainer/ckpt_manager.py)	checkpoint 保存/恢复
[deepspec/utils/optim.py](file:///workspace/deepspec/utils/optim.py)	BF16Optimizer + 调度器
[deepspec/utils/distributed.py](file:///workspace/deepspec/utils/distributed.py)	init_dist + StatelessResumableDistributedSampler
[deepspec/utils/hfai_suspend.py](file:///workspace/deepspec/utils/hfai_suspend.py)	抢占式调度支持

---

分述段（分）

6.1 初始化 10 步详解

BaseTrainer.__init__（[base_trainer.py:151-227](file:///workspace/deepspec/trainer/base_trainer.py#L151-227)）：

1. init_dist(local_rank)：基于 RANK/WORLD_SIZE 环境变量（节点级）+ local_world_size（GPU 数）计算 global rank/world_size，NCCL 后端，默认 60 分钟超时（[distributed.py:12-30](file:///workspace/deepspec/utils/distributed.py#L12-30)）。
2. 设置 precision dtype、checkpoint 目录、SuspendController（[base_trainer.py:154-160](file:///workspace/deepspec/trainer/base_trainer.py#L154-160)）。
3. build_models()：构建 draft 模型 + tokenizer。关键设计：加载完整 target 模型到 CPU 仅为了初始化 draft 的 embedding/lm_head（[base_trainer.py:262-273](file:///workspace/deepspec/trainer/base_trainer.py#L262-273)）：

   draft_model.initialize_embeddings_and_head(
       embed_tokens=target_embed_tokens,
       lm_head=target_lm_head,
       freeze=True  # 冻结
   )
   

   这意味着 DSpark/Eagle3 都从 target 模型继承冻结的 embed_tokens 和 lm_head，论文 Section 3.3 明示：“the draft model shares its embedding layer and language modeling head and keeps them frozen, updating only the backbone drafter, sequential block, and confidence head”。
4. load_resume_draft_model（如有 resume checkpoint）：用 type(draft_model).from_pretrained(resume_checkpoint_dir, ...) 加载。
5. torch.compile(dynamic=True)（可选，DSpark 默认开）：动态 shape 编译。
6. _wrap_with_fsdp：FSDP 包装（详见 6.3）。
7. CacheDataset + validate_train_cache：校验 cache 的 target_layer_ids、hidden_size、target_model_name_or_path 是否与 draft 模型匹配。
8. _compute_training_schedule：计算训练步数（详见 6.2）。
9. BF16Optimizer：构造优化器（详见 6.4）。
10. load_training_state（如有 resume）：加载 optimizer + RNG 状态。

6.2 训练调度

[_compute_training_schedule](file:///workspace/deepspec/trainer/base_trainer.py) （[base_trainer.py:77-135](file:///workspace/deepspec/trainer/base_trainer.py#L77-135)）：

• gradient_accumulation_steps = global_batch_size / (world_size * local_batch_size)
• Qwen3-4B 默认：local_batch_size=1, global_batch_size=512, world_size=8 → grad_accum=64
• 支持 max_train_steps 覆盖 num_train_epochs

local_batch_size=1 是有意为之——DSpark 单样本 anchor 数已达 512，每 anchor 7 个监督位置，单样本就有 3584 个监督位置，batch 维度的并行度足够。

6.3 FSDP 分片策略

[_build_fsdp_kwargs](file:///workspace/deepspec/trainer/base_trainer.py) （[base_trainer.py:55-74](file:///workspace/deepspec/trainer/base_trainer.py#L55-74)）：

图说明： 5 种 FSDP 策略对应不同部署规模。no_shard（默认）适合单节点 8 卡——因为 draft 模型小（5 层 + hidden_size 2560，约 200M 参数），8 卡每卡都能装下完整副本，分片反而增加通信开销。hybrid_shard 适合多节点：节点内 full_shard 节点间 replicate，用 init_device_mesh("cuda", (nodes, devices), mesh_dim_names=("replicate","shard")) 构造 mesh。MixedPrecision(param_dtype, buffer_dtype) 让 forward 在 bf16 下进行，节省显存。

6.4 BF16Optimizer：主从权重同步

[BF16Optimizer](file:///workspace/deepspec/utils/optim.py) （[optim.py:82-142](file:///workspace/deepspec/utils/optim.py#L82-142)）是 SpecForge 风格的 BF16 优化器：

图说明： BF16Optimizer 维护 fp32 master copy of params（[optim.py:94-98](file:///workspace/deepspec/utils/optim.py#L94-98)）。每步：bf16 model grad → 转 fp32 → AdamW.step → fp32 master 回写 bf16 model。这种主从权重同步让 bf16 训练获得 fp32 优化的稳定性，代价是双倍参数显存。AdamW + CosineAnnealingWarmupLR：warmup 阶段线性增长，之后 cosine 退火（[optim.py:64-79](file:///workspace/deepspec/utils/optim.py#L64-79)）。

TwoStageScheduler / WarmupScheduler：可组合的两阶段调度器，正确处理 state_dict 序列化——resume 时调度器状态可恢复。

6.5 训练循环

[train](file:///workspace/deepspec/trainer/base_trainer.py) （[base_trainer.py:348-400](file:///workspace/deepspec/trainer/base_trainer.py#L348-400)）：

1. 计算 remaining samples，构建 dataloader（[base_trainer.py:353-361](file:///workspace/deepspec/trainer/base_trainer.py#L353-361)）。
2. CUDAPrefetcher 包裹 dataloader 实现预取（[base_trainer.py:362](file:///workspace/deepspec/trainer/base_trainer.py#L362)）。
3. with self.suspend_controller.monitoring(): 监听 hfai 抢占信号（[base_trainer.py:365](file:///workspace/deepspec/trainer/base_trainer.py#L365)）。
4. 每个 micro batch：
   • should_sync = (next_micro_step + 1) % gradient_accumulation_steps == 0
   • sync_context = self.model.no_sync() 若不同步（[base_trainer.py:370-371](file:///workspace/deepspec/trainer/base_trainer.py#L370-371)）
   • loss = self.run_batch(batch) / gradient_accumulation_steps; loss.backward()（[base_trainer.py:372-373](file:///workspace/deepspec/trainer/base_trainer.py#L372-373)）
5. 同步步：
   • FSDP.clip_grad_norm_(self.model, max_grad_norm)（[base_trainer.py:379-382](file:///workspace/deepspec/trainer/base_trainer.py#L379-382)）
   • self.optimizer.step()（[base_trainer.py:383](file:///workspace/deepspec/trainer/base_trainer.py#L383)）
   • training_logger.on_optimizer_step 记录 lr、grad_norm、loss、epoch、step（[base_trainer.py:384-391](file:///workspace/deepspec/trainer/base_trainer.py#L384-391)）
   • 每 checkpointing_steps 步 save_and_eval_checkpoint（[base_trainer.py:393-394](file:///workspace/deepspec/trainer/base_trainer.py#L393-394)）
   • 若收到 suspend 信号，_save_and_suspend 保存并挂起（[base_trainer.py:396-398](file:///workspace/deepspec/trainer/base_trainer.py#L396-398)）

6.6 DataLoader 与跨 epoch sampler

[StatelessResumableDistributedSampler](file:///workspace/deepspec/utils/distributed.py) （[distributed.py:62-139](file:///workspace/deepspec/utils/distributed.py#L62-139)）：

• 跨 epoch 流式采样，确定性 shuffle（seed + epoch_idx）
• 支持任意 offset 起始，自动跨 epoch 边界
• _iter_stream（[distributed.py:120-135](file:///workspace/deepspec/utils/distributed.py#L120-135)）：从当前 epoch 偏移开始，yield 完后进入下一 epoch

DataLoader 配置（[base_trainer.py:288-307](file:///workspace/deepspec/trainer/base_trainer.py#L288-307)）：pin_memory=True、drop_last=True、persistent_workers=True、prefetch_factor=4。

6.7 Checkpoint 保存与恢复

[save_checkpoint](file:///workspace/deepspec/trainer/ckpt_manager.py) （[ckpt_manager.py:136-185](file:///workspace/deepspec/trainer/ckpt_manager.py#L136-185)）：

图说明： Checkpoint 保存流程的关键是 FSDP.state_dict_type(FULL_STATE_DICT, offload_to_cpu=True, rank0_only=True) ——把分片参数聚合到 rank0 并 offload 到 CPU，避免 GPU 显存爆炸。主进程剥离 _orig_mod. 前缀（torch.compile 包装残留）后调用 draft_model.save_pretrained，输出 HuggingFace 兼容格式。每 rank 独立保存自己的 optimizer state + RNG 到 training_state.rank{N}.pt，最后创建 step_latest 符号链接（safe_symlink 原子化，[io.py:9-13](file:///workspace/deepspec/utils/io.py#L9-13)）。

Resume 流程：

• load_resume_draft_model（[ckpt_manager.py:64-81](file:///workspace/deepspec/trainer/ckpt_manager.py#L64-81)）：用 from_pretrained 加载 draft。
• load_training_state（[ckpt_manager.py:84-133](file:///workspace/deepspec/trainer/ckpt_manager.py#L84-133)）：加载 optimizer + RNG，验证 saved_rank == global_rank、saved_world_size == world_size、saved_local_batch_size == local_batch_size——resume 时 topology 必须一致。

TrainingResumeState（[ckpt_manager.py:56-61](file:///workspace/deepspec/trainer/ckpt_manager.py#L56-61)）只保留 next_micro_step，作为训练进度的唯一真相源。

6.8 DSpark / Eagle3 Trainer 子类

[Qwen3DSparkTrainer](file:///workspace/deepspec/trainer/dspark_trainer.py) （[dspark_trainer.py:14-39](file:///workspace/deepspec/trainer/dspark_trainer.py#L14-39)）：

• data_collator_cls = CacheCollator
• _build_draft_model：调用 build_qwen3_draft_config → Qwen3DSparkModel
• run_batch：调用 self.model(input_ids, target_hidden_states, loss_mask, target_last_hidden_states)，再 compute_dspark_loss 计算损失

[Qwen3Eagle3Trainer](file:///workspace/deepspec/trainer/eagle3_trainer.py) （[eagle3_trainer.py:16-67](file:///workspace/deepspec/trainer/eagle3_trainer.py#L16-67)）：

• 重写 build_models（[eagle3_trainer.py:19-52](file:///workspace/deepspec/trainer/eagle3_trainer.py#L19-52)），注释"draft head and norm stay frozen / target-independent to match the DSpark setup"
• run_batch：调用 compute_eagle3_loss(model=self.model, batch=batch, ttt_length=draft_model.ttt_length, step_loss_decay=draft_model.step_loss_decay)

6.9 SuspendController：hfai 抢占式调度

[SuspendController](file:///workspace/deepspec/utils/hfai_suspend.py) （[hfai_suspend.py:18-71](file:///workspace/deepspec/utils/hfai_suspend.py#L18-71)）：

• 后台线程轮询 hfai.client.receive_suspend_command()（[hfai_suspend.py:27-34](file:///workspace/deepspec/utils/hfai_suspend.py#L27-34)）
• requested（[hfai_suspend.py:58-66](file:///workspace/deepspec/utils/hfai_suspend.py#L58-66)）：跨 rank broadcast suspend flag
• go_suspend（[hfai_suspend.py:68-71](file:///workspace/deepspec/utils/hfai_suspend.py#L68-71)）：调用 hfai.client.go_suspend()
• 无 hfai 时降级为 no-op（[hfai_suspend.py:38-40](file:///workspace/deepspec/utils/hfai_suspend.py#L38-40)）

这是 DeepSeek 内部 HAI-LLM 集群调度协议的开源版本——当集群需要抢占 GPU 时，训练进程会收到信号，主动保存 checkpoint 后挂起，避免被强杀丢失进度。

---

小结段（总）

BaseTrainer 的设计精髓是"把训练工程的所有痛点抽到基类，让算法子类只关心 forward + loss"。这种抽象让 DSpark / DFlash / Eagle3 三算法共享同一套 FSDP、bf16 master、checkpoint、抢占恢复、跨 epoch sampler 实现，新增算法只需写 _build_draft_model 与 run_batch。

设计要点回顾：

1. 从 target 继承冻结 embed/lm_head：保证 draft 与 target 词表对齐，论文 Section 3.3 明示。
2. local_batch_size=1 + num_anchors=512：anchor 维度提供足够并行度，单样本就够。
3. no_shard 适合小 draft：8 卡都能装下完整副本，分片反而增加通信。
4. BF16 master 权重同步：bf16 forward + fp32 optimize，稳定性与效率兼顾。
5. FULL_STATE_DICT offload 到 CPU：避免 rank0 显存爆炸。
6. resume topology 必须一致：world_size / rank / local_batch_size 都要匹配。
7. hfai 抢占支持：生产训练集群的必备能力。

易踩坑点：

• torch_compile=True 会给 state dict 加 _orig_mod. 前缀，save 时必须剥离。
• resume 时改 world_size 会失败——saved_world_size == world_size 校验不通过。
• local_batch_size 改了不能 resume——saved_local_batch_size == local_batch_size 校验。
• step_latest 是符号链接，删除 checkpoint 要 unlink 而非 rmdir。
• FSDP no_sync 必须在 with 块内，否则梯度同步语义错乱。

延伸阅读：进入 07 评测系统 看 step_latest checkpoint 如何被 evaluator 加载；进入 09 使用指南 案例一看完整训练命令。论文 Section 3.3（Training）描述了"shares embedding layer and language modeling head and keeps them frozen"，见 [DSpark_paper.pdf](file:///workspace/DSpark_paper.pdf)。