maxtext

收录理由

如果你打算用 JAX 训练或微调大语言模型，又不想从零搭建整套训练流程，MaxText 会是个很合适的起点。它内置了 Llama、Gemma、DeepSeek、Qwen、Mistral 等多个开源模型的可用参考实现，覆盖预训练和 SFT、强化学习等后训练环节，并能扩展到数万芯片规模的集群。借助 JAX 和 XLA 编译器，它不需要手工调优的内核就能获得较高的硬件利用率，同时代码保持纯 Python 编写，易读也好改。项目主要面向 Google Cloud TPU 设计，也能跑在 NVIDIA GPU 上，对于已经使用 Google Cloud 基础设施的团队来说上手很顺。由于它本身就是一份参考实现，想了解生产级 LLM 训练如何组织的人，也可以拿它当作学习材料，再据此构建自己的方案。