训练体系总览：从训练语义到静态图落地
这篇文档面向刚进入训练方向的新同学。它不是某个项目的使用手册，而是一张训练体系地图：训练要经过哪些阶段、每个阶段有哪些关键组件、各自承担什么责任，以及新人应该按什么顺序补齐背景。
训练体系的目标不是写出一个能跑的训练脚本，而是把训练这件事本身建模清楚：数据如何流动、数学公式如何落成算子、数值误差如何传播、训练状态如何演化，再让这些语义能被稳定编译、分布式执行、性能调优、实验复现和故障恢复。
1. 这是什么，不是什么
「跑通」只证明在某组参数、某个规模、某次代码状态下 step 能执行；「体系化」要证明训练语义在工程落地后没有被改写，并且可分片、可优化、可恢复、可复现。
所以训练体系不是一个仓库或一个模块，而是一个跨层系统：训练框架、kernel 库、profiling 工具、精度对齐工具、checkpoint 工具、实验平台和项目流程，只是处在生命周期的不同位置。在团队现有项目里，可以把 MaxText 当训练框架的例子、Pallas Kernel 当底层算子/kernel 体系的例子、Falcon 当实验与分析平台的例子。
2. 训练阶段速览
每个阶段的核心概念、关键组件和责任如下表。它不是线性流程，常常要从后段现象回到前段设计。
| 阶段 | 核心概念 | 关键组件 | 责任（要保证什么） |
|---|---|---|---|
| A 训练目标与数学定义 | 先定义要优化什么，以及归一化口径 | primary loss、auxiliary / router / z-loss、regularization、optimizer 更新规则、gradient clip/scale/accumulation、eval metric | 训练目标与归一化口径显式，能作为后续数值对齐的参照 |
| B 数据流与 batch 语义 | 带语义的 batch 定义了这一步训练的含义；mask/position/segment/loss mask 直接进入 attention/loss 计算，与计算语义耦合 | dataloader、sampler、collator、tokenizer、packing/padding、position/segment id、label、loss mask/weight、microbatch、data sharding | batch 形状稳定、语义明确、可分片、可复现；数据进度可恢复 |
| C 模型表达与算子实现 | 把公式/参考实现落成可验证计算；数学公式 ↔ JAX ↔ 任意参考实现三者等价 | 参数树、模块边界、config 字段、Dense/MoE/MLA/GLA/KDA 等结构、custom kernel/VJP | 语义清楚、公式可追溯、图结构稳定、便于分层验证 |
| D Forward / Backward / Loss | 一个训练 step 不只有 forward，语义必须完整对齐 | forward、loss+归一化、backward/grad、clip/scale、optimizer update、metrics、RNG/step、grad accumulation | 跨框架对齐时归一化口径、RNG 粒度、cast 位置、optimizer 语义保持一致 |
| E 数值行为与正确性校验 | 性能之前先证明算得对；dtype/表达范围/舍入/归约顺序/scale 影响数值稳定性 | 分层 dump/compare、golden data、有限差分、误差阈值、regression | 建立可解释的误差边界（区分预期数值路径 vs 语义错误），分层（op→层→loss→grad→optimizer→step）校验 |
| F 静态图构建 | 把训练语义放进可编译 executable（trace/编译/运行三期见 §3） | PyTree state、jit/AOT、lax.scan/cond/while_loop、static args、编译缓存 | side effect 出图、运行期控制流用图内原语、shape/dtype/sharding 编译前明确、避免重复编译 |
| G 分区并行策略 | 一张全局图映射到多设备 mesh，并行语义进入 HLO | DP、FSDP/ZeRO、TP、PP、CP、EP；逻辑轴 → 物理 mesh、sharding constraint、collective | 模型代码与并行策略解耦；能在 HLO/trace 里读懂 collective 与 reshard |
| H 显存与重计算 | 静态内存规划下的取舍；显存/通信/计算互相牵制 | params/opt/grad/activation/attention 中间量/collective buffer/layout copy/kernel scratch；remat、offload | 按来源定位 HBM，用 sharding/remat/microbatch/精度/layout 在算力-通信-显存间取舍 |
| I Profiling 与性能建模 | 没有测量就没有优化；先建基线再按 gap 优化 | step time、tokens/s、MFU/TFLOPs、HBM、collective time、编译时间、kernel 耗时、input pipeline | 用 roofline/通信/内存模型判断瓶颈类别；profile 可追溯（配置/commit/硬件/trace/结论） |
| J Kernel 与编译器优化 | 在默认 codegen 不够的局部瓶颈上补齐，不是越多越好 | tile、HBM↔VMEM 搬运、double buffer/pipeline、fuse、custom VJP、HLO/LLO dump | 仅当 profiling 证明值得、且默认路径达不到目标时，才沉淀成可复用 kernel |
| K Checkpoint 与保存 | 保存既为恢复，也为对比/debug | 可加载训练状态（params/opt/step/RNG/数据进度/scale/layout/元信息）；观测量制品（loss/各类 norm/数值统计/NaN-Inf/关键 batch/trace/dump/config） | 状态可无损接续；关键指标可复现、可对比、可定位；支持跨框架/layout 转换 |
| L 实验复现与生产化沉淀 | 把一次经验变成团队资产 | experiment record、artifact、analysis report、benchmark 规范、项目流程 | 实验可追溯（author/commit/config/data/hardware/env/XLA flags/checkpoint/metrics/结论）、可复现、可比较 |
几条贯穿全表的要点：
- 正确性先于性能：数学和数值基线没建好，端到端 loss 不一致时几乎无法定位。
- 动态数据 vs 静态图是长期矛盾：变长序列、MoE routing、packing 都要在 padding / packing+segment / bucket / mask / varlen kernel 之间显式取舍，没有免费午餐。
- 自动分区不是自动驾驶：GSPMD 降低复杂度，但逻辑轴、mesh、关键 reshard 边界仍要人定，必要时进
shard_map。 - 数据字段与计算耦合：同一段算子喂入不同 mask 行为就不同，很多看似 attention/loss 的 bug 根源在 batch 构造。
3. XLA 静态图：为什么是核心约束
PyTorch eager 风格逐行执行，调试直接但优化空间有限。JAX/XLA 不同：训练 step 通常被 jit 或 AOT 流程捕获成一张图，再由 XLA 对整张图做优化、分区、内存规划和代码生成。
Python / JAX function -> trace -> jaxpr -> HLO
-> XLA optimization + GSPMD partitioning -> executable -> accelerator runtime核心事实：运行在加速器上的不是 Python 函数本身，而是它在某组静态输入规格下 trace 出来的编译产物。 因此很多工程决策都围绕「可被 trace 和编译」展开。
| 阶段 | 发生什么 | 常见误区 |
|---|---|---|
| Trace 期 | Python 函数被 JAX 追踪，生成 jaxpr | 以为 Python 分支会在每个 step 重新执行 |
| 编译期 | jaxpr lower 成 HLO，XLA 做优化、分区、内存规划、代码生成 | 以为所有性能问题都能靠运行时调度解决 |
| 运行期 | 加速器执行编译后的 executable | 以为运行期可以自由改变 shape 和控制流 |
静态图让系统优化成为可能：全局算子视图（跨 op fusion、layout、copy elimination）、静态内存规划（按 buffer 生命周期安排 HBM、估算峰值）、SPMD 分区（按约束传播分片并自动插 collective）、AOT 编译与缓存（固定输入规格下提前生成、复用）。
静态图也会放大难点：动态 shape（变长序列、MoE routing、padding 影响编译缓存与 kernel 形态）、动态控制流（依赖 tensor 值的分支要改写为 lax.cond/lax.scan/mask）、分布式语义进 HLO（sharding/collective/reshard 不再是运行时黑盒）、内存生命周期（save/remat/offload 改变 HBM 峰值）、custom kernel 边界（Pallas/Mosaic 对 XLA 表现为 custom call，内部不透明）。
新人最重要的第一步不是记 API，而是建立直觉：训练语义必须先定义清楚，才能被可靠地放进静态图和硬件执行计划里。
4. 工具链：目标 × 现状 × 缺口
判断一个环节是否已经「工具化」，标准很简单：新人能否根据文档和标准入口复跑同样的检查，并得到结构化、可比较、可追溯的结果。能，才算工具链能力；不能，就还只是个人经验。下面按能力簇说明每一类的目标、现状和缺口。状态按当前 wiki 已文档化的能力判断——即便代码仓里可能有实现，只要没有稳定文档和标准入口，这里仍按「部分/缺口」处理。
4.1 训练语义与数据契约
- 目标：模型接入前冻结训练语义契约（objective/loss/归一化/optimizer/RNG/grad accumulation/metrics/checkpoint state）与数据契约（raw→batch 每步显式：dataloader/sampler/collator/tokenizer/packing/position/mask/label/loss weight/microbatch/sharding），并和 config 绑定。
- 现状：MaxText 配置体系和设计文档可承载部分语义；现有文档讨论了数据和 mask 对训练的影响，但没有独立的数据流工具。
- 缺口：训练语义契约 schema 与 validator（检查 loss 归一化、RNG 粒度、optimizer 语义、并行语义、checkpoint state 是否完整）；data contract manifest、batch inspector（展示某 batch 的 token/segment/position/label/loss mask 以及它们如何耦合 attention/loss）；数据/tokenizer 版本与采样进度的统一记录 + 恢复校验。
4.2 数学与数值正确性
- 目标：每个模块都能把数学公式 ↔ JAX 实现 ↔ 参考实现对齐，分层验证 forward/backward/optimizer，并给出可解释的误差边界（按 dtype/tensor 类别设阈值）。正确性和性能都依赖这层表达是否清晰。
- 现状：精度对齐 workflow、Argus dump、Falcon
precision-compare（支持阈值）、accuracy-alignment thresholds、Ling Megatron dump 指南、multi-artifact analysis 已覆盖跨框架比较主链路；Pallas 模板要求写公式/API/测试/benchmark。 - 缺口：formula↔JAX↔参考实现映射模板与 formula-to-op checklist；custom VJP 验证 harness、backward Jacobian/有限差分抽检；dtype policy manifest、数值范围扫描器（min/max、Inf/NaN、underflow、scale 利用率）、reduction-order 差异解释；统一 golden data registry；从 failed tensor 自动回溯到模块/公式/配置。
4.3 静态图与并行
- 目标：训练 step 可稳定
jit/AOT compile（输入 shape、static args、控制流、PyTree state、编译缓存可控）；先定逻辑轴和 mesh，再估 memory/communication/overlap，用 HLO/trace 检查 collective 与 reshard。 - 现状：MaxText JAX/XLA 原语分析、训练主循环与 AOT 编译文档、HLO dump 指南、sharding 文档、performance System Bound、Falcon profiling artifact。
- 缺口：compile contract checker（PR 阶段查动态 shape、非预期 Python 控制流、host 回读、重复编译风险）、jaxpr/HLO diff；sharding plan linter、从 HLO 自动提取 collective/reshard/通信量并和预期对比、配置 sweep 推荐器。
4.4 显存、性能与 kernel
- 目标：用 memory model 预估可行配置、HLO buffer assignment / XProf 定位峰值 buffer；建 roofline / system bound，跑 benchmark/profiling 后按 gap 排序优化；只在证明值得的局部瓶颈上写 Pallas kernel（补公式、shape/dtype 限制、forward/backward、benchmark、roofline、已知限制）。
- 现状：JAX TPU 显存分析指南、XLA dump、XProf、performance analysis framework、Falcon、Pallas Kernel 经验 + 模板 + benchmark 规范、
pallas-llo-analysis。 - 缺口：自动 HBM attribution（峰值 buffer → Python module/配置/可选优化动作）、remat policy 自动 sweep；把 roofline/实测/历史实验串成优化看板、跨实验自动 diff 与归因；统一 kernel registry、kernel correctness + performance CI、从 HLO 自动识别应替换为 kernel 的候选热点。
4.5 Checkpoint 与保存
- 目标：保存完整训练状态（schema 校验、round-trip、跨框架转换、恢复后 loss/grad/step 平滑性），同时保存可对比、可 debug 的观测量制品。
- 现状：Full-State Safetensor → MaxText Checkpoint Converter 设计已覆盖 Ling3 场景。
- 缺口：能力偏特定模型/路径；通用 checkpoint schema validator、optimizer state coverage checker、恢复训练 smoke test、数据进度/RNG 连续性验证；观测量制品的标准化保存（与下一簇实验平台对接）。
4.6 实验复现与生产化
- 目标：每次实验记录 author/commit/config/data/hardware/env/XLA flags/checkpoint/artifact/metrics/结论，并支持 analyzer 自动运行。
- 现状：Falcon experiment/artifact/analyzer/orchestrator 已覆盖主要基础设施。
- 缺口：训练语义契约、数据契约、数值策略与实验记录的强绑定；从 PR → experiment → analysis report 的端到端模板化工作流。
完成一个具体目标时，先按上面六个能力簇定位它跨哪几类，优先用已文档化的工具，凡是只能靠临时脚本、人肉 diff、口头约定的环节，就记为缺口；把 mapping、threshold、golden case、benchmark config、analysis report、known deviation 沉淀回仓库或实验平台。
5. 新人学习路径
按「训练语义 → 静态图表达 → 分布式执行 → 验证与优化」的主线建立知识结构，不要一开始就深挖所有细节。
第一步：全局地图。 读 团队目标 (docs/onboarding/goals.md) 和本文。读完应能回答：训练体系首先要建模哪些语义？生命周期有哪些阶段？框架/kernel/profiling/checkpoint/实验平台各解决什么问题？
第二步：训练语义、数据流与正确性。 读 MaxText (docs/projects/maxtext/index.md)、checkpoint/converter 设计文档如 Full-State Safetensor to MaxText Checkpoint Converter (docs/designs/2026-05-25-full-state-checkpoint-converter.md)。读完应能回答：raw data 如何经 dataloader/sampler/collator 变成 batch？loss/attention mask、position id 影响什么语义、如何与 attention/loss 耦合？forward/backward/optimizer 如何分层校验？dtype 与表达范围为什么影响稳定性？
第三步：JAX/XLA 训练主循环。 读 JAX/XLA 原语分析 (docs/projects/maxtext/jax-xla/index.md)、训练主循环与 AOT 编译 (docs/projects/maxtext/jax-xla/01-train-loop-aot.md)、分区：逻辑轴 → 物理 mesh (docs/projects/maxtext/jax-xla/02-sharding.md)。读完应能回答：训练 step 如何被 trace 和 compile？Python 控制流和 JAX 控制流的区别？logical axis 如何映射到 mesh？为什么某些 shape 变化触发重新编译？
第四步：显存、重计算与复杂模型结构。 读 Decoder 堆叠与重计算 / Offload (docs/projects/maxtext/jax-xla/03-decoder-remat.md)、注意力 (docs/projects/maxtext/jax-xla/04-attention.md)、MoE / 专家并行 (docs/projects/maxtext/jax-xla/05-moe.md)、上下文并行 / KDA 线性注意力 (docs/projects/maxtext/jax-xla/07-context-parallel-kda.md)。读完应能回答：remat 为什么降低 HBM？MoE 为什么同时是数值、通信、kernel 问题？attention 变体为什么常需自定义 kernel？
第五步：性能分析与 kernel 优化。 读 性能优化 (docs/projects/performance-optimization/index.md)、性能理论分析框架 (docs/projects/performance-optimization/analysis-framework.md)、Pallas Kernel 编写经验总结 (docs/onboarding/pallas-kernel-guide.md)、JAX TPU 显存分析指南 (docs/onboarding/profiling-hbm-guide.md)。读完应能回答：MFU/HBM/collective time/step time 各说明什么？roofline 如何指导优化？HLO buffer assignment 能看出哪些显存问题？什么情况下值得写 Pallas kernel？
第六步：实验复现与生产化流程。 读 Falcon: 性能分析平台 (docs/projects/falcon/index.md)、项目管理框架 (docs/onboarding/project-management.md)、Beaver 项目流转使用指南 (docs/onboarding/beaver-workflow.md)。读完应能回答：一次实验该记录哪些上下文？checkpoint 为什么要含 optimizer state？如何判断恢复训练是否平滑？为什么项目流程和实验平台也属于训练体系？
6. 排查心法与总结
遇到训练问题，不要先问「改哪行代码」，先把现象放回生命周期阶段：
| 现象 | 优先检查方向 |
|---|---|
| 频繁重新编译 | 输入 shape、static arg、控制流、配置是否变化 |
| loss 对不上 | 训练目标、数据、mask、loss 归一化、dtype、optimizer 语义 |
| 梯度异常 | backward 正确性、scale、clip、归约顺序、低精度表达范围 |
| OOM | HLO buffer、activation、remat、sharding、layout copy |
| step time 慢 | profile、roofline、collective、HBM、kernel、input pipeline |
| 大规模才失败 | sharding、collective、checkpoint layout、host/device 环境差异 |
| 恢复后 spike | optimizer state、step counter、LR schedule、数据进度、RNG |
背后的原则一致：训练体系的问题通常不是孤立 bug，而是某个阶段的契约没被满足。建设训练体系的核心，不是让代码更复杂，而是把复杂性放到正确的位置——上层框架稳定表达语义和并行，正确性/数值工具证明没有偏离预期，XLA 静态图提供全局优化与内存规划的基础，kernel 工具补齐局部瓶颈，checkpoint 和实验工具保证结果可信、可恢复、可复现，文档和流程把个人经验变成团队资产。