3作者: Dhyaneesh3 个月前原帖
我已经开源了 awesome-jax-flax-llms,这是一个使用 JAX 和 Flax 从零开始构建的大型语言模型(LLM)实现的精心整理的集合。该仓库旨在支持在 TPU/GPU 上进行高性能训练,非常适合研究人员、机器学习工程师和希望探索或扩展现代变换器模型的好奇者。 主要特点: - 模块化、可读性强且可扩展的代码库 - 纯 JAX/Flax 实现的 GPT-2 和 LLaMA 3 - 使用 XLA + Optax 加速训练 - 支持 Google Colab(TPU 兼容) - 集成 Hugging Face 数据集 - 即将支持微调、Mistral 和 DeepSeek-R 这主要是一个教育资源,但它在设计时考虑了性能,可以适应更严肃的使用。欢迎贡献,无论是提升性能、添加新模型,还是尝试不同的注意力机制。