拉起多卡分布式训练

在单机和多机场景下，有4种方式可拉起分布式训练，分别为shell脚本方式（推荐）、mp.spawn方式、Python方式、torchrun方式。其中torchrun方式仅在PyTorch 1.11.0及以上版本支持使用。以下内容以一个简单模型脚本为样例，展示前3种拉起方式分别需要对脚本代码进行的修改。torchrun方式的代码修改与shell脚本方式完全相同。

  1. 集合通信存在如下约束：
    • 数据并行模式中不同device上执行的图相同。
    • 针对Atlas 训练系列产品：allreduce和reduce_scatter仅支持int8、int32、float16和float32数据类型。
    • 针对Atlas A2 训练系列产品：allreduce和reduce_scatter仅支持int8, int32, float16, float32和bfp16数据类型。
    • 分布式训练场景下，HCCL会使用Host服务器的部分端口进行集群信息收集，需要操作系统预留该部分端口。默认情况下，HCCL使用60000-60015端口，若通过环境变量HCCL_IF_BASE_PORT指定了Host网卡起始端口，则需要预留以该端口起始的16个端口。

      操作系统端口号预留示例：

      sysctl -w net.ipv4.ip_local_reserved_ports=60000-60015
  2. 若用户准备进行2卡训练，可将8卡训练脚本进行改写，改为2卡训练脚本。可参见以下修改方法：
    1. 若8卡脚本的batchsize是单卡脚本的batchsize的8倍，则将8卡训练时的batch size和learning rate同时除以4，作为2卡训练时的batch size和learning rate。
    2. 如果使用for循环启动训练入口脚本，则将for循环的次数改为2次。
    3. world size或者rank size修改为2，并确保训练脚本中dist.init_process_group()中world_size参数为2。
    4. 如果有指定device list参数，且取值取值范围为0-7，则将其改为0-1。

构建简单模型

我们先构建一个简单的神经网络。

# 导入依赖和库
import torch
from torch import nn
import torch_npu
import torch.distributed as dist
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import time
import torch.multiprocessing as mp
import os

torch.manual_seed(0)
# 下载训练数据
training_data = datasets.FashionMNIST(
    root="./data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# 下载测试数据
test_data = datasets.FashionMNIST(
    root="./data",
    train=False,
    download=True,
    transform=ToTensor(),
)

# 构建模型
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

获取超参数

在主函数main中获取训练所需的超参数。

设置地址和端口号

由于昇腾AI处理器初始化进程组时initmethod只支持env:// （即环境变量初始化方式），所以在初始化前需要配置MASTER_ADDR、MASTER_PORT等参数。用户需根据自己实际情况配置。

添加分布式逻辑

不同的拉起训练方式下，device号的获取方式不同：

用户需根据自己选择的方式对代码做不同的修改。

设置参数

在模型脚本中，根据拉起方式不同，设置不同的参数。

拉起训练

以下拉起训练的命令为示例，用户可根据实际情况自行更改。

当屏幕打印类似下图中的Loss数值时，说明拉起训练成功。

父主题： 多卡分布式训练