使用深层转换器架构训练语言模型非常耗时。但是，您可以使用一些技术来加速训练。在本文中，您将了解：
- 使用
torch.compile()加速模型 - 使用梯度累积训练具有更大有效批量大小的模型
让我们开始吧！
使用 torch.compile 和梯度累积更快地训练模型
摄影： 弗朗索瓦·热农。保留一些权利。
概述
本文分为两部分；他们是：
- 使用
torch.compile() - 梯度累积
使用 torch.compile
当您编写模型代码并在 PyTorch 中运行它时，它会以 eager 模式执行。这意味着代码是逐行执行的，结果存储在内存中。这是 Python 的本机特性，因为它是一种解释性语言。您知道是这种情况，因为当您在代码中犯错时，只有运行该行代码后您才会看到该错误。
在 Eager 模式下运行模型速度很慢。从 PyTorch 2.0 开始，您可以使用 torch.compile() 编译模型以提高性能。这会生成一个经过优化的新模型对象。它与您使用创建的模型对象不同 nn.Module，但它与原始模型共享相同的张量。您可以像往常一样使用此编译模型进行前向传递、后向传递和优化器更新。
构建模型并将其编译为计算图就是 TensorFlow 1.0 的工作方式。这使得调试变得更加困难，因为您执行的模型无法与您编写的代码逐行匹配。因此，在运行试验并确认模型没有错误之前，不应编译模型。
并非所有模型都可以编译。但是，如果您的模型支持编译，您将立即受益于加速。要编译模型，您需要做的就是在准备使用模型对象之前替换它：
|
。。。 模型 = 预训练骆驼（模型配置）。到（设备） 模型。加载状态字典（检查站） 模型 = 火炬。编译（模型） 。。。 |
编译后不要加载模型权重。这是因为编译后的模型是一个与原始模型共享相同权重的对象。在编译过程中，计算图是参考原始模型的权重张量构建的。如果在编译后加载权重，模型可能无法按预期工作。
同样，要保存编译后的模型，您应该引用原始模型的状态字典，如下所示：
|
火炬。节省（获取属性（模型, “_original_mod”, 模型）。状态字典（）, “型号.pth”） |
可以使用编译后的模型访问原始模型 model._orig_mod。在上面的代码中，我们使用 getattr(model, "_orig_mod", model) 获取原始模型（如果存在），或使用 model 本身，如果没有。这行代码适用于编译模型和原始模型。
梯度累积
当您训练模型时，您在向后传递上花费的时间可能比向前传递多两到三倍。这是因为向后传递的计算量更大并且使用更多的内存。
加快训练速度的一个简单技巧是减少向后传递。这可以通过增加批大小来实现：在相同数量的数据样本下，更大的批大小意味着要处理的批更少。
但是，较大的批处理大小需要更多的内存。在内存受限的环境中，您可以通过运行多个前向传递并累积梯度来模拟更大的批量大小。这就是所谓的 梯度累积。
用代码来解释这个想法更容易：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 号 18 19 20 21 22 号 |
。。 累积步数 = 4 为了 时代 在 范围（纪元数）: 优化器。零梯度（） 为了 我, 批 在 枚举（数据加载器）: # 获取批量数据 输入ID, 目标ID = 批 # 创建注意力掩码：因果掩码+填充掩码 属性掩码 = 创建因果掩码（输入ID。形状[1], 设备） + \ 创建填充掩码（输入ID, PAD_TOKEN_ID, 设备） # 从模型中提取输出 逻辑数 = 模型（输入ID, 属性掩码） # 计算损失：logits 和目标之间的交叉熵，忽略填充标记 损失 = 损失函数（逻辑数。看法（–1, 逻辑数。尺寸（–1））, 目标ID。看法（–1）） 损失 = 损失 / 积累_步骤 # 向后运行，但每个 `accumulate_steps` 步骤只更新一次 损失。落后（） 如果 （我 + 1） % 累积步数 == 0: 火炬。恩。实用程序。剪辑_梯度_范数_（模型。参数（）, 1.0） 优化器。步（） 优化器。零梯度（） 调度程序。步（） |
上面的训练循环摘录自 上一篇文章 用于在本地 GPU 上训练 Llama 模型。
通常，当您进行前向传球时，您会计算损失。然后你打电话 loss.backward() 通过模型参数反向传播损失梯度。在 PyTorch 中， backward() 方法是累积的，意味着梯度相加。因此，您需要调用 optimizer.zero_grad() 在运行向后传递之前显式清除梯度。
在上面的代码中，你故意不调用 optimizer.zero_grad() 在每次迭代中。相反，您对损失进行反向传播除以 accumulate_steps。这样，梯度会按比例缩小，但会累积 accumulate_steps 迭代。每一次 accumulate_steps 迭代时，您运行优化器来调整模型参数。
这种方法产生的结果与使用较大批量大小获得的结果相当。但是，由于您运行的优化器更新较少，因此应相应调整学习率计划。这意味着您需要使用不同数量的步骤来初始化调度程序：
|
。。。 训练步数 = （伦（数据加载器） // 累积步数) * num_epochs 余弦调度器 = lr_调度程序。余弦退火LR（ 优化器, 最大温度=训练步数 – 预热步骤数, 和最小=0 ） |
进一步阅读
以下是一些您可能感兴趣的材料：
概括
在本文中，您了解到使用 torch.compile() 可以通过编译计算图来帮助您加速模型。您还了解到，梯度累积是一种通过累积多个小批量的梯度来进行更大有效批量大小训练的技术。由于通过这种方式运行的优化器更新较少，因此可以节省向后传递和参数更新的时间。
