FastMoE 系统 易用且高效的基于 PyTorch 的 MoE 模型训练系统 Megatron-LM

FastMoE 系统

简介

FastMoE 是一个易用且高效的基于 PyTorch 的 MoE 模型训练系统.

安装

源代码

依赖

启用了 CUDA 的 PyTorch 是必要的. 当前版本的 FastMoE 在 PyTorch v1.10.0 和 CUDA
11 的平台上经过了测试. 本系统从设计上也支持更旧或更新的 PyTorch 版本.

已知最老的支持的版本是 PyTorch 1.7.0 和 CUDA 10,
但已知某些老版本可能需要修改 FastMoE 的代码以实现支持.

如果需要使能 FastMoE 模型并行特性, 那么支持点对点通信的 NCCL 库 (即不旧于
2.7.5 版本) 也是必需的.

安装

FastMoE 包含一些定制的 PyTorch 算子, 包含一些 C 的组件. 用 python setup.py install
来简单地安装 FastMoE.

FastMoE 分布式模型并行特性默认是被启用的. 如果它需要被禁用,
则需要在运行上述命令时加入环境变量 USE_NCCL=0.

注意, 由于 PyTorch 框架通常仅集成了 NCCL 的运行时组件, 额外的 NCCL
开发包需要被安装在编译环境中, 而且它的版本需要与 PyTorch 的版本相对应. 推荐使用
PyTorch 官方 Docker 镜像,