基于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版本

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

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

安装操作

  1. 安装依赖的软件。

    表2 软件列表

    依赖软件

    软件版本

    NPU驱动

    建议您下载并安装左侧软件栈的最新版本,具体请参见CANN 软件安装指南

    NPU固件

    Toolkit(开发套件)

    Kernels(算子包)

    NNAL(Ascend Transformer Boost加速库)

    PyTorch框架

    建议您下载并安装左侧软件栈的最新版本,具体请参见Ascend Extension for PyTorch 配置与安装

    torch_npu插件

    apex

  2. 执行以下命令,下载并安装MindSpeed。

    • 1.0.0_core_r0.6.0分支
      git clone -b 1.0.0_core_r0.6.0 https://gitee.com/ascend/MindSpeed.git
      pip3 install -e mindspeed
    • 1.0.0_core_r0.7.0分支
      git clone -b 1.0.0_core_r0.7.0 https://gitee.com/ascend/MindSpeed.git
      pip3 install -e mindspeed

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

    • core_r0.6.0分支
      git clone https://github.com/NVIDIA/Megatron-LM.git
      cd Megatron-LM
      git checkout core_r0.6.0
    • core_r0.7.0分支
      git clone https://github.com/NVIDIA/Megatron-LM.git
      cd Megatron-LM
      git checkout core_r0.7.0

  4. “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

  5. “Megatron-LM”目录下修改pretrain_gpt.py文件,在“model_provider”函数中删除“assert(args.context_parallel_size == 1), "Context parallelism is only supported with Megatron Core!"”代码,如下黑体加粗部分。

        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
            )

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

    1. 下载Tokenizer

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

    2. 下载数据集。

      以Alpaca数据集为例,可单击Link获取。

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

  7. 配置环境变量,请用户根据set_env.sh的实际安装路径进行替换。

    source /usr/local/Ascend/ascend-toolkit/set_env.sh

  8. 数据处理,详情可单击对应分支core_r0.6.0core_r0.7.0进行参考。

    1. 语料格式转换。

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

      pip3 install nltk pyarrow pandas
      以下代码段展示了如何读取Parquet格式的原始语料,并将其转换为JSON格式,以便后续处理。
      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')
      
    2. 预训练数据集生成。
      若在昇腾设备上使用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

      通过运行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,代表预处理完成的预训练数据集。

  9. 配置路径。

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

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

  10. 执行如下命令启动预训练。

    bash examples/pretrain_gpt_distributed.sh

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