Pytorch多机多卡的多种打开方式

前言

为了获取最佳阅读体验，推荐移步个人博客
在上一篇介绍多卡训练原理的基础上，本篇主要介绍Pytorch多机多卡的几种实现方式：DDP、multiprocessing、Accelerate。

在介绍具体实现之前，torch.distributed 涉及的分布式概念如下：

• **group：**进程组，通常一个job只有一个组，即一个world，使用多机时，一个group产生了多个world。
• **world_size：**一个job的全局进程数量
• **rank：**进程的序号，一般设置rank=0的主机为master节点。
• **local_rank：**进程内部的GPU序号。

比如，有两台8卡机器，这时具有一个group，2个world，每个world_size为8，第一个主机rank=0，显卡编号依次为0,…,7，第二个主机rank=1，显卡编号依次为0,…,7。

在多机多卡的分布式训练过程中，为每个进程的模型、数据配置好这些参数至关重要。

DDP

Pytorch分布式执行流程如下：

1. init_process_group 初始化进程组，同时初始化 distributed 包。
2. 创建分布式模型model = DDP(model)
3. 创建分布式数据采样的datasampler
4. 利用torch.distributed.launch控制进程训练
5. destory_process_group销毁进程组

进程组初始化

init_process_group(backend, 
                   init_method=None, 
                   timeout=datetime.timedelta(0, 1800), 
                   world_size=-1, 
                   rank=-1, 
                   store=None)

TCP初始化

使用TCP初始化时，需要指定下列参数：

1. rank 为当前进程的进程号
2. word_size 为当前 job 的总进程数
3. init_method 内指定 tcp 模式，且所有进程的 ip:port 必须一致，设定为主进程的 ip:port

初始化时，需要注意下列事项：

1. 在 rank==0 的进程内保存参数,一般是rank0主节点来分发广播梯度。

2. 若程序内未根据 rank 设定当前进程使用的 GPUs，则默认使用全部 GPU，且以数据并行的方式使用。

3. 每条命令表示一个进程，若已开启的进程未达到 word_size 的数量，则所有进程会一直等待。

4. 每台主机上可以开启多个进程。但是，若未为每个进程分配合适的 GPU，则同机不同进程可能会共用 GPU，应该坚决避免这种情况，容易爆显存。

5. 使用 gloo 后端进行 GPU 训练时，会报错。

参考代码如下，需要在args里面添加指定的参数：

import torch.distributed as dist
import torch.utils.data.distributed

# ......
parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,
                    help='rank of current process')
parser.add_argument('--word_size', default=2,
                    help="word size")
parser.add_argument('--init_method', default='tcp://127.0.0.1:23456',
                    help="init-method")
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.word_size)

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
net = Net()
net = net.cuda()
net = torch.nn.parallel.DistributedDataParallel(net)

执行脚本如下：

# Node 1 : ip 192.168.1.201  port : 12345
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 0 --word_size 3

# Node 2 : 
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 1 --word_size 3

# Node 3 : 
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 2 --word_size 3

ENV初始化

在ENV初始化方式中，init中无需指定参数，主要从机器的环境变量中获取参数。

1. 该初始化中需要设定local_rank参数，确定单机进程的序号。

2. 然后，通过torch.distributed.launch设定nnodes节点数，node_rank当前主机进程序号，nproc_per_node每个节点的进程数量，master_addr主节点地址，master_port主节点端口，在环境变量中获取这些参数。

注意事项如下：

1. 使用 torch.distributed.launch 工具时，将会为当前主机创建 nproc_per_node 个进程，每个进程独立执行训练脚本。同时，它还会为每个进程分配一个 local_rank 参数，表示当前进程在当前主机上的编号。例如：rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。

2. 在 rank==0 的进程内保存参数。

3. Env 方式中，在 init_process_group 中，无需指定任何参数

4. 合理利用 local_rank 参数，来合理分配本地的 GPU 资源

5. 每条命令表示一个进程。若已开启的进程未达到 word_size 的数量，则所有进程会一直等待。

参考代码如下：

import torch.distributed as dist
import torch.utils.data.distributed

# ......
import argparse
parser = argparse.ArgumentParser()
# 注意这个参数，必须要以这种形式指定，即使代码中不使用。因为 launch 工具默认传递该参数
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method='env://')

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
# 根据 local_rank，配置当前进程使用的 GPU
net = Net()
device = torch.device('cuda', args.local_rank)
net = net.to(device)
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)

执行脚本如下：

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=0 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=1 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=2 --master_addr="192.168.1.201" --master_port=23456 env_init.py

共享文件系统初始化

使用共享文件系统初始化时，与TCP初始化类似，需要指定下列参数：

1. rank 为当前进程的进程号
2. word_size 为当前 job 的总进程数
3. init_method 内指定 文件系统 模式，以 file:// 为前缀，表示文件系统各式初始化。/xxx 表示共享的文件，各个进程在共享文件系统中通过该文件进行同步或异步。因此，所有进程必须对该文件具有读写权限。

参考代码如下：

mport torch.distributed as dist

# ......
parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,
                    help='rank of current process')
parser.add_argument('--word_size', default=2,
                    help="word size")
parser.add_argument('--init_method', default='file:///mnt/nfs/sharedfile',
                    help="init-method")
args = parser.parse_args()

# rank should always be specified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, rank=args.rank)

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
# 根据 local_rank，配置当前进程使用的 GPU
net = Net()
device = torch.device('cuda', args.local_rank)
net = net.to(device)
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank

执行脚本如下：

Node1:
python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 0 --world_size 2
Node2:
python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 1 --world_size 2

DistributedDataParallel

torch.nn.parallel.DistributedDataParallel(module, 
                                          device_ids=None, 
                                          output_device=None, 
                                          dim=0, 
                                          broadcast_buffers=True, 
                                          process_group=None, 
                                          bucket_cap_mb=25, 
                                          find_unused_parameters=False, 
                                          check_reduction=False)

将给定的 module 进行分布式封装， 其将输入在 batch 维度上进行划分，并分配到指定的 devices 上。

module 会被复制到每台机器的每个 进程 上，每一个模型的副本处理输入的一部分。

在反向传播阶段，每个机器的每个 进程 上的梯度进行汇总并求平均。与 DataParallel 类似，batch size 应该大于 GPU 总数。

主要参数介绍：

1. **module：**将完整的model封装为分布式module,后续需要调用model的方法时，可以采用module.model.xxx

2. **device_ids：**需要并行的设备，在数据并行的情况下，表示模型副本拷贝到哪些GPU上；在模型并行的情况下，表示模型分散在哪些GPU上。

3. **output_device：**输出结果到哪个GPU上。

注意事项如下：

1. 要使用该 class，需要先对 torch.distributed 进行初进程组始化，可以通过 torch.distributed.init_process_group() 实现。

2. 该 module 仅在 gloo和 nccl后端上可用。

DistributedSampler

torch.utils.data.distributed.DistributedSampler(dataset, 
                                                num_replicas=None, 
                                                rank=None)

主要参数介绍：

1. **dataset：**采样的数据集

2. **num_replicas：**参与的总进程数

3. **rank：**当前机器的rank

DistributedSampler将数据集采样为num_replicas份，不同机器根据自己的rank取数据集的子集。

TIPS：在 DataParallel 中，batch size 设置必须为单卡的 n 倍，但是在 DistributedDataParallel 内，batch size 设置于单卡一样即可。

参考代码如下：

# 分布式训练示例
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel

dataset = your_dataset()
datasampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=datasampler)
model = your_model()
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)

torch.distributed.launch

DDP通过torch.distributed.launch辅助实现进程控制。

torch.distributed.launch传入的参数如下：

• **training_script：**执行任务脚本路径
• **–nnodes：**节点数，即分布式机器数量
• **–node_rank：**当前机器的rank序号
• **–nproc_per_node：**每个节点开设的进程数量，最好设置为每个机器GPU数量，使每个GPU在一个进程中
• **–master_addr：**master 节点（rank 为 0）的地址
• **–master_port：**master 节点（rank 为 0）的端口

单机多卡的执行脚本如下：

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

多机多卡执行脚本如下：

Node1:
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes= NUM_MACHINES_YOU_HAVE --node_rank=0 --master_addr="192.168.1.1" --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
...
NodeN:
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes= NUM_MACHINES_YOU_HAVE --node_rank=N-1 --master_addr="192.168.1.1" --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

torch.multiprocessing

通过 torch.multiprocessing手动控制进程，替代 torch.distributed.launch的进程控制工作。

涉及的主要接口为：

def spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method='spawn'):
    r"""Spawns ``nprocs`` processes that run ``fn`` with ``args``.

    If one of the processes exits with a non-zero exit status, the
    remaining processes are killed and an exception is raised with the
    cause of termination. In the case an exception was caught in the
    child process, it is forwarded and its traceback is included in
    the exception raised in the parent process.

    Args:
        fn (function): Function is called as the entrypoint of the
            spawned process. This function must be defined at the top
            level of a module so it can be pickled and spawned. This
            is a requirement imposed by multiprocessing.

            The function is called as ``fn(i, *args)``, where ``i`` is
            the process index and ``args`` is the passed through tuple
            of arguments.

        args (tuple): Arguments passed to ``fn``.
        nprocs (int): Number of processes to spawn.
        join (bool): Perform a blocking join on all processes.
        daemon (bool): The spawned processes' daemon flag. If set to True,
                       daemonic processes will be created.
        start_method (string): (deprecated) this method will always use ``spawn``
                               as the start method. To use a different start method
                               use ``start_processes()``.

    Returns:
        None if ``join`` is ``True``,
        :class:`~ProcessContext` if ``join`` is ``False``

    """

主要参数介绍如下：

• fn:处理的主函数
• args:传递给主函数的参数，主函数第一个参数默认传入进程index
• nprocs:开启的进程数量

结合下列代码介绍torch.multiprocessing多机多卡的使用：

def setup(rank, world_size):
    # dist.init_process_group("gloo", rank=rank, world_size=world_size)
    print("world size:", world_size, " rank:", rank)
    print(os.environ['MASTER_ADDR'])
    print(os.environ['MASTER_PORT'])
    print(os.environ['RANK'])
    print(os.environ['WORLD_SIZE'])
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def main(local_rank, nnodes, args):
    rank = int(os.environ['RANK']) * nnodes + local_rank
    world_size = nnodes * int(os.environ['WORLD_SIZE'])

    print("world size:", world_size, " rank:", rank)
    setup(rank, world_size)
    ……
    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)
    model = torch.nn.parallel.DistributedDataParallel(model.to(local_rank), device_ids=[local_rank])


if __name__ == "__main__":
    args = parse_args()
    world_size = torch.cuda.device_count()
    print('{}:{}'.format(world_size, '---' * 100))
    mp.spawn(main, args=(world_size, args), nprocs=world_size, join=True)

代码流程的解释如下：

1. 根据**torch.cuda.device_count()**获取单机的显卡数量，决定开启的进程数，即一个world的world_size
2. mp.spawn开启多进程
3. 单机的进程index即为local_rank，nnodes代表单机显卡数量，os.environ[‘RANK’]获取机器的rank值，通过rank*nnodes + local_rank 计算全局训练的索引，nnodes * int(os.environ[‘WORLD_SIZE’]) 计算全局训练的进程数量
4. 根据计算的全局索引，全局数量 初始化进程通信
5. **model.to(local_rank)**将模型放置于本地单机的显卡上

Accelerate

Hugging Face发布PyTorch新库「Accelerate」：适用于多GPU、TPU、混合精度训练。「Accelerate」提供了一个简单的 API，将与多 GPU 、 TPU 、 fp16 相关的样板代码抽离了出来，保持其余代码不变。PyTorch 用户无须使用不便控制和调整的抽象类或编写、维护样板代码，就可以直接上手多 GPU 或 TPU。

项目地址：https://github.com/huggingface/accelerate

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

- device = 'cpu'
+ accelerator = Accelerator()

- model = torch.nn.Transformer().to(device)
+ model = torch.nn.Transformer()
  optimizer = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optimizer, data = accelerator.prepare(model, optimizer, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
-         source = source.to(device)
-         targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

模型保存加载：

# 模型保存
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
accelerator.save(unwrapped_model.state_dict(), path)

# 模型加载
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.load_state_dict(torch.load(path))

具体代码可以参考huggingface的transformer代码库

项目地址：https://github.com/huggingface/transformers/tree/main/examples/pytorch

参考链接

https://zhuanlan.zhihu.com/p/462453622

https://zhuanlan.zhihu.com/p/98535650

https://zhuanlan.zhihu.com/p/76638962