Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run
14 hours ago
- 作者根据Sebastian Raschka的书籍笔记，使用JAX、NNX和Optax从零构建了一个LLM训练循环。
- 在构建完整LLM之前，使用“A-to-A”模型（恒等映射）测试了训练循环。
- 主要挑战包括OOM错误、JIT装饰器问题以及学习如何使用JAX进行梯度累积。
- 梯度累积通过Optax的MultiSteps包装器实现，梯度裁剪使用了clip_by_global_norm。
- 使用Safetensors和Orbax实现了检查点保存和从检查点恢复，并在学习率提取代码中使用了调度器。
- 该训练循环在A-to-A任务上实现了110k tokens/秒的吞吐量和0.001的最终损失。
- 学习率调度使用了Optax的warmup_cosine_decay_schedule，并自动绘制了学习率随时间变化的图表。
- 下一篇文章将介绍使用该训练框架构建和训练实际的小型LLM。