:robot: MegaBlocks

MegaBlocks 是一个轻量级的混合专家模型（MoE）训练库。其核心是高效的 "无丢弃-MoE"（dMoE，论文）和标准 MoE 层。

MegaBlocks 已集成到 Megatron-LM 中，支持 MoE 的数据并行、专家并行和流水线并行训练。敬请期待与 Databricks 库和工具的进一步集成！

:rocket: 性能

相比使用 Tutel 训练的最佳 capacity_factor 配置，MegaBlocks dMoE 的性能提升高达 40%。MegaBlocks dMoE 采用基于块稀疏操作的 MoE 重构，无需牺牲硬件效率即可避免令牌丢弃。除了速度更快，MegaBlocks 通过完全移除 capacity_factor 超参数简化了 MoE 训练。与使用 Megatron-LM 训练的密集 Transformer 相比，MegaBlocks dMoE 可将训练速度提升至多 2.4 倍。更多细节请查阅我们的论文！

:building_construction: 安装

注意：假设您已安装 numpy 和 torch。

使用 Megatron-LM 训练模型： 推荐使用 NGC 的 nvcr.io/nvidia/pytorch:23.09-py3 PyTorch 容器。Dockerfile 在此镜像基础上添加了额外依赖。构建镜像请运行 docker build . -t megablocks-dev，然后运行 bash docker.sh 启动容器。进入容器后，通过 pip install . 安装 MegaBlocks。关于使用 MegaBlocks + Megatron-LM 训练 MoE 的说明，请参阅用法。

在其他包中使用 MegaBlocks： 要在其他框架中使用 MegaBlocks 包，请运行 pip install megablocks。例如，通过此安装方法，可使用 vLLM + MegaBlocks 运行 Mixtral-8x7B。

可选依赖： MegaBlocks 具有可启用附加功能的可选依赖。

安装 megablocks[gg] 可使用分组 GEMM 进行 dMoE 计算。此功能通过将 mlp_impl 参数设置为 grouped 启用。目前这是针对 Hopper 代 GPU 的推荐路径。

安装 megablocks[dev] 可参与 MegaBlocks 开发并在本地测试。安装 megablocks[testing] 可通过 Github Actions 进行测试。如果您已安装 megablocks[dev]，可以运行 pre-commit install 配置预提交钩子以自动格式化代码。

通过 megablocks[all] 包可安装所有依赖（testing 除外）。

:steam_locomotive: 用法

我们在顶级目录中提供了预训练 Transformer MoE 和 dMoE 语言模型的脚本。最快的入门方式是使用 实验启动脚本。这些脚本需要 Megatron-LM 格式的数据集，您可按照其说明创建。

:writing_hand: 引用

@article{megablocks,
  title={{MegaBlocks: Efficient Sparse Training with Mixture-of-Experts}},
  author={Trevor Gale and Deepak Narayanan and Cliff Young and Matei Zaharia},
  journal={Proceedings of Machine Learning and Systems},
  volume={5},
  year={2023}
}