提交PyTorch分布式作业#

PAI支持用户提交分布式PyTorch训练作业，本文将介绍如何使用PAI Python SDK，以PyTorch DDP(DistributedDataParallel)模式提交分布式PyTorch训练作业。

安装和配置SDK#

我们需要首先安装PAI Python SDK以运行本示例。

!python -m pip install --upgrade alipai

!python -m pip install pygments

SDK需要配置访问阿里云服务需要的AccessKey，以及当前使用的工作空间和OSS Bucket。在PAI SDK安装之后，通过在 命令行终端 中执行以下命令，按照引导配置密钥、工作空间等信息。

# 以下命令，请在 命令行终端 中执行.

python -m pai.toolkit.config

我们可以通过以下代码验证配置是否已生效。

import pai
from pai.session import get_default_session

print(pai.__version__)

sess = get_default_session()

# 获取配置的工作空间信息
assert sess.workspace_name is not None
print(sess.workspace_name)

PyToch 分布式作业介绍#

PyTorch DDP（Distributed Data Parallel）是PyTorch提供的分布式数据并行训练功能，支持模型在多台机器上进行并行训练，从而提高训练效率。

PyTorch DDP基于多进程的方案实现，支持单机多卡模式和多机多卡模式。在单机多卡模式下，用户可以使用同一台机器下的多个GPU来加速模型的训练。在多机多卡模式下，可以将计算任务分配到多台机器上进行并行计算，加速训练速度。对于DDP的详细介绍，可以参考PyTorch的官方文档链接。

PyTorch DDP

PyTorch提供的DataParallel和DistributedDataParallel模块都支持数据并行训练，PyTorch官方推荐不论是单机多卡还是多机多卡，都使用DistributedDataParallel模块进行训练。

代码适配DDP改造#

使用PyTorch DDP进行分布式训练需要对原先的PyTorch训练代码（使用单机单卡）进行的修改，具体可以参考PyTorch官方文档。

主要包括：

  • 初始化分布式训练配置:

需要在训练迭代开始之前完成训练环境初始化。


from torch.distributed import init_process_group, destroy_process_group

def ddp_setup()
    init_process_group(backend="nccl")

初始化需要指定机器之间的通讯方式，当使用GPU进行训练时，通常使用nccl作为通讯后端，而使用CPU训练时，使用gloo，详细的介绍可以参考PyTorch文档: Which Backend To Use?

  • 使用DDP封装模型：


from torch.nn.parallel import DistributedDataParallel as DDP

# model是原始单机单卡训练的PyTorch模型
model = DDP(model)

  • 修改DataLoader的采样方式：

当使用DDP进行数据并行训练，不同的worker进程需要读取不同的数据分片进行训练。当不同机器上通过共享存储的方式使用同一份数据集时，可以使用torch.utils.data.distributed.DistributedSampler来对数据进行采样，从而保证不同的worker进程读取不同的数据分片。


from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

train_sampler = DistributedSampler(
	train_dataset,
	shuffle=True)

trainloader = DataLoader(
	train_dataset,
	batch_size=args.per_device_train_batch_size,
	sampler=train_sampler,
	num_workers=2,
	drop_last=True)

PAI支持PyTorch DDP分布式训练#

PAI原生支持的PyTorch的分布式训练，当用户提交训练作业，指定作业类型为PyTorch训练作业时(job_type="PyTorchJob")，PAI的训练服务会在机器节点上设置环境变量，包含作业机器数量，机器RANK，机器之间的通讯地址等信息。

环境变量名

描述

MASTER_ADDR

Master机器节点的服务地址

MASTER_PORT

Master机器节点端口号，如：23456

WORLD_SIZE

分布式作业的机器节点总数，例如提交的训练作业申请了4台机器，则WORLD_ISZE=4

RANK

机器节点的RANK，例如启动了一个4个节点的作业，则各个机器节点的RANK分别为0,1,2,3

PyTorch提供了分布式训练启动工具，torchrun(PyTorch 1.1.0及以上版本) 和 torch.distributed.launch(PyTorch 1.1.0版本以下)，支持训练作业的拉起。配合以上PAI预置的环境变量，我们可以便利得启动分布式训练作业。

使用torch.distributed.launch拉起训练作业示例：

# for PyTorch<1.1.0

python -m torch.distributed.launch \
--nproc_per_node=<NumberOrProcessPerNode> \
--master_addr=$MASTER_ADDR \
--master_port=$MASTER_PORT \
--nnodes=$WORLD_SIZE \
--node_rank=$RANK \
<YourTrainingScript> training_arguments...

使用torchrun拉起训练作业示例：

# for PyTorch>=1.1.0
torchrun \
--nproc_per_node=<NumberOrProcessPerNode> \
--master_addr=$MASTER_ADDR \
--master_port=$MASTER_PORT \
--nnodes=$WORLD_SIZE \
--node_rank=$RANK \
<YourTrainingScript> training_arguments...

用户需要修改<NumberOfProcessPerNode为每一个机器节点需要启动的进程数，通常设置为机器节点的GPU数量。

以上的作业启动命令，同样适用于单机多卡的训练作业启动。

提交训练作业#

PAI Python SDK提供了Estimator的接口，用于提交训练作业，以下示例中，我们将通过Estimator提交一个PyTorch分布式训练作业。

  • 准备训练代码

PyTorch提供了多机多卡的训练代码示例，在修改了模型和checkpoints保存路径后，我们既可以将其用于提交到PAI进行训练。

# 通过以下代码查看准备提交的训练代码
!pygmentize train_src/train_multinode.py

  • 提交训练作业

我们将使用PAI提供的PyTorch 1.12版本的GPU镜像完成多机多卡的作业训练。使用estimator.fit提交训练作业之后，SDK会打印作业的控制台链接，用户可以通过控制台查看作业状态，日志详情等信息。

from pai.estimator import Estimator
from pai.image import retrieve

# 使用PAI提供的PyTorch 1.12 GPU镜像
image_uri = retrieve(
    "pytorch",
    framework_version="1.12",
    accelerator_type="GPU",
).image_uri
print("Training Image URI: ", image_uri)


# 每一个机器实例的GPU数量，需要根据用户选择的机器型号(instance_type)进行修改
gpu_count_per_instance = 2

# 训练脚本使用torchrun命令启动
command = f"""torchrun --master_addr=$MASTER_ADDR \
--master_port=$MASTER_PORT \
--nnodes=$WORLD_SIZE --node_rank=$RANK \
--nproc_per_node={gpu_count_per_instance} \
train_multinode.py --total_epochs 10 --save_every 5 \
"""


# 提交训练作业
est = Estimator(
    image_uri=image_uri,
    source_dir="./train_src",  # 训练代码所在目录
    command=command,
    job_type="PyTorchJob",
    instance_type="ecs.gn6i-c24g1.12xlarge",  # 2 * NVIDIA T4 GPU
    instance_count=2,  # 机器实例数量
    base_job_name="pytorch-ddp",
)

# fit方法提交训练作业，默认等待到作业执行完成
est.fit()


# 查看作业的输出模型

est.model_data()

参考：#

  • PyTorch Distributed Overview: https://pytorch.org/tutorials/beginner/dist_overview.html