Torch 编译优化 Llama 3.2 可以事半功倍Torch 编译最初是在 PyTorch 2.0 中引入的, 但经

Preview image

Torch 编译优化 Llama 3.2 可以事半功倍

但这取决于你的 GPU

Torch 编译最初是在 PyTorch 2.0 中引入的, 但经过多次更新和优化后, 它才能可靠地支持大多数大型语言模型(LLM).

在推理方面, torch.compile可以真正加快解码速度, 而内存使用量却只有很小的增加.

在本文中, 将介绍 torch.compile 的工作原理, 并衡量它对 LLM 推理性能的影响. 要在代码中使用 torch.compile, 只需添加一行即可. 在本文中, 我使用 Llama 3.2 进行了测试, 并使用两个不同的 GPU, 尝试了bitsandbytes量化: Google Colab 的 L4 和 A100.

Torch 编译: 它是如何让模型更快的?

torch.compile提供了一种通过将标准 PyTorch 代码转换为优化的机器代码来加速模型的方法. 这种方法被称为 JIT(Just-In-Time, 即时编译), 它能让代码在特定硬件上更高效地运行, 也就是比普通 Python 代码更快. 对于复杂的模型, 这种方法尤其有效, 即使是微小的速度改进也能节省大量时间.

torch.compile背后的核心工具包括 PyTorch 编译器中的几个重要组件:

TorchDynamo: 该工具在执行过程中修改 Python 代码, 使 PyTorch 能够以图谱的形式捕获操作, 然后由torch.compile进行优化.
AOTAutograd: 该工具根据捕获的图谱生成梯度, 用于训练.
PrimTorch: 它将复杂的操作简化为较小的, 易于管理的组件.
TorchInductor: 该组件获取优化后的图谱, 并为 GPU, CPU 和其他硬件设置生成硬件优化代码.

torch.compile的一大优势是易于使用. 它只需用 torch.compile 对模型进行封装, 就能生成优化版本. 它能与现有的 PyTorch 代码顺利集成.

当你第一次使用 torch.compile 运行一个模型时, 它会执行初始优化, 随后的调用将受益于这个更快的版本. 在内部, torch.compile 遵循一个三步流程:

它将模型分解成更小的, 经过优化的片段, 而将任何不可优化的代码保留为原始形式.
调整 PyTorch 操作以适应特定的硬件后端.
最后, 它会对这些部分进行编译, 以在设备上最大限度地提高效率, 主要是通过最大限度地减少内存传输.

其他优化还包括将多个操作合并为单个内核调用, 以及使用 CUDA 图谱捕捉来提升 GPU 性能. 虽然并非模型的每个部分都能得到优化, 但 torch.compile通常能加速大多数模型, 而无需进行结构性修改.

不过, 也存在一些限制. 输入图谱不同的模型可能会触发重复的重新编译, 从而导致速度减慢. 一致的输入图谱有助于避免这个问题, 不过对于使用 LLM 的推理和微调来说, 这通常不是问题.

编译后的模型也可能比预期消耗更多内存或运行更慢, 因此建议进行基准测试, 以确认 torch.compile是否能提高特定设置的性能. 在分布式设置中, 优化可能不会总是统一应用, 因此最好在配置分布式进程之前编译模型.

使用 Transformer 编译 Torch 模型

要充分利用 torch.compile 加速, 我建议使用最新的 PyTorch 2.5 版本. 然后, 你可以通过将模型传递给torch.compile来轻松启用torch.compile, 如下所示:

import torch
model = torch.compile(model)

就是这么简单.

Hugging Face 发表了一些有趣的关于 torch.compile 对视觉模型影响的基准结果.

让我们看看它在 LLM 上的效果如何.

Torch Compile 用于 LLM

我在 Llama 3.2 3B 中试用了 torch.compile.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
checkpoint = "meta-llama/Llama-3.2-3B"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="cuda")
model = torch.compile(model)

在使用 torch 编译进行推理基准测试时, 我使用了 optimum benchmark (Apache 2.0 license), 批量大小为 1, 序列长度为 1024. 我们将特别关注预填充(KV 缓存创建, 编码)和解码(令牌生成)阶段的吞吐量(令牌/秒), 以及模型编译后所消耗的内存.

请记住, torch.compile在首次使用时需要一些初始时间来编译模型. 要获得准确的基准, 最好先运行几次热身迭代. 这样, 推理速度和内存使用量等性能指标只能在模型完全编译后才能测量.

我还测试了 torch.compile 对使用 bitsandbytes(BnB)量化为 4 位的 Llama 3.2 (3B) 的影响.

使用 A100 GPU(Google Colab)得出的结果:

None

torch.compile大大提高了解码速度, 吞吐量几乎翻了一番(如中间图表所示). 它还提高了预填充阶段的性能, 但程度较低. 不过, 在使用 bitsandbytes量化模型的情况下使用 torch.compile时, 对性能的影响微乎其微, 甚至可能略微减慢 A100 GPU 预填充阶段的速度.

在内存使用方面, 正如预期的那样, 编译模型会消耗更多内存. 增加的内存约为 +200MB(约为模型大小的 3%), 但相对较小, 仍在可控范围内.

使用 L4 GPU 的结果(Google Colab):

None

在内存只有 A100 一半的 L4 GPU 上, 我们观察到内存消耗有类似的增加, 但加速效果微乎其微. 这与在 A100 上的结果形成了鲜明对比, 凸显出 torch.compile的有效性会因使用的 GPU 不同而大相径庭.

总之, 值得在你的特定 GPU 上对 torch.compile 进行基准测试, 看看它是否能带来预期的性能提升.

总结一下

你应该使用 `torch.compile` 吗?

在标准设置中, torch.compile可以大大加快推理速度, 尤其是在不涉及量化的情况下. 不过, 我建议在模型开发过程的最后阶段, 也就是在配置了用于生产的所有功能和技术之后, 才启用 torch.compile. 这种方法非常重要, 因为根据模型和 GPU 的不同, torch.compile 在某些配置下可能无法达到预期效果.

例如, 我遇到了 FlashAttention 和 bfloat16 的兼容性问题, 它们在使用 torch.compile 时并不总是能顺利运行.

如果使用 LoRA Adapter, 请注意加速收益可能会较低. LoRA 组件是高度动态的, 因此有效编译具有挑战性.

简而言之, torch.compile是一个强大的加速工具, 但最好还是仔细测试并进行基准测试, 以确保它符合你的特定设置和要求. 它还会随着 PyTorch 的每个新版本而不断改进.

好吧, 今天的内容就分享到这里啦!

一家之言, 欢迎拍砖!

Happy Coding! Stay GOLDEN!

Torch 编译优化 Llama 3.2 可以事半功倍