編譯

概述

Pytorch 2.0 引入了 torch.compile，這是一個強大的功能，它透過 JIT（即時編譯）將 PyTorch 程式碼編譯成最佳化的核心，從而使 PyTorch 程式碼執行得更快。torch.compile 的主要特性包括：

• 效能提升：透過最佳化計算圖，顯著加快模型執行速度。
• 易於使用：只需極少的程式碼更改即可實現，使其非常易於上手。
• 相容性：與現有的 PyTorch 程式碼和模型無縫協作。

當與 Accelerate 一起使用時，torch.compile 可以平滑地整合到分散式訓練工作流中，讓您同時受益於分散式執行和編譯最佳化。

編譯後代碼的首次執行通常需要更長的時間，因為它包含了編譯時間，但後續的執行會快得多。為了在不同場景下獲得最佳效能，torch.compile 提供了多種模式，如 "default"、"reduce-overhead"（它使用 CUDA 圖來進一步減少開銷）和 "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 倉庫中找到。我們基準測試的主要發現是：

1. 效能相當：區域編譯提供了與完全編譯相似的效能加速，特別是對於較大的模型。
2. 更快的編譯速度：區域編譯顯著減少了編譯模型所需的時間，使其成為更高效的部署選擇。
3. 批次大小的影響：隨著批次大小的增加，編譯策略之間的效能差異會減小，這表明在這些情況下，編譯的開銷影響較小。
4. 模型大小的考慮：區域編譯的好處在較大的模型中更為明顯，因為編譯時間的節省可能非常可觀。
5. 實際應用：對於實際應用，區域編譯是最佳化訓練冷啟動時間的實用選擇，尤其是在處理大型模型時。

結論

完全編譯和區域編譯都可以顯著加速您的模型。區域編譯在編譯時間和執行時效能之間提供了一個實際的平衡，尤其適用於訓練具有大批次的大型模型。