Megatron-LM是NVIDIA提供的大模型训练加速库,提供了多种针对GPU的训练优化技术。MindSpeed提供了Megatron-LM的适配功能,用户可使用MindSpeed并配合Megatron-LM,在NPU上使用其分布式特性。
当前MindSpeed已适配的特性主要包括:模型并行、流水线并行、序列并行及分布式优化器。使用Megatron-LM训练的大模型,如GPT-3,可参考以下步骤进行模型的迁移与训练。
git clone -b 1.0 https://gitee.com/ascend/MindSpeed.git pip install -e MindSpeed
git clone https://github.com/NVIDIA/Megatron-LM.git cd Megatron-LM git checkout bcce6f54e075e3c3374ea67adefe54f3f2da2b07
import os import torch import mindspeed.megatron_adaptor from torch import Tensor from functools import partial from typing import Union
bash examples/pretrain_gpt_distributed.sh