2-PyTorch 多 GPU 训练总结（1）

小轻贵

已于 2023-01-18 18:48:46 修改

阅读量676

点赞数 1

分类专栏：学习笔记文章标签： pytorch 深度学习人工智能

于 2023-01-01 20:23:24 首次发布

原文链接：https://www.zhihu.com/people/xiang-hong-xin-6

版权

学习笔记专栏收录该内容

2 篇文章 1 订阅

订阅专栏

前言

过去一直TITAN XP坚挺单卡训练，，但是训练时间太久，一直想研究一下如何实现多GPU训练，为此在这里做了一个总结。来帮助自己迅速了解 PyTorch 多 GPU 的基本用法。并且比较不同的模型训练版本，使得我能够选择出合适的实现方式和简易的修改方式来完成自己的项目。

本篇文章浅入深地介绍了关于多 GPU 相关的知识和代码细节，其中包括需要修改的代码细节、运行命令及实验结果，以供读者动手实践并检查自己的操作是否正确。最后，附上翻阅别的大佬总结的 Github 仓库来展示所有教程中的代码。

一、单机单GPU

首先从基于 PyTorch 的单机器单 GPU 训练出发，探索和实现单机多 GPU、多机多 GPU 的训练方式，涉及的内容如下：

首先定义了最简单的单机单 GPU 训练管道。内容参考于 Pytorch 官方教程。在之后的教程中将逐步修改代码为不同的多 GPU 训练方式。

数据集

torchvision.datasets 模块中的 Dataset 对象定义了许多现实世界中的视觉数据集，例如CIFAR, COCO（完整数据集列表可以在 Datasets - Torchvision 0.13 documentation 中查看）。本篇博客中使用 FashionMNIST 数据集。每个 torchvision 中的数据集包含了两个参数：transform 和 target_transform，分别用于处理样本和标签。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# Download training data from open datasets.
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)
# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

数据加载器

在定义数据加载器时，只需将定义好的 Dataset 作为参数传递给 DataLoader即可。

batch_size = 64

# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

创建模型

定义一个简单的神经网络。

# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__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

# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

model = NeuralNetwork().to(device)
print(model)

定义优化器

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

单 GPU 训练

单 GPU 训练时的两个重点：

使用device = "cuda" if torch.cuda.is_available() else "cpu" 获取当前的实验设备，如果CUDA环境和Torch安装正确的话，torch.cuda.is_available() 将返回 True，表示支持 GPU 训练。如果是 False 则使用 CPU 训练；
在遍历 dataloader 时，使用 X, y = X.to(device), y.to(device) 将数据 X, y 复制到 GPU 设备上去计算。

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

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")

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

保存模型

训练完成后，保存模型。

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

二、单机多GPU（Data Parallel (DP)）

总的来说，DP（DataParallel）不是效率最高的（DistributedDataParalle 更高效），但一定是修改代码最少的，需要修改的代码片段仅 3 处。

修改1. 获取机器上的所有 GPU 设备。

# [*] Get multiple GPU device for training.
n_gpu = torch.cuda.device_count()
device = torch.device('cuda:0' if n_gpu > 0 else 'cpu')
device_ids = list(range(n_gpu))

修改2. 将模型放入多个 GPU 中

声明 DataParallel 后，模型将被拷贝到 n 个 GPU 上。同时，在前向传播时，一个batch_size 被平均拆分成 n 份并输入到各个 GPU 的模型中。在前向传播后，每个 GPU 上的输出结果都会放到默认的 GPU (默认一般为 0 号GPU) 中，并在这个 GPU 上计算损失和反向传播（详情参见官方运行示例：DATA PARALLELISM）。由于需要把每个 GPU 上的结果放入到默认 GPU 上，会存在效率瓶颈的问题，例如 GPU 显存占用不平衡、GPU 利用率不平衡、以及通信损耗等。

if len(device_ids) > 1:
    model = torch.nn.DataParallel(model, device_ids=device_ids)

修改3. 模型保存

在使用多GPU训练后，模型其实是 torch.nn.DataParallel 类型。为了保存模型，我们需要使用 model.module 来获取该模型。

# [*] save model with multi-GPU
if isinstance(model, torch.nn.DataParallel):
    model_state_dict = model.module.state_dict()
else:
    model_state_dict = model.state_dict()
torch.save(model_state_dict, "model.pth")

三、多服务器多GPU

使用多个服务器的多个 GPU 训练模型，所涉及到的技术为 Distributed Data Parallel (DDP)。本节主要讲解了 DDP 的基本概念和用例。

官方定义

（DistributedDataParallel，DDP）在模块级实现数据并行，可以跨多台机器运行。使用 DDP 的应用程序应该生成多个进程，并为每个进程创建一个DDP实例。DDP在 torch.distributed中使用集体通信来同步梯度和缓冲区。更具体地说，DDP 为 model.parameters() 给定的每个参数注册一个 autograd 钩子，当在反向传递中计算相应的梯度时，该钩子将触发。然后 DDP 使用该信号触发跨进程的梯度同步。详情请参阅 DDP设计说明。使用 DDP 的推荐方式是为每个模型副本生成一个进程，其中模型副本可以跨越多个设备。 DDP 进程可以放在同一台机器上，也可以跨机器放置，但 GPU 设备不能跨进程共享。

本节从一个基础的 DDP 用例开始，然后演示更高级的用例，包括处理模型的检查点以及将 DDP 与模型并行相结合。

DataParallel 和 DistributedDataParallel 的区别

在进行讲解之前，先对比一下 DataParallel 和 DistributedDataParallel 的区别。在这里，我们将理解为什么 DistributedDataParallel 增加了代码的复杂性，但我们还是更容易考虑使用它而不是 DataParallel：

首先，DataParallel 是单进程、多线程的，只在一台机器上工作，而 DistributedDataParallel 是多进程的，适用于单机和多机训练。由于线程间 GIL（Global Interpreter Lock，全局解释锁）的争用、每次迭代复制模型以及分散输入和收集输出带来的额外开销，即使在单机上，DataParallel 通常也比 DistributedDataParallel 慢；
回想一下之前的教程，如果您的模型太大，无法安装在单个GPU上，则必须使用 model parallel 将其拆分到多个 GPU 上。DistributedDataParallel 支持模型并行，而 DataParallel 此时不支持。当 DDP 与模型并行相结合时，每个 DDP 进程将使用模型并行，所有进程将共同使用数据并行。

如果您的模型需要跨越多台机器，或者代码不适合数据并行范式时，参阅 RPC API 以获得更通用的分布式训练支持。

基础DDP用例

要创建 DDP 模块，必须首先正确地设置进程组。更多细节可以在用PyTorch写分布式应用程序中找到。

import os
import sys
import tempfile
import torch
import torch.nn as nn
import torch.optim as optim

# 导入分布式数据并行所需要的包
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

"""
# On Windows platform, the torch.distributed package only supports Gloo backend, FileStore and TcpStore. 
# For FileStore, set init_method parameter in init_process_group to a local file. Example as follow:
init_method="file:///f:/libtmp/some_file"
dist.init_process_group(
   "gloo",
   rank=rank,
   init_method=init_method,
   world_size=world_size)
For TcpStore, same way as on Linux.
"""

# 初始化分布式进程组
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # 如果操作系统是 Windows 或 macOS, 使用 gloo 代替 nccl
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)

# 摧毁进程组
def cleanup():
    dist.destroy_process_group()

在这里，我们使用 PyTorch 官方提供的模型举例说明，并在之后将我们的代码修改为 DDP 版本的代码。在官方的例子中，他们创建了一个玩具模型并尝试用 DDP 封装并向其提供一些输入数据。请注意，由于 DDP 将模型状态从 rank 0 进程广播到 DDP 构造函数中的所有其他进程，因此不需要担心不同 DDP 进程中的模型有着不同的初始化参数值。

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    # [*]
    setup(rank, world_size)

    # [*] create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)  # [*]
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup() # [*]

def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)

如上所见，DDP 封装了较低级别的分布式通信细节，并提供了一个干净的 API，就好像它是一个本地模型一样。梯度同步通信发生在反向传递期间并与反向计算重叠。当 backward() 返回时，param.grad 已经包含了同步的梯度张量。对于基础用例，DDP 只需要更多的 LoC 来设置进程组。在将 DDP 应用于更高级的用例时，需要注意一些警告。

处理速度不一致产生的问题

在 DDP 中，构造函数、前向传播和反向传播是在分布式的同步点上执行。不同的进程预计会启动相同数量的同步，并以相同的顺序到达这些同步点，并在大致相同的时间进入每个同步点。否则，快速进程可能会提前到达并在等待落后者时超时。因此，用户需要平衡跨进程间的工作负载分布。有时，由于网络延迟、资源争用或不可预测的工作负载峰值等原因，处理速度的偏差是不可避免的。为避免在这些情况下超时，请确保在调用 init_process_group 时传递足够大的超时值。

保存和加载检查点

在训练期间使用 torch.save 和 torch.load 来保存检查点以及从检查点恢复是很常见的。有关详细信息请参阅保存和加载模型。使用 DDP 时，一种优化手段是仅在一个进程中保存模型，然后将其加载到所有进程中，从而减少写入开销。这是正确的，因为所有进程都从相同的参数开始，并且梯度在反向传递中是同步的，因此优化器应该将参数设置为相同的值。如果您使用此优化方式，请确保在检查点保存完成之前没有进程开始加载这个检查点。此外，在加载模块时，需要提供适当的 map_location 参数以防止进程进入其他设备。如果缺少 map_location，torch.load 首先将模块加载到 CPU，然后将每个参数复制到它所保存的位置（可能是默认的 GPU，一般 GPU index 为 0），这有可能导致同一台机器上的所有进程使用同一组设备。有关更高级的故障恢复和弹性支持，请参阅 TorchElastic。

def demo_checkpoint(rank, world_size):
    # [*]
    print(f"Running DDP checkpoint example on rank {rank}.")
    setup(rank, world_size)

    # [*]
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    CHECKPOINT_PATH = tempfile.gettempdir() + "/model.checkpoint"
    # [*]
    if rank == 0:
        # All processes should see same parameters as they all start from same
        # random parameters and gradients are synchronized in backward passes.
        # Therefore, saving it in one process is sufficient.
        torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)

    # [*]
    # Use a barrier() to make sure that process 1 loads the model after process 0 saves it.
    dist.barrier()
    # configure map_location properly
    map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
    ddp_model.load_state_dict(
        torch.load(CHECKPOINT_PATH, map_location=map_location))

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank) # [*]

    loss_fn(outputs, labels).backward()
    optimizer.step()

    # [*]
    # Not necessary to use a dist.barrier() to guard the file deletion below
    # as the AllReduce ops in the backward pass of DDP already served as
    # a synchronization.
    if rank == 0:
        os.remove(CHECKPOINT_PATH)

    cleanup()

DDP 与模型并行的结合

DDP也适用于多 GPU 模型。包装多 GPU 模型的 DDP 在训练具有大数据量的大容量模型时特别有用。

class ToyMpModel(nn.Module):
    def __init__(self, dev0, dev1):
        super(ToyMpModel, self).__init__()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)  # [*]
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)  # [*]

    def forward(self, x):
        x = x.to(self.dev0)  # [*]
        x = self.relu(self.net1(x))
        x = x.to(self.dev1)  # [*]
        return self.net2(x)

将多 GPU 模型传递给 DDP 时，不得设置 device_ids 和 output_device。输入和输出数据将通过应用程序或模型的 forward() 方法放置在适当的设备中。

def demo_model_parallel(rank, world_size):
    print(f"Running DDP with model parallel example on rank {rank}.")
    setup(rank, world_size)

    # setup mp_model and devices for this process
    dev0 = (rank * 2) % world_size
    dev1 = (rank * 2 + 1) % world_size
    mp_model = ToyMpModel(dev0, dev1)
    ddp_mp_model = DDP(mp_model)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    # outputs will be on dev1
    outputs = ddp_mp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(dev1)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus
    run_demo(demo_basic, world_size)
    run_demo(demo_checkpoint, world_size)
    run_demo(demo_model_parallel, world_size)

使用 torch.distributed.run/torchrun 初始化 DDP

我们可以利用 PyTorch Elastic 简化 DDP 代码并更容易地初始化程序。这里，我们仍然使用 Toymodel 的例子并创建一个名为 elastic_ddp.py 的文件。

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    print(f"Start running basic DDP example on rank {rank}.")

    # create model and move it to GPU with id rank
    device_id = rank % torch.cuda.device_count()
    model = ToyModel().to(device_id)
    ddp_model = DDP(model, device_ids=[device_id])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_id)
    loss_fn(outputs, labels).backward()
    optimizer.step()

if __name__ == "__main__":
    demo_basic()

我们可以运行如下命令在所有节点（节点可以理解为服务器）上去初始化上面创建的 DDP 作业：

torchrun --nnodes=2 --nproc_per_node=8 --rdzv_id=100 
    --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400 elastic_ddp.py

我们在两台主机上运行 DDP 脚本，每台主机运行 8 个进程（ 1 个进程使用 1 个GPU），也就是我们在 16 个 GPU 上运行它。请注意，所有节点的 $MASTER_ADDR 必须相同。

在这里，torchrun 将启动 8 个进程并在启动它的节点上的每个进程上调用 elastic_ddp.py，但用户还需要应用 slurm 等集群管理工具（没有的话使用 torch.distributed.launch 或 torch.multiprocessing 也可以）才能在 2 个节点上实际运行此命令。

例如，在启用 SLURM 的集群上，我们可以编写一个脚本来运行上面的命令并将 MASTER_ADDR 设置为：

export MASTER_ADDR=$(scontrol show hostname ${SLURM_NODELIST} | head -n 1)

然后我们可以使用 SLURM 命令运行这个脚本：srun --nodes=2 ./torchrun_script.sh。当然，这只是一个例子；您可以选择自己的集群调度工具来启动 torchrun 作业。

后续

后续内容见链接

代码

完整的代码可以查看
pytorch多GPU总结

在此对闪闪红星闪闪同志的博客表达最真挚的感谢

小轻贵

关注

1
点赞
踩
3

收藏

觉得还不错? 一键收藏
0
评论
2-PyTorch 多 GPU 训练总结（1）

PyTorch多GPU训练总结
复制链接

扫一扫

专栏目录