编译

概述

PyTorch 2.0 引入了 torch.compile，这是一个强大的功能，它通过 JIT 编译将 PyTorch 代码转换为优化后的内核，从而提升 PyTorch 代码的运行速度。torch.compile 的主要特性包括：

性能提升：通过优化计算图，显著加快模型执行速度。
易用性：实现所需的代码改动极少，因此上手门槛很低。
兼容性：可与现有的 PyTorch 代码和模型无缝配合使用。

与 Accelerate 配合使用时，torch.compile 可以无缝集成到分布式训练工作流中，使您能够同时受益于分布式执行和编译优化。

首次执行编译后的代码通常需要更长的时间，因为它包含了编译耗时，但后续运行会显著加快。为了在不同场景下获得最佳性能，torch.compile 提供了多种模式，例如 "default"、"reduce-overhead"（使用 CUDA graphs 进一步降低开销）以及 "max-autotune"（执行广泛的自动调优以找到适合您模型的最佳内核）。

使用 torch.compile 与 Accelerate

Accelerate 提供了 TorchDynamoPlugin，以便轻松且无缝地将 torch.compile 集成到您的训练脚本中。

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

它与 Accelerate 的所有其他功能和插件兼容，包括混合精度、分布式训练（DDP、FSDP、Deepspeed）等。

区域编译

与其尝试编译整个模型（这通常会导致优化问题空间过大），不如采用区域编译。区域编译针对同一类的重复块，并按顺序编译它们以命中编译器缓存。例如，在 GPT2LMHeadModel 中，重复的块/类是 GPT2Block，可以通过 model.transformer.h[0] 访问。模型的其余部分（例如 model.lm_head）则单独编译。

这使我们能够加速 LLM 和 Transformers 等模型的编译开销/冷启动。有关更多详细信息，请参阅 https://pytorch.org/tutorials/recipes/regional_compilation.html。

如何使用区域编译

可以通过在 TorchDynamoPlugin 配置中设置 use_regional_compilation=True 来启用该功能：

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

您也可以像使用 torch.compile 一样，直接使用 accelerate.utils.compile_regions 实用工具。

区域编译的优势

我们使用 PyTorch 中的 torch.compile 功能，对完整编译和区域编译进行了广泛的基准测试。完整结果可在 accelerate 仓库中查看。我们基准测试的主要发现如下：

性能相当：区域编译带来的性能加速与完整编译相似，尤其对于较大的模型。
更快的编译：区域编译显著缩短了模型编译所需的时间，使其成为部署时更高效的选择。
批大小影响：随着批大小的增大，不同编译策略之间的性能差异会减小，这表明在这些场景下编译带来的开销影响较小。
模型规模考量：区域编译的优势在更大的模型中更为显著，因为编译时间的节省可能非常可观。
实际应用：在实际应用中，区域编译是优化训练冷启动时间的实用选择，尤其是在处理大型模型时。

结论

完整编译和区域编译都能显著提升模型速度。区域编译在编译时间和运行时性能之间提供了实用的平衡，尤其适用于使用较大批量大小训练大型模型的场景。

Update on GitHub