最新

我已经开源了 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 这主要是一个教育资源，但它在设计时考虑了性能，可以适应更严肃的使用。欢迎贡献，无论是提升性能、添加新模型，还是尝试不同的注意力机制。