KuangjuX’s Notes

外观

PithTrain 代码解读与训练系统学习路线 ​

基于 KuangjuX/pith-train（fork 自 mlc-ai/pith-train，commit 889a408，2026-07）逐文件阅读整理。 配套论文笔记见 pithtrain.md。

0. 代码库总览 ​

全库约 11.1K 行 Python（不含测试），三层架构：

pithtrain/
├── tasks/       # Application: 可启动的入口（pretrain_lm.py 510 行）
├── contexts/    # 全局运行时状态（模块级单例，总共不到 100 行）
├── dualpipe/    # 引擎核心: DualPipeV 流水线调度器 (~2050 行)
├── models/      # 每个模型家族一个自包含文件 (~1720 行)
├── modules/     # 分布式 + 训练基础设施 (~1700 行)
└── operators/   # Triton/库 backed 算子 (~4000 行)
examples/        # 每个模型: script.py(配置) + config.json；launch.sh(torchrun/SLURM)
benchmarks/      # 算子 micro-bench + 论文用预训练配置
.agents/skills/  # 9 个 agent 技能
tools/           # memory_estimator 显存估算器
docs/            # user-guide.md / architecture.md（官方开发者导览，值得先读）

核心文件行数（学习优先级排序）：dualpipe/execution.py 829、dualpipe/dualpipev.py 655、modules/training.py 478、tasks/pretrain_lm.py 510、models/qwen3_moe.py 337、operators/ep_dispatch.py 676、operators/ring_attention.py 990、operators/deepgemm_quantize.py 899。

环境要求：Hopper(SM90)/Blackwell(SM100) GPU、CUDA ≥ 13.0、Python ≥ 3.12，用 uv 管依赖（uv venv && uv pip install .，开发用 uv sync）。

1. 核心抽象逐个解读 ​

1.1 contexts —— 全局状态的"单例模式"（pithtrain/contexts/） ​

distributed、training、logging 三个模块就是三个命名空间对象，setup_* 函数往里填字段，全库直接 from pithtrain.contexts import distributed 读 distributed.pp_rank / training.model。这是 PithTrain 消除间接层的关键设计之一：没有依赖注入、没有 registry，状态就放在全局，跳转定义一击即中。

1.2 4D Device Mesh —— 一切的起点（modules/distributed.py，170 行） ​

setup_device_mesh 建一个 (PP, DP, CP, EP) 四维 mesh，DP 不需要配置，由 world_size 除出来。轴序外→内为 PP→DP→CP→EP，注释写明动机：CP（ring K/V 交换）和 EP（MoE all-to-all）通信最频繁，放最内层让流量尽量留在 NVLink 域内。

world_size = pp_size × dp_size × cp_size × ep_size   (dp_size = world_size / 其余)

另有两个值得学的工程细节：fail-fast excepthook（未捕获异常直接 os._exit(1)，绕过 NCCL drain 防止整机挂死）；recompile_limit=64（配合 5-stage 的多入口编译）。

1.3 模型协议 —— 5-stage 切分（models/interface.py，116 行） ​

全库唯一的抽象接口。每个 decoder 层实现：

python
class LayerProtocol(Protocol):
    def forward_stage1(hidden_states, rotary_posemb, cu_seqlens)  # LN+Attn+LN+topk路由 → RoutingInfo
    def forward_stage3(gathered_tokens, expert_idxs, expand_idx)  # 专家计算
    def forward_stage5(moe_outs, moe_local_idxs, topk_weight, residual)  # 加权聚合+残差
    def reference_forward(hidden_states, ...)                     # 不切分的朴素参考实现

Stage 2（dispatch all-to-all）和 Stage 4（combine all-to-all）是框架持有的——模型只返回路由元数据 RoutingInfo，调度器驱动通信。这就是论文说的"每层在 EP 通信边界切 5 刀"。reference_forward 是自校验基准：分布式实现必须逐数值对齐它（tests/ 就是这么测的），这招在自己写训练系统时极 worth 抄。

1.4 以 Qwen3 MoE 为例读一个模型（models/qwen3_moe.py，337 行） ​

  • V 形双 chunk：Qwen3MoeModel(config, phase=0/1)，rank r 持有 chunk r 和 chunk 2·pp_size-1-r（layer_partition 切层，边缘 chunk 更轻因为要背 embedding/lm_head）。这就是 DualPipeV 的 "V"。
  • Gate（Qwen3MoeGate）：softmax→topk→归一化权重；负载均衡 loss 通过 MoELoadBalanceLossInjector.apply(topk_weight, lb_loss * n_tokens) 以自定义 autograd 函数注入到 topk_weight 的梯度通路，绕过 loss 直接把 lb 梯度混进反向——这样 train_step 统一除以 token 数后 lb 梯度仍然正确归一化。这是一个很聪明的 trick。
  • Experts（Qwen3MoeExperts）：三个 GroupedLinear（gate/up/down），silu_mul 融合激活，按 grouped_mm_offs 分组 GEMM。每个 rank 只持有 num_experts // ep_size 个专家。
  • Stage 切分点：stage1 结束于 prepare_dispatch（本地先做 token 去重），stage3 是 scatter→grouped GEMM→gather，stage5 是 topk 加权 scatter_add + 残差。stage1/5 包了 @torch.compile(fullgraph=True)，stage3 故意不包（数据依赖 shape）。
  • zigzag CP：forward_posemb 里 position_ids 按"前块+镜像后块"拼接，与 pretrain_lm.get_global_batch 的数据切分、ring_attention.py 的注意力实现三方严格对齐——CP 是最容易写错的功能，PithTrain 用三方共守一个约定来保证正确。

1.5 DualPipeV 调度器（dualpipe/，共 ~2050 行，引擎心脏） ​

  • dualpipev.py：V 形调度本体，8 步调度算法改编自 DeepSeek DualPipe（MIT），原创新增：5-stage 重叠循环、FSDP2 集成、FP8 weight cache、ChunkRecord 零分配复用。
  • overlap.py（331 行）：overlapped_forward_backward —— 真正把"一个 micro-batch 的计算"与"另一个 micro-batch的 stage2/4 通信"重叠起来的主循环。
  • execution.py（829 行）：chunk 生命周期管理、model_forward/model_backward、显式 post_backward 调用（FSDP hook 在流水线循环里被抑制，循环结束后手动触发规约分片）。
  • utils.py：WeightGradStore（wgrad delay：把反向权重梯度缓存到 stage5 之后才放行，用 micro-batch 重叠掉 wgrad 的通信）；FP8WeightCacheControl（跨 micro-batch 缓存 FP8 量化权重，避免重量化）；gather/scatter（all-gather 参数进 forward）。
  • comm.py：P2P 批量收发（预分配 buffer + bf16 固定 shape，set_p2p_tensor_shapes/dtype 在建模型时设定）。

1.6 训练基础设施（modules/training.py，478 行） ​

  • FSDP2 组装（apply_fsdp）：MoE 专家沿 dp×cp mesh 分片（每个 EP rank 持有不同专家，本来就是复制的），非 MoE 参数沿 dp×cp×ep 分片；逐层 fully_shard(reshard_after_forward=False)（层内多次用，摊薄 all-gather）；支持 fsdp/hsdp 两种策略。
  • 优化器：make_muon_optimizer 把参数按 is_muon_param 分流——2D hidden 权重走 Muon，embedding/lm_head/router/bias 走 AdamW（Moonlight 论文的可扩展 Muon 配方）；也提供纯 AdamW。
  • 调度器：WSD（warmup-stable-decay，cosine/linear 两种衰减）+ constant。
  • 初始化：输入层 N(0, 0.02)，o_proj/down_proj 按 1/sqrt(2L) 缩放（GPT-2 传下来的残差流方差控制）。

1.7 训练主循环（tasks/pretrain_lm.py，510 行） ​

launch → setup 三连 → 循环 train_step。值得精读的段落：

  1. get_global_batch：PP rank 0 负责取数，显式算出 DP×EP 各 rank 的样本区间 + zigzag CP 前后块，单次 HtoD 传输。
  2. criterion 返回 sum 而非 mean，反向累积 d(sum loss)；step 末统一除以全局非忽略 token 数 → 任意 micro-batch 划分都得到正确的 token 加权均值。token 计数在 PP rank 0，broadcast 到所有 stage。
  3. clip_grad_norm_：本地 L2 → all-reduce 平方和 → 全局范数，处理 DTensor（.to_local()）。
  4. checkpoint（DCP）：to_canonical_model 把 FQN 变成 PP 无关的全局层号 + 展开堆叠专家 → checkpoint 可跨并行度 reshard；从 HF 转来的 ckpt 只有 model 时用 model_only 跳过 optimizer 键。
  5. 每 step 打点 loss/lb_loss/grad_norm/tokens-per-second/peak-mem，可选 wandb；nsys 与 memory snapshot 的起停步数都做成配置（nsys_start/stop、memory_profile_start/stop）。

1.8 算子层（operators/，性能来源） ​

文件干什么亮点
ep_dispatch.py 676EP dispatch/combine3 个 Triton kernel 替换 ~22 个小 PyTorch 算子；token 去重（一个 token选的多个专家在同卡只发一次）；counting sort O(n) 替 argsort
token_scatter.py 384分组 GEMM 前后 scatter/gatherpadded_index_gather、预分配 pinned buffer
grouped_linear.py 284BF16 grouped GEMM / FP8(DeepGEMM)同一接口两条后端
deepgemm_quantize.py 899FP8 128 元素块量化Hopper/Blackwell 两条路径，power-of-2 scales
ring_attention.py 990CP 的 K/V 环交换注意力zigzag 布局负载均衡
linear.py 208FP8/BF16 Linearweight cache 控制
silu_mul.py/clamped_swiglu.py融合激活SwiGLU 前后向融合
cross_entropy.py 156融合 CEvocab 级 logits 的省显存实现

modules/load_balance.py（274 行）：三种 MoE 均衡 loss——micro-batch（Switch 式）、global-batch（DeepSeek 式，DP×EP 频率同步）、sequence 三选一。

1.9 其余部分 ​

  • tasks/tokenize_corpus.py：HF tokenizer → memmap packed 序列（workspace/datasets/...）。
  • tasks/convert_checkpoint/：DCP ↔ HF 双向转换，每模型一个映射文件。
  • .agents/skills/：9 个技能——add-new-model（含模型骨架模板）、validate-correctness（对比分支 loss 曲线）、capture-nsys-profile/analyze-nsys-profile、add-memory-prints、estimate-memory、setup-benchmark-inputs、launch-with-slurm、wandb-tracking。
  • tools/memory_estimator/：不启 GPU 就能估算各并行度下的显存（含流水线模拟器），改配置前先跑它。

2. Step-by-Step 学习路线 ​

设计原则：每阶段有明确产出物，能验证自己学会了；从"能跑"到"能读"到"能改"到"能扩展"。

Phase 0 · 跑通（半天） ​

bash
git clone https://github.com/KuangjuX/pith-train && cd pith-train
uv venv && uv pip install .      # 需 H100/H200/B200, CUDA>=13, Python>=3.12
bash examples/tokenize_corpus/launch.sh dclm-qwen3     # 或先只 tokenize 一个小子集
bash examples/pretrain_lm/launch.sh qwen3-30b-a3b      # 单机 8 卡, PP1/EP8

产出：看到 step 00000001 ... cross-entropy-loss ... tokens-per-second ... 日志。没有 8 卡 H 系 GPU 的话从 deepseek-v2-lite 起（同样配置单机可跑）。

Phase 1 · 读懂一个模型（1–2 天） ​

  1. 读 pithtrain/models/interface.py（116 行，协议全文）。
  2. 精读 models/qwen3_moe.py，对照 HF 的 Qwen3MoeForCausalLM 实现，逐模块对参数名。
  3. 跑/读 tests/ 中 reference vs 分布式的一致性测试，理解 reference_forward 的自校验模式。
  4. 自测：手画 stage1→3→5 的数据流图，标出每个张量 shape 和在哪一侧 EP group。

Phase 2 · 理解并行（2–3 天） ​

  1. 精读 modules/distributed.py（170 行）+ modules/training.py 的 apply_fsdp。
  2. 精读 tasks/pretrain_lm.py 的 get_global_batch/train_step，搞清 DP×EP×CP 的样本划分算术和 token 加权 loss。
  3. 实验：固定单机 8 卡，改 script.py 的 expert_parallel_size（1→2→8）和 pipeline_parallel_size（1→2），观察 tokens/s 与 peak memory 变化；用 tools/memory_estimator 先预测再验证。
  4. 自测：论文 ATE-Bench 的 Q1/Q3/Q4/Q10/Q11/Q12 就是这部分的自测清单（device mesh 怎么建、数据怎么切、seed 怎么设、FSDP 怎么包、梯度裁剪在哪、checkpoint 怎么存）。

Phase 3 · 啃下 DualPipeV（3–5 天，最难也最值） ​

  1. 先读 docs/architecture.md 第 2–3 节（官方导览）。
  2. 按序精读：dualpipev.py（调度骨架）→ overlap.py（重叠主循环）→ execution.py（chunk 生命周期）→ utils.py（WeightGradStore / FP8WeightCache）。
  3. 实验：开 nsys_start=N, nsys_stop=N+1 抓单步 profile，用 analyze-nsys-profile 技能看计算/通信流重叠；对比 PP=1（无流水线）与 PP=2 的 bubble。
  4. 自测：解释为什么 V 形切法能减 bubble；为什么 stage2/4 必须在独立 comm stream；wgrad delay 推迟的是什么、和什么重叠。

Phase 4 · 算子层（按需，2–4 天） ​

  1. 用 benchmarks/operators/ 单测各算子：ep_dispatch、token_scatter、grouped_linear、ring attention。
  2. 精读 ep_dispatch.py 头部注释（把 22 个小算子折叠成 3 个 kernel 的完整思路写在 docstring 里，是全库最好的 Triton 教材）。
  3. 开 fp8=True 跑对照实验，读 deepgemm_quantize.py 理解 128-block scaling。

Phase 5 · 扩展练习（论文的 New Feature 任务，1 周） ​

按 .agents/skills/add-new-model 的模板（model_skeleton.py + inference_test.py）走一遍完整流程：

  1. 入门版：把一个 dense 小模型（如 Qwen3-0.6B dense）接进框架（无 EP，跳过 MoE 机制）。
  2. 进阶版：复现论文的 MoBA / DynMoE / MoE++ 移植任务——论文的正确性判据（64 步 loss 下降且有限 + 三条架构规则）可以直接当验收标准。
  3. 每次改动后用 validate-correctness 技能对比主分支 loss 曲线——这是论文 skill 三属性里 "verifiable success" 的落地。

Phase 6 · 与生产框架对照（可选） ​

用同一模型在 Megatron-LM（NVIDIA MoE 最佳实践）和 PithTrain 各跑 25 步，对比 tokens/s 和 step-time 曲线；再挑 ATE-Bench 的 Q&A 12 题在两个框架上自答，体感"无间接层"的差距。

3. 命令速查 ​

bash
# 数据
bash examples/tokenize_corpus/launch.sh dclm-qwen3
# 训练（单机；多节点走 SLURM: srun -W 0 examples/pretrain_lm/launch.sh <model>）
bash examples/pretrain_lm/launch.sh qwen3-30b-a3b
# checkpoint 导出 HF / 从 HF 导入续训
bash examples/convert_checkpoint/launch.sh qwen3-30b-a3b
# 显存估算（改并行度前）
python -m tools.memory_estimator   # 具体参数见 tools/memory_estimator/__main__.py

训练配置全在 examples/pretrain_lm/<model>/script.py，改完即跑，无 YAML 无 CLI flag（无隐式间接原则的贯彻）。

4. 阅读时建议记的问题清单 ​

  • criterion 为什么返回 sum 而不是 mean？（→ token 加权全局均值）
  • lb_loss 为什么乘 n_tokens 再注入？（→ 绕过 token 缩放后的归一化）
  • MoE 专家为什么沿 dp×cp 分片而非 dp×cp×ep？（→ 专家在 EP 间已互异）
  • zigzag CP 需要哪三处代码协同？（→ 数据切分 / position_ids / ring attention）
  • reshard_after_forward=False 为什么对层成立、对 embed/norm/lm_head 不成立？（→ 层在 5-stage 中被多次进出的复用次数不同）

记录原理、连接知识、积累实践。

Copyright © KuangjuX