MaxText

MaxText 是团队面向 TPU 的 JAX 大模型训练引擎（fork 自 Google MaxText，仓库 primatrix/maxtext），隶属 Goal 1：计算引擎与工具链，承担 Ling 系列模型的预训练。其定位围绕三条主线：

• 快速模型架构适配 — 以标准化流程接入新模型架构（Dense / MoE / 新注意力变体），缩短从设计到引擎落地的周期。
• 与 Megatron-LM 数值对齐 — 相同权重 + 相同输入下产出与 Megatron-LM 数值一致的 forward / backward / loss。
• Ling 系列训练极致性能 — 以 roofline 分析为指导，在 TPU 上持续逼近硬件理论上限（MFU 最大化）。

文档方向

本项目文档围绕以下五个方向展开：

框架整体架构

MaxText 的代码结构、对 Jax/XLA 的使用与原理、切分的方式与逻辑、训练主循环与配置驱动的能力开关，帮助理解引擎如何从 config 一路落到 TPU 上的训练 step。

新人入门

环境准备、跑通第一个训练任务、代码导览与开发协作流程，帮助新成员快速上手并参与贡献。

精度对齐

与 Megatron-LM 的数值一致性验证方法论：两阶段对齐（CPU 逻辑对齐 → TPU 精度对齐）、分层测试与「捕获 → 比对 → 定位 → 收敛」流程。详见 精度对齐。

性能优化

以 roofline 两级上界分析为指导的 MFU 最大化：并行 / 重计算 sweep、XLA flag 调优、IR dump 分析与高性能算子集成。详见 性能优化 与 pallas-kernel。

工具链建设

包括但不限于精度对比工具, HLO 分析工具, 自动化实验体系, 数学等价性保证工具等