torch.compile

torch.compile 將 PyTorch 程式碼編譯成最佳化的核心，顯著加快推理速度。此功能依賴於 TorchDynamo 將程式碼編譯成圖，並依賴於 TorchInductor 將圖進一步編譯成最佳化的核心。它是一個強大的最佳化工具，在許多情況下，只需新增一行程式碼即可。

使用 `torch.compile` 封裝模型，即可編譯並返回最佳化後的模型。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

有幾個引數可以自定義編譯過程。下面列出了其中兩個更重要的引數。有關引數的完整列表，請參閱 torch.compile 文件。

模式

`mode` 引數提供了多種編譯效能選項。嘗試不同的模式，看看哪種最適合你的用例。

• `default` 是速度和記憶體之間的平衡選項。
• `reduce-overhead` 以犧牲少量記憶體為代價減少 Python 開銷，但它可能更快。
• `max-autotune` 提供最快的速度，但編譯時間更長。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph 嘗試將整個模型編譯成單個圖以最大限度地提高效能。如果遇到圖中斷，`torch.compile` 將引發錯誤，這意味著它無法將模型編譯成單個圖。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

基準測試

請參閱下表，瞭解在不同 GPU 和批處理大小下，針對不同視覺任務，在同一影像上啟用和停用 torch.compile 時平均推理時間（毫秒）的效能基準比較。

在下表中選擇**子集**以切換不同的 GPU，以及 PyTorch nightly 2.1.0dev 和啟用 `reduce-overhead` 模式的 torch.compile 的基準測試。