在 AWS Trainium 上训练 MoE
简介
本文将介绍如何在 AWS Trainium 上训练混合专家 (MoE) 模型。
我们将使用的代码基于用于开发KARAKURI LM 8x7B Chat v0.1 的代码，该模型是世界上第一个在 Trainium 上训练的 MoE 模型。
什么是 Trainium？
AWS Trainium 是 AWS 专门为训练机器学习模型而开发的加速器。据说 Trainium 可以将训练成本降低高达 50%，与同等 GPU 实例相比。
有关 Trainium 的更多详细信息，请参阅 AWS 的这篇博文。
什么是 MoE？
在 Transformer 模型中，MoE 是一种将前馈层划分为多个独立单元（专家）的技术，并在每个前向传播期间仅激活这些专家的一部分。这种方法有助于减少通常伴随参数数量增加而产生的计算和内存占用量的增加。
有关 MoE 的更多详细信息，请参阅 Hugging Face 的这篇博文。
分布式训练库
在 AWS Trainium 上，常用的分布式训练库是 AWS Neuron Reference for NeMo Megatron (neuronx-nemo-megatron)。
本文中使用的库是 neuronx-nemo-megatron 的修改版本。具体来说，它包含用于 Mixtral 模型中使用的稀疏 MoE 层的实现。此外，它还已适配支持 Hugging Face Datasets，从而更易于使用 Hugging Face Hub 中的数据集。
1. 设置环境
首先，我们将设置 AWS 基础设施。请参考以下链接配置 VPC 和 ParallelCluster
2. 安装所需工具
接下来，我们将按照 官方 AWS 教程 安装所需工具。
连接到头节点
首先，通过 SSH 连接到头节点
ssh -i YOUR_KEY.pem ubuntu@HEAD_NODE_IP_ADDRESS
激活虚拟环境
激活虚拟环境
cd ~
source ./aws_neuron_venv_pytorch/bin/activate
克隆存储库
克隆仓库
cd ~
git clone https://github.com/karakuri-ai/neuronx-nemo-megatron.git
cd neuronx-nemo-megatron
请注意，虽然 AWS 教程使用 aws-neuron/neuronx-nemo-megatron，但我们将使用 karakuri-ai/neuronx-nemo-megatron。
构建包并安装依赖项
构建包并安装依赖项
pip install wheel
./build.sh
pip install ./build/*.whl
pip install -r requirements.txt protobuf==3.20.3
cd ~
python -c "from nemo.collections.nlp.data.language_modeling.megatron.dataset_utils import compile_helper; compile_helper()"
3. 准备数据集
下载并预处理数据集
cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling
python create_mixtral_sft_dataset.py
这里我们使用 No Robots 数据集。如果您希望使用不同的数据集，请修改 create_mixtral_sft_dataset.py 脚本。
4. 训练模型
转换检查点 (HF -> NeMo)
将检查点从 Hugging Face 格式转换为 NeMo 格式
cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling/checkpoint_conversion
python convert_hf_checkpoint_to_nemo_mixtral.py \
--path_to_checkpoint /path/to/hf_checkpoint \
--config_file /path/to/hf_checkpoint/config.json \
--model_bin_file /path/to/hf_checkpoint/pytorch_model.bin.index.json \
--output_path /path/to/nemo_checkpoint \
--tp_degree 8 \
--pp_degree 8 \
--save_bf16 True \
--num_shards 19
如果在 ParallelCluster 之外运行，请将转换后的检查点上传到链接到 FSx 的 S3 存储桶。
编辑路径
编辑 mixtral_8x7b.sh 以指定数据集和检查点的路径。
预编译模型
要预编译模型，请运行以下命令
sbatch --nodes 2 compile.slurm ./mixtral_8x7b.sh
开始训练
开始训练过程
sbatch --nodes 2 run.slurm ./mixtral_8x7b.sh
转换检查点 (NeMo -> HF)
训练完成后，将检查点转换回 Hugging Face 格式
cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling/checkpoint_conversion
python convert_nemo_checkpoint_to_hf_mixtral.py \
--path_to_checkpoints /path/to/nemo_checkpoint \
--config_file /path/to/hf_config_file \
--output_path /path/to/hf_checkpoint \
--is_xser True \
--dtype bfloat16
5. 推理
推理可以使用 GPU 或 AWS Inferentia2 进行。有关基于 Inferentia2 的推理的实现细节，请参阅 AWS 提供的示例代码。
结论
KARAKURI LM 8x7B Chat v0.1 是首个在 Trainium 上训练的 MoE 模型，现已在 Hugging Face Hub 上可用。发布时，该模型在 MT-Bench-jp（一个用于评估日语多轮对话能力的基准）上的性能位居开放模型前列。
此外，我们还发布了一个演示，展示了使用 Inferentia2 进行推理。