多GPU模型训练的监测和优化教程

方石剑
1,007 阅读16分钟

你是否为在多个GPU上监控和优化深度神经网络的训练而苦恼？如果是的话，你就来对地方了。

在这篇文章中，我们将讨论Pytorch Lightning的多GPU训练，并找出应该采用的最佳做法来优化训练过程。我们还将看到如何在训练过程中监控所有GPU的使用情况。

让我们从一些基础知识开始。

什么是使用GPU的分布式训练？

有时，对于复杂的任务，如计算机视觉或自然语言处理，训练深度神经网络（DNN）涉及解决数百万或数十亿参数的梯度下降，因此它成为一个复杂的计算过程，可能需要几天甚至几周才能完成。

例如，Generative Pre-trained Transformer 3（GPT-3）是一个自回归语言模型，它有1750亿个参数，在一个NVIDIA Tesla V100 GPU上训练它大约需要355年。但如果使用1024个英伟达A100 GPU并行训练同一模型，我们可以估计训练时间约为34天。因此，在GPU上进行并行训练是目前广泛使用的加速过程的方法。

为了清楚地了解我们如何使用多个GPU来训练深度神经网络，让我们简单地看看神经网络是如何训练的。

  • 深度神经网络模型通常使用迷你批次梯度下降法进行训练，其中训练数据被随机抽样为迷你批次。
  • 小型批次被送入模型，分两个阶段遍历模型。
    • 前传
    • 后向通道
  • 前向传递产生预测，并计算预测和地面真相之间的损失。
  • 后向通道将误差通过网络的各层（称为反向传播），以获得梯度来更新模型权重。
  • 通过前向和后向阶段的小型批次被称为迭代，而一个历时被定义为通过整个训练数据集进行前向-后向传递。
  • 训练过程会持续多个 epochs，直到模型收敛。

如果这对你来说似乎不堪重负，我建议你在这篇文章中更深入地了解神经网络是如何训练的。

为了加快训练过程，我们使用多个GPU来并行化训练过程，数据并行和模型并行是用于并行化任务的两种技术。

数据并行

在数据并行中，每个GPU持有一个模型的副本，数据被分为n个分区，每个分区都用于在每个GPU上训练一个模型的副本。

当异步数据并行被应用时，参数服务器负责权重的更新。每个GPU将其梯度发送到参数服务器，然后参数服务器更新权重并将更新的权重送回给该GPU。

通过这种方式，GPU之间不存在同步问题。这种方法解决了分布式计算环境中不稳定的网络问题，但它引入了不一致的问题。此外，这种方法并没有减少GPU之间的数据传输数量。

模型并行化

模型并行将一个模型划分给多个GPU，每个GPU负责模型的指定层的权重更新。中间数据，如前向传递的神经网络层的输出和后向传递的梯度，在GPU之间传输。

由于这些分区有依赖性，在模型并行的天真实施中，每次只有一个GPU处于活动状态，导致GPU利用率低。为了实现并行执行，流水线并行将输入的微型批处理分成多个微型批处理，并在多个GPU上对这些微型批处理进行流水线执行。这在下图中有所概述。

Pipeline parallelism

管线并行｜来源

上图表示一个有4层的模型放在4个不同的GPU上（纵轴）。横轴表示通过时间来训练这个模型，表明GPU的利用效率要高得多。然而，仍然存在一个泡沫（如图所示），即某些GPU没有被利用。

使用PyTorch Lightning的多GPU训练

在本节中，我们将重点讨论如何使用PyTorch Lightning在多个GPU上进行训练，因为它在过去一年中越来越受欢迎。PyTorch Lightning真的很简单，使用起来也很方便，它可以帮助我们扩展模型，不需要模板。锅炉模板代码是大多数人在扩展模型时容易出错的地方。

分布式模式

在本节中，我们将介绍Pytorch Lightning提供的不同分布式模式。

数据并行

我们可以在一台拥有多个GPU的机器上训练一个模型。通过DataParallel（DP）方法，一个批次被平均分配给一个节点的所有选定的GPU，然后根节点汇总所有结果。但是Pytorch闪电开发人员并不推荐这种方法，因为它还不稳定，如果你在forward()或*_step()方法中给模块分配一个状态，你可能会看到错误或错误行为。

# train on 4 GPUs (using DP mode)

分布式数据并行（Distributed Data Parallel

分布式数据并行（DDP）的工作方式如下。

  • 每个节点上的每个GPU都有自己的进程。
  • 每个GPU可以看到整个数据集的一个子集。它只能看到那个子集。
  • 每个进程初始化模型。
  • 每个进程并行地执行一个完整的前向和后向通道。
  • 梯度在所有进程中被同步和平均化。
  • 每个进程更新其优化器。

我们有两种方法可以使用这种方法，即 "ddp "和 "ddp_spawn"。在'ddp'方法中，该脚本在引擎盖下被多次调用，并带有正确的环境变量。

# train on 8 GPUs (same machine (ie: node))

虽然在大多数情况下这似乎是一个不错的选择，但它有一些局限性，因为它不能在Jupyter Notebook、Google COLAB和Kaggle等工具中工作。另外，当有一个没有根包的嵌套脚本时，它似乎不工作。在这些情况下，'ddp_spawn'方法是首选。

ddp_spawn "和ddp完全一样，只是它使用torch.multiprocessing.spawn()方法来启动训练进程。因此，人们可能认为总是喜欢'ddp_spawn'方法而不是'ddp'，但是'ddp_spawn'也有这些限制。

  • spawn方法在子进程中训练模型，主进程中的模型不会被更新。
  • Dataloader(num_workers=N)，如果N很大的话，就会对DDP的训练造成瓶颈，也就是说，它将会非常慢或者根本无法工作。这是PyTorch的一个限制。
  • 这种方法迫使所有的东西都是可提取的。
# train on 8 GPUs (same machine (ie: node))

就速度和性能而言，"ddp "方法应该总是优于 "ddp_spawn"。

分布式数据并行2

DDP2在单台机器中的行为与DP类似，但当在多个节点上使用时，它的行为与DDP相同。有时，在同一台机器上使用所有的批次而不是一个子集可能是有用的，ddp2方法在这种情况下可能很方便。DDP2做了以下工作。

  • 将数据的一个子集复制到每个节点上。
  • 在每个节点上建立一个模型。
  • 使用DP运行一个前向和后向通道。
  • 在各节点上同步梯度。
  • 应用优化器的更新。
# train on 32 GPUs (4 nodes)

目前不建议使用这种技术，因为它在所有PyTorch版本>=1.9的情况下都是坏的，不清楚如何让它在PyTorch>=1.9的情况下工作，而且没有对这种方法进行功能测试。

Horovod

Horovod是一个分布式深度学习训练框架，适用于TensorFlow、Keras、PyTorch和Apache MXNet，它使分布式深度学习变得快速和易于使用。

  • 每个过程都使用一个GPU来处理一个固定的数据子集。
  • 在后向传递过程中，梯度在所有GPU上被平行平均化。
  • 在进入下一阶段之前，这些梯度在开始下一步骤之前被同步应用。
  • 在训练脚本中，Horovod将从环境中检测工人的数量，并自动调整学习率，以补偿增加的总批次规模。

Horovod支持使用同一训练脚本的单GPU、多GPU和多节点训练。它可以在训练脚本中进行配置，以便在任何数量的GPU/进程中运行，具体如下。

# train Horovod on GPU (number of GPUs / machines provided on command-line)

# train Horovod on CPU (number of processes / machines provided on command-line)

当启动训练工作时，驱动应用程序将被用来指定工作进程的总数。

# run training with 4 GPUs on a single machine
horovodrun -np 4 python train.py

# run training with 8 GPUs on two machines (4 GPUs each)
horovodrun -np 8 -H hostname1:4,hostname2:4 python train.py

碎片化训练

在训练大型模型或尝试更大的批次规模时，你可能会遇到一些内存问题。人们可能会想到在这种情况下使用模型并行，但目前，由于与之相关的复杂实现，我们使用分片训练来代替。

在引擎盖下，分片训练与数据并行训练类似，只是优化器的状态和梯度在GPU之间分片。在内存受限的多GPU设置中，或者在训练大型模型（500M以上参数的模型）时，强烈推荐使用这种方法。

要使用分片训练，你需要首先使用下面的命令安装FairScale。

pip install fairscale
# train using Sharded DDP
trainer = Trainer(strategy="ddp_sharded")

在使用分片训练策略时，在内存和性能之间有一个权衡，因为由于设备之间的高分布式通信，训练可能会变得更慢。

如何在多个GPU上优化训练

当涉及到在多个GPU上用大型数据集训练巨大的模型时，我们可能会遇到一些内存或性能瓶颈的问题。在本节中，我们将看到我们如何优化训练过程。

FairScale Activation Checkpointing

激活检查点在前向传递过程中一旦不需要激活，就将其从内存中释放。然后，在需要的时候，它们会被重新计算到后向传递中。当你有中间层产生大量的激活时，激活检查点是非常有用的。

与PyTorch的实现不同，FairScales的检查点封装器还可以正确处理批量规范层，确保由于多次前向传递而正确跟踪统计。

这在训练大型模型时可以节省内存，但是需要包装你想使用的激活检查点的模块。你可以在这里阅读更多关于它的信息。

from fairscale.nn import checkpoint_wrapper

class Model(LightningModule):
    def init(self):
        super().__init__()
        self.block1 = checkpoint_wrapper(nn.Sequential(nn.Linear(32, 32),       nn.ReLU()))
        self.block_2 = nn.Linear(32, 2)

混合精度（16位）训练

默认情况下，PyTorch和大多数深度学习框架一样，使用32位浮点（FP32）运算。另一方面，许多深度学习模型可能会通过较低位的浮点（如16位）实现完全的准确性。由于它们需要较少的内存，因此有可能训练和部署大型神经网络，这导致了由于较少的内存带宽要求而增强数据传输操作。

但是要使用16位浮点，你必须有

  • 支持16位精度的GPU，如NVIDIA pascal架构或更新版本。
  • 你的优化算法，即training_step在数值上应该是稳定的。

混合精度结合了32位和16位浮点的使用，这提高了性能，同时也摆脱了我们可能面临的任何内存问题。Lightning为使用本地或APEX amp后端的GPU提供混合精度训练。

Trainer(precision=16, amp_backend="native")
# To use NVIDIA APEX amp_backend
Trainer(amp_backend="apex", precision=16)

除非你需要更精细的控制，否则建议始终使用本地amp_backend。

优先选择分布式数据并行（DDP）而不是数据并行（DP）。

正如上一节已经提到的，我们应该优先使用DDP策略而不是DP。这背后的原因是，DP对每个批次使用3个传输步骤，而DDP只使用2个传输步骤，从而使其更快。

DP执行以下步骤：

  1. 将模型复制到设备上。
  2. 将数据复制到设备上。
  3. 将每个设备的输出复制到主设备上。

GPU data parallel strategy

GPU数据并行策略｜来源

DDP执行以下步骤。

  1. 将数据移动到设备上。
  2. 传输和同步梯度。

GPU distributed data parallel strategy

GPU分布式数据并行策略 |来源

在改变技术和策略的同时，我们也可以在代码中做一些改变，以优化我们的训练过程，以下是其中的几个例子。

增加批量大小

如果你在训练时使用的是小批量，那么更多的时间会花在加载和卸载训练数据上，而不是计算操作上，这会降低训练速度。因此，最好是使用更大的批处理量来提高GPU的利用率。但增加批处理量可能会对模型的准确性产生不利影响，因此我们应该尝试不同的批处理量，以找到一个最佳的批处理量。

使用PyTorch的DataLoader方法更快地加载数据

Pytorch中的DataLoader类是一种快速和简单的方法来加载和批处理你的数据。我们可以使用参数 "num_workers"，通过设置其值为多于1来更快地加载数据进行训练。当使用PyTorch闪电时，它会为你推荐num_workers的最佳值。

但是如果你的数据集非常大，你可能会遇到内存问题，因为加载器的工作进程和父进程将为父进程中所有从工作进程中访问的Python对象消耗相同的CPU内存。避免这个问题的一个方法是在Dataloader __getitem__方法中使用Pandas、Numpy或PyArrow对象而不是Python对象。

将梯度设置为零

我们可以通过覆盖optimizer_zero_grad()方法并将其设置为None来提高性能和速度，而不是将梯度设置为零，一般来说，这将有更低的内存占用。

class Model(LightningModule):
    def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):
        optimizer.zero_grad(set_to_none=True)

模型切换

当我们必须在分布式环境下用多个优化器进行梯度累积时，这个方法特别有用。当执行梯度累积时，将sync_grad设置为False会阻止这种同步，并提高你的训练速度。

LightningOptimizer为高级用户提供了一个toggle_model()函数作为contextlib.contextmanager()。以下是官方PyTorch Lightning文档中的一个例子。

# Scenario for a GAN with gradient accumulation every two batches and optimized for multiple GPUs.
class SimpleGAN(LightningModule):
    def __init__(self):
        super().__init__()
        self.automatic_optimization = False

    def training_step(self, batch, batch_idx):
        # Implementation follows the PyTorch tutorial:
        # https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html
        g_opt, d_opt = self.optimizers()

        X, _ = batch
        X.requires_grad = True
        batch_size = X.shape[0]

        real_label = torch.ones((batch_size, 1), device=self.device)
        fake_label = torch.zeros((batch_size, 1), device=self.device)

        # Sync and clear gradients
        # at the end of accumulation or
        # at the end of an epoch.
        is_last_batch_to_accumulate = (batch_idx + 1) % 2 == 0 or self.trainer.is_last_batch

        g_X = self.sample_G(batch_size)

        ##########################
        # Optimize Discriminator #
        ##########################
        with d_opt.toggle_model(sync_grad=is_last_batch_to_accumulate):
            d_x = self.D(X)
            errD_real = self.criterion(d_x, real_label)

            d_z = self.D(g_X.detach())
            errD_fake = self.criterion(d_z, fake_label)

            errD = errD_real + errD_fake

            self.manual_backward(errD)
            if is_last_batch_to_accumulate:
                d_opt.step()
                d_opt.zero_grad()

        ######################
        # Optimize Generator #
        ######################
        with g_opt.toggle_model(sync_grad=is_last_batch_to_accumulate):
            d_z = self.D(g_X)
            errG = self.criterion(d_z, real_label)

            self.manual_backward(errG)
            if is_last_batch_to_accumulate:
                g_opt.step()
                g_opt.zero_grad()

        self.log_dict({"g_loss": errG, "d_loss": errD}, prog_bar=True)

正如你在代码中所看到的，我们将sync_grad参数设置为False，并且只在一个纪元结束时或每两个批次后将其设置为True。通过这样做，我们实际上是在每两批或在历时结束后积累梯度。

避免调用.item(), .numpy(), .cpu()。

避免在代码中调用.item()、.numpy()、.cpu()。如果你必须删除连接图的调用，你可以使用.detach()方法代替。这是因为这些调用的每一次都会导致数据从GPU传输到CPU，会使性能急剧下降。

清除缓存

每次调用torch.cuda.empty_cache()方法时，所有的GPU都要等待同步。所以要避免不必要地调用这个方法。

如何监控多个GPU上的训练

在训练模型时监控GPU的使用情况是非常重要的，因为它可能会提供一些有用的见解来改善训练，如果GPU没有被充分利用，我们可以相应地处理它。有各种工具，如Neptune和Wandb，可以用来监测多个GPU的训练情况。

在本节中，我们将使用Neptune来监控GPU和GPU内存，同时通过多个GPU进行训练。

用Neptune监控训练情况

Neptune是一个元数据存储，可以在任何MLOps工作流程中使用。它允许我们在训练时监控我们的资源，所以我们可以用它来监控训练时使用的不同GPU的使用情况。

在PyTorch Lightning代码中加入Neptune真的很简单，你所要做的就是创建一个NeptuneLogger对象并将其传递给Trainer对象，如下图所示。

from pytorch_lightning import Trainer
from pytorch_lightning.loggers import NeptuneLogger

# create NeptuneLogger
neptune_logger = NeptuneLogger(
    api_key="ANONYMOUS",  # replace with your own
    project="common/pytorch-lightning-integration",  # "<WORKSPACE/PROJECT>"
    tags=["training", "resnet"],  # optional
)

# pass it to the Trainer
trainer = Trainer(max_epochs=10, logger=neptune_logger)

# run training
trainer.fit(my_model, my_dataloader)

如果你是第一次接触到Neptune，我强烈建议你通过这个步骤指南来安装所有必要的库，使其发挥作用。之后，请查看Neptune + PyTorch Lightning整合文档。

运行该文件后，你应该得到一个控制台的链接。你可以看到监控部分（下图中的包围），在那里你可以看到训练时所有GPU的使用情况以及其他一些指标。

GPU model training

监控多个GPU上的训练 |来源

让我们看看我们可以从GPU利用率图中推断出什么样的有意义的见解。

GPU model monitoring

监测多个GPU上的训练情况 |来源

  • 正如你在上面看到的，GPU的使用情况是波动的，有一些短暂的时间没有被使用，要解释其中的原因其实并不容易。
  • 这可能发生在验证阶段，因为我们在这个阶段不计算梯度，或者这可能是由于其他的瓶颈，例如，你可能在你的数据上使用一些数据预处理技术，使用CPU可能真的很慢。
  • 此外，在一些框架中，如Caffe，默认情况下，在验证阶段只使用一个GPU，所以你可能会发现在这种情况下只有一个GPU的使用率很高。
  • 因此，根据你训练神经网络的方式，你可能会发现一个不同的图表，表明不同的GPU是如何被利用的。

总结

这篇文章讨论了为什么我们用多个GPU来训练机器学习模型。我们还发现用Pytorch闪电在多个GPU上训练是多么容易，以及优化训练过程的最佳方法。最后，我们发现Neptune如何在训练时用于监控GPU的使用。

avatar
文章
阅读
粉丝