基于Megatron-LM的大模型迁移

Megatron-LM是NVIDIA提供的大模型训练加速库，提供了多种针对GPU的训练优化技术。MindSpeed提供了Megatron-LM的适配功能，用户可使用MindSpeed并配合Megatron-LM，在NPU上使用其分布式特性。

当前MindSpeed已适配的特性主要包括：模型并行、流水线并行、序列并行及分布式优化器。使用Megatron-LM训练的大模型，如GPT-3，可参考以下步骤进行模型的迁移与训练。

1. 准备环境，具体请参考《Ascend Extension for PyTorch 配置与安装》。
2. 执行以下命令，下载安装MindSpeed。

   git clone -b 1.1 https://gitee.com/ascend/MindSpeed.git 
   pip install -e MindSpeed

3. 获取Megatron-LM并切换为配套版本。

   # 目前版本基于core_r0.6.0的release版本
   git clone https://github.com/NVIDIA/Megatron-LM.git 
   cd Megatron-LM 
   git checkout core_r0.6.0

4. 以GPT-3为例，在“Megatron-LM”目录下修改“pretrain_gpt.py”文件，在“import torch”下新增一行“import mindspeed.megatron_adaptor”代码，如下黑体加粗部分。

   import os  
   import torch 
   import mindspeed.megatron_adaptor
   from torch import Tensor  
   from functools import partial  
   from typing import Union

5. 在“Megatron-LM”目录下，参考Megatron-LM官方文档准备好训练数据，并在示例脚本中填写对应路径，然后执行以下命令。

   bash examples/pretrain_gpt_distributed.sh

   

   更多模型迁移实例请参考《MindSpeed 迁移开发指南》中的“模型迁移最佳实践”章节。