快速上手

按照如下步骤操作,即可实现Megatron-LM在昇腾设备上的高效运行,且无缝集成并充分发挥MindSpeed所提供的丰富加速与优化技术。

  1. “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
  2. 数据准备,参考Megatron-LM官方文档准备训练数据。
    1. 下载Tokenizer

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

    2. 下载数据集。

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

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

  3. 配置环境变量,当前以root用户安装后的默认路径为例,请用户根据set_env.sh的实际路径执行如下命令。
    source /usr/local/Ascend/ascend-toolkit/set_env.sh
  4. 数据处理,详情可单击对应分支core_r0.8.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

      新建“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,代表预处理完成的预训练数据集。

  5. “Megatron-LM”目录下准备预训练脚本train_distributed.sh,脚本示例如下:
     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    #!/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
    
  6. 配置路径。
    请编辑示例脚本train_distributed.sh,并设置如下环境变量以指定必要的路径:
    1
    2
    3
    4
    CKPT_DIR=./ckpt
    VOCAB_FILE=./gpt-tokenizer/vocab.json
    MERGE_FILE=./gpt-tokenizer/merges.txt
    DATA_PATH=./gpt_pretrain_data/alpaca_text_document
    

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

  7. 执行如下命令启动预训练。
    bash ./train_distributed.sh
  • 示例脚本train_distributed.sh中的部分参数(如--hidden-size、--num-layers)需根据实际场景进行适配,避免OOM等现象。
  • 更多模型开发实例请参见模型开发实践