基于Megatron-LM的大模型迁移

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

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

前提条件

为了便于用户理解和选择合适的MindSpeed版本，我们提供了详细的版本配套表，如表1所示。

该表详细列出了MindSpeed版本与对应的PyTorch、Ascend Extension for PyTorch版本以及CANN版本之间的匹配关系，确保用户能够根据自身软件环境准确选择相匹配的版本，以实现最优的性能与功能支持。

表1 版本配套表
MindSpeed版本	2.0.0
MindSpeed代码分支名称	2.0.0_core_r0.8.0：配套Megatron-LM的core_r0.8.0分支。
CANN版本	8.1.RC1
PyTorch版本	2.1.0
Ascend Extension for PyTorch版本	7.0.0
Ascend Extension for PyTorch安装包版本	2.1.0.post12
Python版本	Python3.8.x，Python3.9.x，Python3.10.x

用户请根据需要选择MindSpeed代码分支下载源码并进行安装。
更多关于PyTorch和CANN的版本信息可点击昇腾辅助软件查看。

安装操作

安装依赖的软件。

表2 软件列表
依赖软件	软件版本
NPU驱动	建议您下载并安装左侧软件栈的最新版本，具体请参见《CANN 软件安装指南》。
NPU固件
Toolkit（开发套件）
Kernels（算子包）
NNAL（Ascend Transformer Boost加速库）
PyTorch框架	建议您下载并安装左侧软件栈的最新版本，具体请参见《Ascend Extension for PyTorch 软件安装指南》。
torch_npu插件
apex

执行以下命令，下载并安装MindSpeed源码2.0.0_core_r0.8.0分支。

git clone -b 2.0.0_core_r0.8.0 https://gitee.com/ascend/MindSpeed.git
pip3 install -e MindSpeed

获取Megatron-LM并指定core_r0.8.0分支。

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_r0.8.0

在“Megatron-LM”目录下修改pretrain_gpt.py文件，在“import torch”下新增一行“import mindspeed.megatron_adaptor”代码，如下黑体加粗部分。
```
import os
import torch
import mindspeed.megatron_adaptor
from functools import partial
from typing import Union
```
数据准备，参考Megatron-LM官方文档准备训练数据。
1. 下载Tokenizer。
  新建“Megatron-LM/gpt-tokenizer”目录，并将vocab.json和merges.txt文件下载至该目录。
2. 下载数据集。
  以Alpaca数据集为例，可单击Link获取。
用户需要自行设置代理，以便访问或下载数据集。
配置环境变量，当前以root用户安装后的默认路径为例，请用户根据set_env.sh的实际路径执行如下命令。
```
source /usr/local/Ascend/ascend-toolkit/set_env.sh
```

数据处理，详情可单击对应分支core_r0.8.0进行参考。

语料格式转换。

数据处理依赖于多个第三方库，请确保已正确安装如下依赖：

pip3 install nltk pyarrow pandas

以下代码段展示了如何读取Parquet格式的原始语料，并将其转换为JSON格式，以便后续处理。

import json
import pandas as pd
data_df = pd.read_parquet("train-00000-of-00001-a09b74b3ef9c3b56.parquet")
data_df['text'] = data_df['text'].apply(lambda v: json.dumps({"text": v}))
with open("alpaca_json.json", encoding='utf-8', mode='w') as f:
    for i, row in data_df.iterrows():
        f.write(row['text'])
        f.write('\n')

预训练数据集生成。
若在昇腾设备上使用preprocess_data.py脚本处理数据，须在“Megatron-LM”目录下修改“tools/preprocess_data.py”脚本，在“import torch”下新增一行“import mindspeed.megatron_adaptor”代码，如下黑体加粗部分。
```
import torch
import mindspeed.megatron_adaptor
import numpy as np
```
新建“Megatron-LM/gpt_pretrain_data”目录，通过运行preprocess_data.py脚本，可以将转换后的JSON文件进一步处理为适合Megatron-LM预训练的二进制格式。
```
python tools/preprocess_data.py \
   --input alpaca_json.json \
   --output-prefix ./gpt_pretrain_data/alpaca \
   --tokenizer-type GPT2BPETokenizer \
   --vocab-file ./gpt-tokenizer/vocab.json \
   --merge-file ./gpt-tokenizer/merges.txt \
   --append-eod \
   --log-interval 1000 \
   --workers 8
```
执行成功后，将在gpt_pretrain_data目录下生成两个文件：alpaca_text_document.bin和alpaca_text_document.idx，代表预处理完成的预训练数据集。

在“Megatron-LM”目录下准备预训练脚本train_distributed.sh，脚本示例如下：

#!/bin/bash
export CUDA_DEVICE_MAX_CONNECTIONS=1
NPUS_PER_NODE=8
MASTER_ADDR=localhost
MASTER_PORT=6001
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
CKPT_DIR=./ckpt
VOCAB_FILE=<Specify path to file>/vocab.json
MERGE_FILE=<Specify path to file>/merges.txt
DATA_PATH=<Specify path and file prefix>_text_document
TP=2
PP=2
CP=1
EP=1
DISTRIBUTED_ARGS="
    --nproc_per_node $NPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"
GPT_ARGS="
    --tensor-model-parallel-size ${TP} \
    --pipeline-model-parallel-size ${PP} \
    --num-layers-per-virtual-pipeline-stage 1 \
    --num-layers 8 \
    --hidden-size 4096 \
    --ffn-hidden-size 14336 \
    --num-attention-heads 64 \
    --seq-length 4096 \
    --max-position-embeddings 4096 \
    --micro-batch-size 1 \
    --global-batch-size 16 \
    --make-vocab-size-divisible-by 1 \
    --lr 1.0e-6 \
    --train-iters 1000 \
    --init-method-std 0.01 \
    --no-masked-softmax-fusion \
    --attention-softmax-in-fp32 \
    --min-lr 1.0e-7 \
    --weight-decay 0.1 \
    --clip-grad 1.0 \
    --initial-loss-scale 4096.0 \
    --disable-bias-linear \
    --lr-warmup-fraction 0.01 \
    --fp16
"
DATA_ARGS="
    --split 990,5,5
    --data-path $DATA_PATH \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
"
OUTPUT_ARGS="
    --log-throughput \
    --log-interval 1 \
    --save-interval 10000 \
    --eval-interval 10000 \
    --eval-iters 10 \
"
torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
set +x

配置路径。

请编辑示例脚本train_distributed.sh，并设置如下环境变量以指定必要的路径：

CKPT_DIR=./ckpt
VOCAB_FILE=./gpt-tokenizer/vocab.json
MERGE_FILE=./gpt-tokenizer/merges.txt
DATA_PATH=./gpt_pretrain_data/alpaca_text_document

执行如下命令启动预训练。
```
bash ./train_distributed.sh
```
- 示例脚本train_distributed.sh中的部分参数（如--hidden-size、--num-layers）需根据实际场景进行适配，避免OOM等现象。
- 更多模型开发实例请参见《MindSpeed迁移开发指南》中的“模型开发实践”章节。

父主题： LLM大模型迁移案例