使用 torch.compile 和梯度累积更快地训练模型

📢 转载信息

原文链接：https://machinelearningmastery.com/train-a-model-faster-with-torch-compile-and-gradient-accumulation/

原文作者：Adrian Tam

使用深度Transformer架构训练语言模型非常耗时。然而，你可以使用一些技巧来加速训练。在本文中，你将了解到：

• 使用 torch.compile() 来加速模型
• 使用梯度累积来训练具有更大有效批次大小的模型

让我们开始吧！

使用 torch.compile 和梯度累积更快地训练模型
图片作者：François Genon。保留部分权利。

概述

本文分为两个部分：

• 使用 torch.compile()
• 梯度累积

使用 torch.compile

当你编写模型代码并使用PyTorch运行时，代码以eager mode（即时模式）执行。这意味着代码逐行执行，结果存储在内存中。由于Python是一种解释型语言，这是其原生特性。你知道这是真的，因为如果代码中出现错误，直到执行到那一行代码时你才会看到错误。

以即时模式运行模型速度很慢。从PyTorch 2.0开始，你可以使用 torch.compile() 来编译模型以提高性能。它会生成一个新的、经过优化的模型对象。它不是你使用 nn.Module 创建的同一个模型对象，但它与原始模型共享相同的张量。你可以像平常一样使用这个编译后的模型进行前向传播、反向传播和优化器更新。

构建模型并将其编译成计算图是TensorFlow 1.0的预期工作方式。这使得调试变得更加困难，因为你执行的模型无法与你编写的代码逐行对应。因此，你应该在试运行并确认代码没有错误后再编译模型。

并非所有模型都可以被编译。但是，如果你的模型支持编译，你将立即受益于速度提升。要编译模型，你只需要在准备使用它之前替换模型对象即可：

... model = LlamaForPretraining(model_config).to(device) model.load_state_dict(checkpoint) model = torch.compile(model) ...

不要在编译后加载模型权重。这是因为编译后的模型是一个共享原始模型相同权重的对象。在编译期间，计算图的构建引用了原始模型的权重张量。如果在编译后加载权重，模型可能无法按预期工作。

类似地，要保存编译后的模型，你应该引用原始模型的 state_dict，如下所示：

torch.save(getattr(model, "_orig_mod", model).state_dict(), "model.pth")

可以通过 model._orig_mod 从编译后的模型中访问原始模型。在上面的代码中，我们使用 getattr(model, "_orig_mod", model) 来获取原始模型（如果存在），否则使用 model 本身。此代码适用于编译后的模型和原始模型。

梯度累积

当你训练模型时，你可能会花费两到三倍于前向传播的时间在反向传播上。这是因为反向传播在计算上更密集，并且使用更多的内存。

加速训练的一个简单技巧是减少反向传播的次数。这可以通过增加批次大小来实现：在相同数量的数据样本下，更大的批次大小意味着需要处理的批次更少。

然而，更大的批次大小需要更多的内存。在内存受限的环境中，你可以通过多次运行前向传播并累积梯度来模拟更大的批次大小。这就是所谓的梯度累积。

用代码解释这个想法更容易：

.. accumulate_steps = 4
for epoch in range(num_epochs):
    optimizer.zero_grad()
    for i, batch in enumerate(dataloader):
        # get batched data
        input_ids, target_ids = batch
        # create attention mask: causal mask + padding mask
        attn_mask = create_causal_mask(input_ids.shape[1], device) + \ 
                        create_padding_mask(input_ids, PAD_TOKEN_ID, device)
        # extract output from model
        logits = model(input_ids, attn_mask)
        # compute loss: cross-entropy between logits and target, ignoring padding tokens
        loss = loss_fn(logits.view(-1, logits.size(-1)), target_ids.view(-1))
        loss = loss / accumulate_steps
        # Run backward, but update only once every `accumulate_steps` steps
        loss.backward()
        if (i + 1) % accumulate_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()

上面的训练循环摘自上一篇文章，用于在本地GPU上训练Llama模型。

通常，当你运行前向传播时，你会计算损失。然后调用 loss.backward() 将损失梯度反向传播到模型参数。在PyTorch中，backward() 方法是累积的，这意味着梯度会被累加起来。因此，在运行反向传播之前，你需要显式调用 optimizer.zero_grad() 来清除梯度。

在上面的代码中，你故意不在每次迭代中调用 optimizer.zero_grad()。相反，你对除以 accumulate_steps 的损失进行反向传播。这样，梯度被按比例缩小，但在 accumulate_steps 次迭代中累积起来。每隔 accumulate_steps 次迭代，你运行优化器来调整模型参数。

这种方法产生的效果与使用更大的批次大小相当。但是，由于你进行的优化器更新次数更少，学习率调度器应相应调整。这意味着你需要使用不同的步数初始化调度器：

... num_training_steps = (len(dataloader) // accumulate_steps) * num_epochs
cosine_scheduler = lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=num_training_steps - num_warmup_steps,
    eta_min=0
)

进一步阅读

以下是一些你可能会感兴趣的资料：

• torch.compile 文档
• 来自PyTorch文档的自动混合精度示例

总结

在本文中，你了解到使用 torch.compile() 可以通过编译计算图来帮助加速模型。你还了解到，梯度累积是一种通过累积来自多个小批量（mini-batches）的梯度来使用更大有效批次大小进行训练的技术。通过这种方式减少优化器更新的次数，可以节省反向传播和参数更新的时间。

🚀 想要体验更好更全面的AI调用？

欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。