使用Megatron-Deepspeed训练GPT-2并生成文本

本文介绍如何使用GPU云服务器,使用Megatron-DeepSpeed框架训练GPT-2模型并生成文本。

背景信息

GPT-2模型是OpenAI于2018年在GPT模型的基础上发布的新的无监督NLP模型,当时被称为“史上最强通用NLP模型”。该模型可以生成连贯的文本段落,并且能在未经预训练的情况下,完成阅读理解、问答、机器翻译等多项不同的语言建模任务。GPT-2模型尤其在文本生成上有着惊艳的表现,其生成的文本在上下文连贯性和情感表达上都超过了人们的预期。

GPT-2模型由多层单向Transformer的解码器部分构成,根据模型层次规模,GPT-2模型有以下几个规格:

image

Megatron-Deepspeed框架结合了两种主要技术:

  • Megatron-LM是由NVIDIA开源的Transformer 模型框架。

  • DeepSpeed是由Microsoft开源的一个深度学习优化库。

DeepSpeed团队通过将DeepSpeed库中的ZeRO分片(ZeRO sharding)数据并行(Data Parallelism)和管道并行(Pipeline Parallelism)与Megatron-LM中的张量并行(Tensor Parallelism)相结合,开发了一种基于3D并行的实现,这就是Megatron-DeepSpeed,它使得千亿级参数量以上的大规模语言模型(LLM)的分布式训练变得更简单、高效和有效。

本文基于阿里云GPU服务器,使用Megatron-DeepSpeed框架训练GPT-2模型并生成文本。

image..png

重要
  • 阿里云不对第三方模型“GPT-2”的合法性、安全性、准确性进行任何保证,阿里云不对由此引发的任何损害承担责任。

  • 您应自觉遵守第三方模型的用户协议、使用规范和相关法律法规,并就使用第三方模型的合法性、合规性自行承担相关责任。

操作步骤

创建ECS实例

ECS实例创建页面,创建ECS实例。

关键参数说明如下,其他参数的配置,请参见自定义购买实例

  • 实例规格:选择ecs.gn7i-c8g1.2xlarge(单卡NVIDIA A10)。

    image

  • 镜像:使用云市场镜像,名称为aiacc-train-solution,您可以直接通过名称搜索该镜像,选择最新版本即可。选择镜像..png

  • 公网IP:选中分配公网IPv4地址,带宽计费方式选择按使用流量,带宽峰值选择100 Mbps,以加快模型下载速度。

    image..png

安装Megatron-Deepspeed框架

  1. 使用root用户远程登录ECS实例。具体操作,请参见通过密码或密钥认证登录Linux实例

  2. 执行以下命令,启动容器。

    docker run -d -t --network=host --gpus all --privileged --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --name megatron-deepspeed -v /etc/localtime:/etc/localtime -v /root/.ssh:/root/.ssh nvcr.io/nvidia/pytorch:21.10-py3
  3. 执行以下命令,进入容器终端。

    docker exec -it megatron-deepspeed bash
  4. 执行以下命令,下载Megatron-DeepSpeed框架。

    git clone https://github.com/bigscience-workshop/Megatron-DeepSpeed
  5. 执行以下命令,安装Megatron-DeepSpeed框架。

    cd Megatron-DeepSpeed
    pip install -r requirements.txt

处理数据

本指南使用1GB 79K-record的JSON格式的OSCAR数据集。

  1. 执行以下命令,下载数据集。

    wget https://huggingface.co/bigscience/misc-test-data/resolve/main/stas/oscar-1GB.jsonl.xz
    wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
    wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
  2. 执行以下命令,解压数据集。

    xz -d oscar-1GB.jsonl.xz
  3. 执行以下命令,预处理数据。

    python3 tools/preprocess_data.py \
        --input oscar-1GB.jsonl \
        --output-prefix meg-gpt2 \
        --vocab gpt2-vocab.json \
        --dataset-impl mmap \
        --tokenizer-type GPT2BPETokenizer \
        --merge-file gpt2-merges.txt \
        --append-eod \
        --workers 8

    如果回显信息类似如下所示,表示预处理数据完成。

    image..png

  4. 执行以下命令,新建data目录。

    mkdir data
  5. 执行以下命令,将处理好的数据移动到data目录下。

    mv meg-gpt2* ./data
    mv gpt2* ./data

预训练

本示例使用单机单卡的GPU实例完成GPT-2 MEDIUM模型的预训练。

  1. 创建预训练脚本文件。

    1. 执行以下命令,创建预训练脚本文件。

      vim pretrain_gpt2.sh
    2. i键,进入编辑模式,在文件中添加以下信息。

      #! /bin/bash
      
      # Runs the "345M" parameter model
      
      GPUS_PER_NODE=1
      # Change for multinode config
      MASTER_ADDR=localhost
      MASTER_PORT=6000
      NNODES=1
      NODE_RANK=0
      WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
      
      DATA_PATH=data/meg-gpt2_text_document
      CHECKPOINT_PATH=checkpoints/gpt2
      
      DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
      
      python -m torch.distributed.launch $DISTRIBUTED_ARGS \
             pretrain_gpt.py \
             --tensor-model-parallel-size 1 \
             --pipeline-model-parallel-size 1 \
             --num-layers 24 \
             --hidden-size 1024 \
             --num-attention-heads 16 \
             --micro-batch-size 4 \
             --global-batch-size 8 \
             --seq-length 1024 \
             --max-position-embeddings 1024 \
             --train-iters 5000 \
             --lr-decay-iters 320000 \
             --save $CHECKPOINT_PATH \
             --load $CHECKPOINT_PATH \
             --data-path $DATA_PATH \
             --vocab-file data/gpt2-vocab.json \
             --merge-file data/gpt2-merges.txt \
             --data-impl mmap \
             --split 949,50,1 \
             --distributed-backend nccl \
             --lr 0.00015 \
             --lr-decay-style cosine \
             --min-lr 1.0e-5 \
             --weight-decay 1e-2 \
             --clip-grad 1.0 \
             --lr-warmup-fraction .01 \
             --checkpoint-activations \
             --log-interval 10 \
             --save-interval 500 \
             --eval-interval 100 \
             --eval-iters 10 \
             --fp16
    3. Esc键,输入:wq后,按Enter键保存文件。

  2. 修改测试代码。

    Megatron源码有一个断言需要注释掉,以保证代码正常运行。

    1. 执行以下命令,打开测试代码文件。

      vim /workspace/Megatron-DeepSpeed/megatron/model/fused_softmax.py +191
    2. i键,进入编辑模式,在assert mask is None, "Mask is silently ignored due to the use of a custom kernel"前加#

      image..png

    3. Esc键,输入:wq后,按Enter键保存文件。

  3. 预训练。

    1. 执行以下命令,开始预训练。

      nohup sh ./pretrain_gpt2.sh &

      image..png

    2. 执行如下命令,可以持续地查看nohup.out的输出,达到监控程序的效果。

      tail -f nohup.out

      如果回显信息类似如下所示,表示预训练完成。

      image..png

      说明

      预训练完成大概需要1小时30分钟,如果超时断开了ECS连接,重新远程登录ECS实例后,执行以下命令,继续查看预训练进度。预训练完成后,可以执行Ctrl+Z命令退出。

      docker exec -it megatron-deepspeed bash
      cd Megatron-DeepSpeed
      tail -f nohup.out
  4. (可选)执行以下命令,查看生成的模型checkpoint路径。

    本示例生成的模型checkpoint路径设置在/workspace/Megatron-DeepSpeed/checkpoints/gpt2

    ll ./checkpoints/gpt2

    image..png

使用GPT-2模型生成文本

  1. 执行以下命令,安装相关依赖。

    说明

    由于网络原因,执行命令后可能会失败,建议您多次尝试。

    pip install mpi4py

    如果回显信息类似如下所示,表示依赖安装完成。

    image..png

  2. 创建文本生成脚本。

    1. 执行以下命令,创建文本生成脚本。

      vim generate_text.sh
    2. i键,进入编辑模式,在文件中增加以下内容。

      #!/bin/bash
      
      CHECKPOINT_PATH=checkpoints/gpt2
      VOCAB_FILE=data/gpt2-vocab.json
      MERGE_FILE=data/gpt2-merges.txt
      
      python tools/generate_samples_gpt.py \
             --tensor-model-parallel-size 1 \
             --num-layers 24 \
             --hidden-size 1024 \
             --load $CHECKPOINT_PATH \
             --num-attention-heads 16 \
             --max-position-embeddings 1024 \
             --tokenizer-type GPT2BPETokenizer \
             --fp16 \
             --micro-batch-size 2 \
             --seq-length 1024 \
             --out-seq-length 1024 \
             --temperature 1.0 \
             --vocab-file $VOCAB_FILE \
             --merge-file $MERGE_FILE \
             --genfile unconditional_samples.json \
             --num-samples 2 \
             --top_p 0.9 \
             --recompute
    3. Esc键,输入:wq后,按Enter键保存文件。

  3. 执行以下命令,生成文本。

    sh ./generate_text.sh

    如果回显信息类似如下所示,表示生成文本完成。

    image..png

  4. 执行以下命令,查看生成的JSON格式的文本文件。

    vim unconditional_samples.json

    回显信息类似如下所示。

    image..png

了解更多AIGC实践和GPU优惠

活动入口立即开启AIGC之旅

image..png