从零编写LLM，第34a部分——为LLM训练构建JAX训练循环
作者以Sebastian Raschka的《从零构建大语言模型》为教材学习现代AI，在独立完成PyTorch模型训练后，为了验证理解而非机械复现代码，选择用JAX框架编写LLM训练循环。本文记录了从零构建A-to-A（输入等于输出）的最小模型与训练流程的过程，涉及数据加载、设备内存管理、优化器配置等关键技术细节，以及从PyTorch迁移到JAX/Flax NNX/Optax生态的经验教训。
背景速读
- **JAX**: 由 Google 开发的数值计算框架，采用函数式编程风格，自动微分且默认在 GPU/TPU 上运行。与 PyTorch 最大的不同是它不内置神经网络层或优化器，需要搭配生态库使用。
- **Flax NNX / Optax**: JAX 生态中的核心库——Flax NNX 提供神经网络组件（类似 PyTorch 的 nn.Module），Optax 提供优化器（如 AdamW）。本文作者用这两者搭建训练循环。
- **Sebastian Raschka 的《从零构建大语言模型》**: 一本广受关注的 AI 教材，带领读者逐步实现 GPT-2 级别的模型。作者之前一直跟着这本书学习，现在想"脱稿"独立用 JAX 复现一遍，以检验自己的理解。
- **A-to-A 模型**: 作者为了先验证训练循环是否正确，设计了一个极端简化的"语言模型"——只有词嵌入层和输出投影层，没有 Transformer 层，训练目标是让输出直接复制输入（而不是预测下一个 token）。如果能训练到接近零损失，就说明训练循环本身没问题。
- **背景动因**: 作者用了一年多时间跟随 Raschka 的书学习，先后基于 PyTorch 训练出接近 GPT-2 small 质量的基础模型。这次改用 JAX 框架重写全部训练代码，目的是确认自己真正理解了整个流程，而不是死记硬背代码。