pytorch多GPU训练教程

pytorch多GPU训练教程

1. Torch 的两种并行化模型封装

1.1 DataParallel

DataParallel 是 PyTorch 提供的一种数据并行方法，用于在单台机器上的多个 GPU 上进行模型训练。它通过将输入数据划分成多个子部分（mini-batches），并将这些子部分分配给不同的 GPU，以实现并行计算。
在前向传播过程中，输入数据会被划分成多个副本并发送到不同的设备（device）上进行计算。模型（module）会被复制到每个设备上，这意味着输入的批次（batch）会被平均分配到每个设备，但模型会在每个设备上有一个副本。每个模型副本只需要处理对应的子部分。需要注意的是，批次大小应大于GPU数量。在反向传播过程中，每个副本的梯度会被累加到原始模型中。总结来说，DataParallel会自动将数据切分并加载到相应的GPU上，将模型复制到每个GPU上，进行正向传播以计算梯度并汇总。
注意：DataParallel是单进程多线程的，仅仅能工作在单机中。
封装示例：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 初始化模型
model = SimpleModel()

# 使用 DataParallel 将模型分布到多个 GPU 上
model = nn.DataParallel(model)

1.2 DistributedDataParallel

DistributedDataParallel (DDP) 是 PyTorch 提供的一个用于分布式数据并行训练的模块，适用于单机多卡和多机多卡的场景。相比于 DataParallel，DDP 更加高效和灵活，能够在多个 GPU 和多个节点上进行并行训练。
DistributedDataParallel是多进程的，可以工作在单机或多机器中。DataParallel通常会慢于DistributedDataParallel。所以目前主流的方法是DistributedDataParallel。
封装示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def main(rank, world_size):
    # 初始化进程组
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    
    # 创建模型并移动到GPU
    model = SimpleModel().to(rank)
    
    # 包装模型为DDP模型
    ddp_model = DDP(model, device_ids=[rank])
    

if __name__ == "__main__":
    import os
    import torch.multiprocessing as mp

    # 世界大小：总共的进程数
    world_size = 4
    
    # 使用mp.spawn启动多个进程
    mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)

2. 多GPU训练的三种架构组织方式

由于上一节的讨论，本节所有源码均由DDP封装实现。
###2.1 数据拆分，模型不拆分（Data Parallelism）
数据并行（Data Parallelism）将输入数据拆分成多个子部分（mini-batches），并分配给不同的 GPU 进行计算。每个 GPU 上都有一份完整的模型副本。这种方式适用于模型相对较小，但需要处理大量数据的场景。
由于下面的代码涉及了rank、world_size等概念，这里先做一下简要普及。
** Rank**
rank 是一个整数，用于标识当前进程在整个分布式训练中的身份。每个进程都有一个唯一的 rank。rank 的范围是 0 到 world_size - 1。

• 用于区分不同的进程。
• 可以根据 rank 来分配不同的数据和模型部分。

World Size
world_size 是一个整数，表示参与分布式训练的所有进程的总数。

• 确定分布式训练中所有进程的数量。
• 用于初始化通信组，确保所有进程能够正确地进行通信和同步。

Backend
backend 指定了用于进程间通信的后端库。常用的后端有 nccl（适用于 GPU）、gloo（适用于 CPU 和 GPU）和 mpi（适用于多种设备）。

• 决定了进程间通信的具体实现方式。
• 影响训练的效率和性能。

Init Method
init_method 指定了初始化分布式环境的方法。常用的初始化方法有 TCP、共享文件系统和环境变量。

• 用于设置进程间通信的初始化方式，确保所有进程能够正确加入到分布式训练中。

Local Rank
local_rank 是每个进程在其所在节点（机器）上的本地标识。不同节点上的进程可能会有相同的 local_rank。

• 用于将每个进程绑定到特定的 GPU 或 CPU。
  

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

def train(rank, world_size):
    dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:29500', rank=rank, world_size=world_size)
    
    model = SimpleModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    criterion = nn.MSELoss().to(rank)
    optimizer = optim.SGD(ddp_model