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版本之间的匹配关系,确保用户能够根据自身软件环境准确选择相匹配的版本,以实现最优的性能与功能支持。
MindSpeed版本 |
1.0.0 |
---|---|
MindSpeed代码分支名称 |
1.0.0_core_r0.6.0 说明:
配套Megatron-LM的core_r0.6.0分支。 |
1.0.0_core_r0.7.0 说明:
配套Megatron-LM的core_r0.7.0分支。 |
|
CANN版本 |
8.0.0 |
PyTorch版本 |
2.1.0 |
Ascend Extension for PyTorch版本 |
6.0.0 |
Ascend Extension for PyTorch安装包版本 |
2.1.0.post10 |
Python版本 |
Python3.8.x,Python3.9.x,Python3.10.x |
依赖软件 |
软件版本 |
---|---|
NPU驱动 |
建议您下载并安装左侧软件栈的最新版本,具体请参见《CANN 软件安装指南》。 |
NPU固件 |
|
Toolkit(开发套件) |
|
Kernels(算子包) |
|
NNAL(Ascend Transformer Boost加速库) |
|
PyTorch框架 |
建议您下载并安装左侧软件栈的最新版本,具体请参见《Ascend Extension for PyTorch 配置与安装》。 |
torch_npu插件 |
|
apex |
git clone -b 1.0.0_core_r0.6.0 https://gitee.com/ascend/MindSpeed.git pip3 install -e mindspeed
git clone -b 1.0.0_core_r0.7.0 https://gitee.com/ascend/MindSpeed.git pip3 install -e mindspeed
git clone https://github.com/NVIDIA/Megatron-LM.git cd Megatron-LM git checkout core_r0.6.0
git clone https://github.com/NVIDIA/Megatron-LM.git cd Megatron-LM git checkout core_r0.7.0
import os import torch import mindspeed.megatron_adaptor from functools import partial from typing import Union
else: assert(args.context_parallel_size == 1), "Context parallelism is only supported with Megatron Core!" model = megatron.legacy.model.GPTModel( config, num_tokentypes=0, parallel_output=True, pre_process=pre_process, post_process=post_process )
新建“Megatron-LM/gpt-tokenizer”目录,并将vocab.json和merges.txt文件下载至该目录。
以Alpaca数据集为例,可单击Link获取。
用户需要自行设置代理,以便访问或下载数据集。
source /usr/local/Ascend/ascend-toolkit/set_env.sh
pip3 install nltk pyarrow pandas
1 2 3 4 5 6 7 8 | 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') |
import torch import mindspeed.megatron_adaptor import numpy as np
通过运行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,代表预处理完成的预训练数据集。
请编辑示例脚本“examples/pretrain_gpt_distributed.sh”,并设置如下环境变量以指定必要的路径:
1 2 3 4 | CHECKPOINT_PATH=./ckpt VOCAB_FILE=./gpt-tokenizer/vocab.json MERGE_FILE=./gpt-tokenizer/merges.txt DATA_PATH=./gpt_pretrain_data/alpaca_text_document |
注意,上述路径需根据您的实际情况进行适当调整。
bash examples/pretrain_gpt_distributed.sh