jax-training-recipe

JAX 训练方案

JAX 并不是换了语法的 PyTorch。函数是纯函数，状态是显式的，而且只有在函数边界清晰时，变换（jit、vmap、grad、pmap）才能组合使用。

如需查看具体配置和检查清单，请阅读 references/training-recipe-details.md。

信息收集

请说明：

• 任务、数据集、损失函数和指标；
• 目标设备：CPU、单 GPU、多 GPU、TPU；
• 模型框架：Flax/Linen、Equinox 或原始 JAX；
• 随机性需求：dropout、采样、数据增强；
• 状态需求：batch norm、优化器状态、EMA、检查点。

工作流程

1. 定义实验和最小数据批次。
2. 规划 PRNG key 的归属与拆分方式。
3. 定义模型和状态的边界。
4. 初始化模型和优化器。
5. 编写纯函数 train_step。
6. 仅在未启用 jit 的步骤运行正常后，再添加 jit。
7. 仅在单设备行为正确后，再添加 vmap/pmap。
8. 添加检查点和恢复测试。
9. 执行首次运行安全检查。
10. 按从小批次到完整训练的顺序分阶段扩展。

JAX 专属规则

• 绝不重复使用 PRNG key。
• 保持状态显式；不要在闭包中隐藏可变状态。
• 避免在 jit 中执行 Python 副作用。
• 将 jit 下的形状变化视为设计决策。
• 先在不启用 jit 的情况下调试，然后逐一重新启用变换。
• 在多设备训练中，有意复制状态。

常见故障模式

• 重复使用 PRNG key，导致 dropout/采样结果重复。
• Python 控制流或副作用导致 tracer 错误。
• batch norm/dropout 状态未与参数分离。
• vmap 或 pmap 下出现形状不匹配。
• 编译循环掩盖 NaN。
• 检查点缺少优化器状态或 batch 统计状态。

输出

请返回：

• 实验摘要；
• PRNG 策略；
• 模型/状态设计；
• 训练/评估步骤计划；
• 最小配置；
• 首次运行检查；
• JAX 专属故障模式；
• 扩展计划。