Megatron-Deepspeed 预训练 GPT-小白踩坑与解决

记录在使用 megatron deepspeed 训练大模型 gpt 时，遇到的坑

1. 成功运行的最终步骤

1.1 配置分布式环境：

1.1.1 硬件配置：

3090 单机四卡

root@xxxx:/workspace# nvidia-smi -L
GPU 0: NVIDIA GeForce RTX 3090 (UUID: GPU-xxxx)
GPU 1: NVIDIA GeForce RTX 3090 (UUID: GPU-xxxx)
GPU 2: NVIDIA GeForce RTX 3090 (UUID: GPU-xxxx)
GPU 3: NVIDIA GeForce RTX 3090 (UUID: GPU-xxxx)

1.1.2 基础配置：

• 1.1.2.1 拉取 nvidia 的 docker 镜像，通过命令：
  其中参数 --shm-size 在训练大模型的 TP（模型并行） 和 PP（流水并行） 特别重要，xxxx 是自己选择的 nvidia docker 的版本，通过在官网选择合适的填入，例如 21.07：

  docker run --gpus all -it --rm --runtime=nvidia --shm-size=16g --ulimit memlock=-1 --privileged --net host nvcr.io/nvidia/pytorch:21.09-py3
  

• 1.1.2.2 下载各种包，reference：https://github.com/bigscience-workshop/Megatron-DeepSpeed

  # 基础包
  git clone https://github.com/bigscience-workshop/Megatron-DeepSpeed
  cd Megatron-DeepSpeed
  pip install -r requirements.txt
  
  # 可恶的 apex 包，由于在 docker 环境里，所以安装不会有问题
  git clone https://github.com/NVIDIA/apex
  cd apex
  pip install --global-option="--cpp_ext" --global-option="--cuda_ext" --no-cache -v --disable-pip-version-check .  2>&1 | tee build.log
  
  
  # 安装 deepspeed 
  git clone https://github.com/microsoft/deepspeed
  cd deepspeed
  rm -rf build
  TORCH_CUDA_ARCH_LIST="7.0" DS_BUILD_CPU_ADAM=1 DS_BUILD_AIO=1 DS_BUILD_UTILS=1 pip install -e . --global-option="build_ext" --global-option="-j8" --no-cache -v --disable-pip-version-check
  

1.2 首先测试机器上是否有分布式运行的环境，使用简单的 alexnet 的流水并行做测试：

microsoft 的 deepspeedexample 库提供的 alexnet 的 pp 可以直接用，（过于简单所以可以用来测试分布式环境），也可以用来理解 deepspeed 实现流水的逻辑，测试步骤：
默认配置的 dp=2,pp=2

```bash
git clone https://github.com/microsoft/DeepSpeedExamples.git
cd DeepSpeedExamples/pipeline
sh run.sh
```

1.3 下载数据以及预处理：

主要参考：

```bash
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
xz -d oscar-1GB.jsonl.xz
python tools/preprocess_data.py \
    --input oscar-1GB.jsonl \
    --output-prefix my-gpt2 \
    --vocab gpt2-vocab.json \
    --dataset-impl mmap \
    --tokenizer-type GPT2BPETokenizer \
    --merge-file gpt2-merges.txt \
    --append-eod \
    --workers 8
```

1.4 RUN

根据这个提供的 https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/start_fast.md 的快速 run 脚本，配置好对应的数据路径，和自定义 pp dp mp 数量。

1.4.1 主要需要修改以下的数据路径：

VOCAB_FILE=data/gpt2-vocab.json
MERGE_FILE=data/gpt2-merges.txt
DATA_PATH=data/meg-gpt2-oscar-en-10k_text_document
TENSORBOARD_PATH=output_dir/tensorboard

1.4.2 快速运行脚本：

CHECKPOINT_PATH=checkpoints/gpt2_4

VOCAB_FILE=data/gpt2-vocab.json
MERGE_FILE=data/gpt2-merges.txt
DATA_PATH=data/meg-gpt2-oscar-en-10k_text_document
TENSORBOARD_PATH=output_dir/tensorboard

N_GPUS=4
MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=16
TP_SIZE=1
PP_SIZE=2

NLAYERS=2
NHIDDEN=8
NHEADS=2
SEQ_LEN=512
VOCAB_SIZE=50257

SAVE_INTERVAL=50

TRAIN_SAMPLES=10_000

GPT_ARGS=" \
    --num-layers $NLAYERS \
    --hidden-size $NHIDDEN \
    --num-attention-heads $NHEADS \
    --seq-length $SEQ_LEN \
    --max-position-embeddings $SEQ_LEN \
    --micro-batch-size $MICRO_BATCH_SIZE \
    --rampup-batch-size 2 2 1_000 \
    --global-batch-size $GLOBAL_BATCH_SIZE \
    --train-samples $TRAIN_SAMPLES \
    --optimizer adam \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --adam-eps 1e-8 \
    --lr 1e-4 \
    --lr-warmup-samples 5 \
    --min-lr 1e-6 \
    --lr-decay-style cosine \
    --lr-decay-samples 12 \
    --clip-grad 1.0 \
    --weight-decay 1e-1 \
    --embed-layernorm \
    --fp16 \
    --partition-activations \
    --seed 42 \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
    "

OUTPUT_ARGS=" \
    --exit-interval 100 \
    --log-interval 10 \
    --save-interval $SAVE_INTERVAL \
    --eval-interval 100 \
    --eval-iters 10 \
    --checkpoint-activations \
    "

DATA_ARGS=" \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH \
    --data-path $DATA_PATH \
    --tensorboard-dir $TENSORBOARD_PATH \
    --tensorboard-queue-size 5 \
    --log-timers-to-tensorboard \
    --log-batch-size-to-tensorboard \
    --log-validation-ppl-to-tensorboard \
    --kill-switch-path /tmp/kill-switch \
    "

ZERO_STAGE=1

config_json="./ds_config.json"

# Deepspeed figures out GAS dynamically from dynamic GBS via set_train_batch_size()
cat <<EOT > $config_json
{
  "train_micro_batch_size_per_gpu": $MICRO_BATCH_SIZE,
  "train_batch_size": $GLOBAL_BATCH_SIZE,
  "gradient_clipping": 1.0,
  "zero_optimization": {
    "stage": $ZERO_STAGE
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 500,
    "hysteresis": 2,
    "min_loss_scale": 1,
    "initial_scale_power": 12
  },
  "steps_per_print": 2000,
  "wall_clock_breakdown": false
}
EOT

DEEPSPEED_ARGS=" \
    --deepspeed \
    --deepspeed_config ${config_json} \
    --zero-stage ${ZERO_STAGE} \
    --deepspeed-activation-checkpointing \
    "

ALL_ARGS="$GPT_ARGS $OUTPUT_ARGS $DATA_ARGS $DEEPSPEED_ARGS"

MASTER_ADDR=localhost
MASTER_PORT=6777

export LAUNCHER="python -u -m torch.distributed.run \
    --nproc_per_node $N_GPUS \
    --nnodes 1 \
    --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \
    --rdzv_backend c10d \
    --max_restarts 0 \
    --tee 3 \
    "
export CMD=" \
    $LAUNCHER pretrain_gpt.py \
    --tensor-model-parallel-size $TP_SIZE \
    --pipeline-model-parallel-size $PP_SIZE \
    --distributed-backend nccl \
    $ALL_ARGS \
    "

echo $CMD

$CMD

2. 遇到的坑以及解决方案

2.1 包的转换使用

一开始用的是 microsoft 的官方包，（地址：https://github.com/microsoft/Megatron-DeepSpeed），奈何他实在是写的太简单了，难以入门开始 run。
于是换一个库：（地址：https://github.com/bigscience-workshop/Megatron-DeepSpeed） 这里的文档比较清楚。

但是 microsoft 提供了一个 base 就是需要拉取 nvidia 的 docker 在 docker 上跑，这样可以避免掉一些包冲突不匹配的问题。比如接下来的 apex 的问题

2.2 complex apex

使用 megetron deepspeed 运行 bert 总是报 apex 包的问题。
按照 https://github.com/bigscience-workshop/Megatron-DeepSpeed 的安装方法是不行der，原因是 apex 不被维护很久了 有很多坑。

于是按照 microsoft 官方提供的 docker，这里配置好了

https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch

find/check in this link and run:

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:xx.xx-py3

in this code,it will give u a contaier,

• find a docker image at the left bar,the container opened for u will avalible at the top,
• right click and click the attach choice.
• get into the vscode open the workspace directory
• check version by run

conda list
nvcc -V

2.3 nccl bug in docker

在 docker 里可以运行数据并行但没办法运行模型并行和流水并行，报了很多 nccl 的错误

• 跑两次就报 no space left on devices,通过重启 docker 容器解决 restart
• 接下来就一直报 nccl 的错误，bus error。
  • 按照官方文档要求，通过在终端设置环境变量 export NCCL_DEBUGINFO 查看更加详细的报错
  • 报错 but store->get('0') got error: Connection reset by peer。。。 网上没有关于这个报错的解释
• ncclInternalError: Internal check failed. This is either a bug in NCCL or due to memory corruption 就是这个错误让我怀疑是容器内分配的空间不够
  通过查找 nccl 文档 ,发现这么一段话

In order to share data between ranks, NCCL may require shared system memory for IPC and pinned (page-locked) system memory resources. The operating system’s limits on these resources may need to be increased accordingly. Please see your system’s documentation for details. In particular, Docker containers default to limited shared and pinned memory resources. When using NCCL inside a container, it is recommended that you increase these resources by issuing:

–shm-size=1g –ulimit memlock=-1
in the command line to nvidia-docker run.

发现需要配置和扩展 –shm-size 的大小。
所以在拉取和run docker 需配置好 docker 内的内存大小，所以应该运行以下语句，重开一个 docker。
bash docker run --gpus --runtime=nvidia --shm-size=16g --ulimit memlock=-1 --privileged --net host all -it --rm nvcr.io/nvidia/pytorch:xx.xx-py3

2.4 生成 bin 和 idx 文件时，出现 premission denied 的问题：

通过命令 ls -l 查看文件夹的权限，发现megatron 包是 qlchen 的权限，然后megatron 内的 data 权限是 root 权限，需要 root 用到 chmod 修改 data 文件夹的权限。到底为什么没有 root 权限却能创建 data 文件夹，目前还不知道。