快速入门

本手册指导用户完成安装后，可实现Megatron-LM在昇腾设备上的高效运行，且无缝集成并充分发挥MindSpeed所提供的丰富加速与优化技术。

环境准备

1. 请先参考MindSpeed安装指导进行环境准备。

2. 在Megatron-LM中导入MindSpeed适配器。

   在“Megatron-LM”目录下修改pretrain_gpt.py文件，在“import torch”下新增一行“import mindspeed.megatron_adaptor”代码，即如下修改：

   import torch
   import mindspeed.megatron_adaptor # 新增代码行
   from functools import partial
   from contextlib import nullcontext
   import inspect
   

数据准备

参考Megatron-LM官方文档准备训练数据。

1. 下载Tokenizer

   新建“Megatron-LM/gpt-tokenizer”目录，并将vocab.json和merges.txt文件下载至该目录。

2. 下载数据集，以Alpaca数据集为例

   Note

   用户需要自行设置代理，以便访问或下载数据集。

3. 语料格式转换

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

   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')
   

4. 预训练数据集生成

   若在昇腾设备上使用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
   

   表 1 preprocess_data.py参数说明

   参数	说明
   --input	数据集
   --output-prefix	处理后的数据集
   --tokenizer-type	tokenizer类型
   --vocab-file	tokenizer文件
   --merge-file	tokenizer文件
   --append-eod	添加<-eod>结尾token
   --log-interval	log迭代数
   --workers	并行数

   执行成功后，将在gpt_pretrain_data目录下生成两个文件：alpaca_text_document.bin和alpaca_text_document.idx，代表预处理完成的预训练数据集。

启动预训练

1. 配置环境变量

   当前以root用户安装后的默认路径为例，请用户根据set_env.sh的实际路径执行如下命令。

   source /usr/local/Ascend/cann/set_env.sh
   

2. 准备预训练脚本

   在“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="
       --transformer-impl local \
       --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
   
   

3. 配置路径

   请编辑示例脚本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
   

   表 2 train_distributed.sh参数配置

   参数	说明
   CKPT_DIR	权重文件路径
   VOCAB_FILE	tokenizer文件
   MERGE_FILE	tokenizer文件
   DATA_PATH	数据集文件

   上述路径需根据您的实际情况进行适当调整。

4. 运行脚本启动预训练

   bash ./train_distributed.sh
   

   Note

   示例脚本train_distributed.sh中的部分参数（如--hidden-size、--num-layers）需根据实际场景进行适配，避免OOM等现象。