文档

基于ACCL优化套件进行多机训练

更新时间:

本文以BERT-Large模型为例,介绍如何使用ACCL优化套件进行PyTorch分布式训练。

前提条件

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

  • 已安装和配置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来指定。