笔记本电脑 台式电脑

申请回电

支持

利用 GPU 训练大型模型

2021 年春
中国 Yuanhao Wu

高级算法工程师

作为惠普 Z 系列全球数据科学大使，Yuanhao Wu 的内容由惠普赞助，且向他提供了惠普产品。

 

在上一篇博文中，我分享了自己参加 Kaggle Jigsaw 多语言恶毒评论分类竞赛的故事。当时我只有一张具有 11 GB VRAM 的 1080Ti 显卡，用它来训练 SOTA Roberta-XLM 大型模型无异于痴人说梦，这类模型需要更大的 VRAM。在这篇博文中，我想分享一些关于如何降低 VRAM 使用量的技巧，帮助大家使用 GPU 训练更大的深度神经网络。

 

第一个技巧是使用混合精度训练。在训练模型时，每项参数通常都会存储在 VRAM 中。我们都知道，VRAM 的总使用量等于存储的参数数量乘以单个参数的 VRAM 使用量。模型越大不仅意味着性能越好，使用的 VRAM 也会越多。由于性能相当重要，比如在 Kaggle 竞赛中，我们不希望减小模型的规模。因此减少内存使用量的唯一方法是减少每个变量的内存使用量。默认情况下，变量采用 32 位浮点格式，也就是说一个变量占用 4 个字节。幸运的是，人们发现某些变量可以采用 16 位浮点格式，而不会损失太多的精度。这意味着我们可以将内存使用量减半！此外，使用低精度还可以提高训练速度，特别是在支持 Tensor Core 的 GPU 上。

 

在 1.5 版本之后，PyTorch 开始支持自动混合精度训练。该框架可以识别需要全精度的模块，并对其使用 32 位浮点格式，而对其他模块则使用 16 位浮点格式。下面是 PyTorch 官方文档中的一段代码示例。

第二个技巧是使用梯度累积。梯度累积的原理很简单：在优化器更新参数之前，用相同的模型参数进行几次反向传播。每次反向传播所计算的梯度将被累积（加总）。如果实际的批量大小为 N，累积的梯度有 M 步，则等效批量大小为 N*M。然而，训练结果并不会严格相等，因为有些参数（如批标准化统计数据）无法很好地累积。

 

关于梯度累积，有几点需要注意：

 

  1. 在混合精度训练中使用梯度累积时，应针对有效批次校准其比例，且应该以有效批次的粒度进行比例更新。
  2. 在分布式数据并行 (DDP) 训练中使用梯度累积时，可以使用 no_sync() 上下文管理器来禁用前 M-1 步的梯度全还原。

 

最后一个非常重要的技巧是使用梯度检查点。梯度检查点的基本思路是只将一些节点的中间结果保存为检查点，然后在反向传播过程中对这些节点之间的其他部分进行重新计算。《梯度检查点》的作者声称，通过使用梯度检查点，他们可以把 10 倍大的模型放到 GPU 上训练，而计算时间只增加了 20%。PyTorch 从 0.4.0 版本开始正式支持这一功能，一些常用的库（如 Huggingface Transformers）也支持这一功能。

 

以下示例展示了如何利用梯度检查点来训练 Transformer 模型：

在这篇博文的最后，我想分享一个之前我在 HP Z4 工作站上做过的简单基准测试。正如我之前提到的，这台工作站配备了 2 个 24 GB VRAM RTX6000 GPU，我在实验中只使用了一个 GPU。我利用不同的配置来训练 XLM-Roberta Base/Large，并使用了 Kaggle Jigsaw 多语言有毒评论分类竞赛中使用的训练集。

大家可以看到，混合精度训练不仅减少了内存使用量，而且还带来了显著的速度提升。梯度检查点的功能也非常强大。它将 VRAM 的使用量从 23.5 GB 削减到了 11.8 GB！我真希望我能够早点发现这项技术，这样我就可以用我之前的 11G 2080Ti GPU 训练 XLM-Roberta Large 模型了 :)

有疑问？
联系销售支持。 

不知道该怎么选
Z 系列工作站？

申请回电

 需要 Z 系列工作站方面的支持？ 

前往支持页面

免责声明

  1. 产品的Logo及位置可能与图片有所差异，不影响产品性能和功能，请以实物为准。

     

    产品图片仅供参考，因不同国家可能存在稍许差异，实际产品以销售为准。

     

    本文所载信息如有变更，恕不另行通知。惠普产品与服务的完整保修条款见此类产品和服务附带的正式保修声明。本文中的信息概不构成额外的保修条款。惠普对本文包含的技术或编辑方面的错误或遗漏概不负责。