阿里云人工智能平台PAI基于ACCL优化套件进行多机训练-云淘科技

前提条件

• 可以访问灵骏节点并拥有管理员权限。

• 已安装和配置GPU驱动、RDMA驱动和Docker CE。

步骤一：准备容器镜像

• 本文推荐的镜像为registry.cn-hangzhou.aliyuncs.com/eflops/pytorch:22.02py38，该镜像已打包集成ACCL后端的PyTorch及Torchvision、Tensorboard、MPI等常用依赖与工具。

• 如果使用其他镜像，需要在NGC镜像安装部署ACCL，Dockerfile示例如下。

  # syntax=docker/dockerfile:experimental
  FROM nvcr.io/nvidia/pytorch:21.09-py3
  RUN curl http://eflops.oss-cn-beijing.aliyuncs.com/accl/ubuntu20.04/accl-v1.3.2-4-ubuntu20.04-cuda 11.4.sh -o /tmp/accl.sh && \
       bash /tmp/accl.sh && \
       rm -rf /tmp/accl.sh
  RUN (cd /tmp && curl -O http://eflops.oss-cn-beijing.aliyuncs.com/accl/torch/ngc/21.09/torch_accl-0.3.0-cp38-cp38-linux_x86_64.whl) && \
       pip install /tmp/torch_accl-0.3.0-cp38-cp38-linux_x86_64.whl && \
       rm -rf /tmp/torch_accl-0.3.0-cp38-cp38-linux_x86_64.whl
  RUN sed -i \
       -e '/is_nccl_backend =/s/=.*/= group_backend != Backend.MPI and group_backend != Backend. GLOO/g' \
       -e '/GroupNCCL/ {N; /):/ s/:/ or get_backend(group) == "accl":/}' \
       -e '/!= Backend.NCCL:/s/:/ and get_backend(group) != "accl":/' \
       /opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py

步骤二：下载训练示例代码与数据集

1. 下载PyTorch版BERT示例代码。本示例代码基于DeepLearningExamples，仅修改DDP版本、训练脚本选项等参数，增加通信后端使用ACCL的支持。

2. 在PyTorch中，使用ACCL需要对训练脚本做两点修改。更多信息，请参见安装ACCL库。

   1. 导入PyTorchACCL插件：import torch_accl。

   2. 在分布式初始化时指定后端为ACCL：init_process_group(backend='accl')。

3. 下载数据集。

   • 使用示例代码data目录下的脚本create_datasets_from_start.sh下载、处理Wikipedia语料库。

   • 下载非完整的英文语料数据集（约3.8GB）以便于快速验证，数据集已转换为可供BERT训练直接读取的格式。下载后解压至BERT示例代码所在路径下的data目录。

步骤三：多机训练

• 使用BERT示例代码中的run_pretraining.sh脚本启动多机多卡训练任务。

  1. 以2台灵骏节点16张GPU为例，假设其中某个灵骏节点的管理网络接口IP地址为10.35.110.11，在该灵骏节点运行：

     # docker run --name=bert_accl_test-0 --rm -it
     --net=host --gpus=all --shm-size=8g
     --device=/dev/infiniband --ulimit memlock=-1:-1
     -e MASTER_ADDR=10.35.110.11 -e MASTER_PORT=29500
     -e RANK=0 -e WORLD_SIZE=2 -e NPROC_PER_NODE=8
     -e PHASE=1 -e BATCH_SIZE=384 -e BACKEND=accl -e ACCL_LOG_LEVEL=INFO
     registry.cn-hangzhou.aliyuncs.com/eflops/pytorch:22.02py38
     /PATH/TO/BERT/run_pretraining.sh

  2. 在另一个灵骏节点运行：

     # docker run --name=bert_accl_test-1 --rm -it
     --net=host --gpus=all --shm-size=8g
     --device=/dev/infiniband --ulimit memlock=-1:-1
     -e MASTER_ADDR=10.35.110.11 -e MASTER_PORT=29500
     -e RANK=1 -e WORLD_SIZE=2 -e NPROC_PER_NODE=8
     -e PHASE=1 -e BATCH_SIZE=384 -e BACKEND=accl -e ACCL_LOG_LEVEL=INFO
     registry.cn-hangzhou.aliyuncs.com/eflops/pytorch:22.02py38
     /PATH/TO/BERT/run_pretraining.sh

• 其中，/PATH/TO/BERT代表BERT训练脚本所在路径。假设BERT训练代码位于/home/accl/BERT目录下，容器启动时挂载至目录-v/home/accl/BERT:/workspace/BERT，那么训练脚本的具体路径为/workspace/BERT/run_pretraining.sh。

说明

• 在任务调度系统中，主节点Pod的选取、Rank相关环境变量的预设由任务平台自动配置，用户只需在训练代码中根据环境变量确定所使用的GPU设备、读取对应的数据分片。底层的设置例如GPU与网卡之间的亲和关系等由ACCL自动探测拓扑来决定，用户通常无需感知。

• 在BERT论文中，预训练分成了Phase-1与Phase-2两个阶段。二者的主要区别为Sequence length不同。虽然更长的序列有助于提高准确性，但会大幅增加计算耗时。为了控制整体的训练耗时，原文中Phase-1使用128的序列长度完成了36个Epoch的训练，Phase-2使用512的序列长度完成了剩余的4个Epoch。本文提供了与原文一致的序列长度128与512的两组数据集，使用run_pretraining.sh脚本运行时可通过环境变量PHASE来指定。