開始使用

概觀

隆重推出 PyTorch 2.0，這是我們邁向下一代 PyTorch 2 系列版本的首要步驟。 在過去幾年中，我們從 PyTorch 1.0 創新和迭代到最新的 1.13，並轉移到新成立的 PyTorch 基金會，該基金會是 Linux 基金會的一部分。

PyTorch 除了擁有令人驚嘆的社群之外，其最大的優勢是我們繼續作為一流的 Python 整合、命令式風格、API 的簡潔性和選項。 PyTorch 2.0 提供相同的 eager-mode 開發和使用者體驗，同時從根本上改變和強化 PyTorch 在底層編譯器層級的運作方式。 我們能夠提供更快的效能，並支援動態形狀 (Dynamic Shapes) 和分散式 (Distributed)。

在下面，您將找到所有需要的信息，以便更好地了解 PyTorch 2.0 是什麼、它的發展方向，更重要的是，如何立即開始使用它（例如，教程、需求、模型、常見的常見問題）。 還有很多東西需要學習和開發，但我們期待社群的回饋和貢獻，使 2 系列更好，並感謝所有使 1 系列如此成功的人。

PyTorch 2.x：更快、更 Pythonic，並且一如既往地動態

今天，我們宣布 torch.compile，這是一項將 PyTorch 效能推向新高度的功能，並開始將 PyTorch 的部分程式碼從 C++ 移回 Python。 我們相信這是 PyTorch 的一個重要的全新方向 – 因此我們稱之為 2.0。torch.compile 是一個完全附加（且可選）的功能，因此根據定義，2.0 是 100% 向後相容的。

支援 torch.compile 的是新技術 – TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

• TorchDynamo 使用 Python 框架評估鉤子安全地捕獲 PyTorch 程式，這是一個重大的創新，是我們 5 年來對安全圖捕獲進行研發的結果

• AOTAutograd 重載 PyTorch 的 autograd 引擎，作為追蹤自動微分，用於產生提前 (ahead-of-time) 的向後追蹤。

• PrimTorch 將 2000 多個 PyTorch 運算子規範化為一組封閉的約 250 個原始運算子，開發人員可以針對這些運算子建構完整的 PyTorch 後端。 這大大降低了編寫 PyTorch 功能或後端的障礙。
• TorchInductor 是一個深度學習編譯器，能為多種加速器和後端產生快速程式碼。對於 NVIDIA 和 AMD GPU，它使用 OpenAI Triton 作為關鍵的建構模組。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 都是以 Python 編寫，並支援動態形狀 (dynamic shapes) (也就是說，可以傳入不同大小的 Tensor 而無需重新編譯)，這使得它們具有彈性、易於修改，並降低了開發人員和供應商的進入門檻。

為了驗證這些技術，我們使用了一組涵蓋各種機器學習領域的 163 個開源模型。我們精心打造了這個基準測試，以包含圖像分類、物件偵測、圖像生成、各種 NLP 任務 (例如語言模型、問答、序列分類)、推薦系統和強化學習等任務。我們將基準測試分為三個類別

• 來自 HuggingFace Transformers 的 46 個模型
• 來自 TIMM 的 61 個模型：Ross Wightman 提供的最先進 PyTorch 圖像模型集合
• 來自 TorchBench 的 56 個模型：來自 Github 各處精選的熱門程式碼庫

我們沒有修改這些開源模型，除了添加一個 torch.compile 呼叫來包裝它們。

然後，我們測量加速 (speedup) 並驗證這些模型的準確性。由於加速可能取決於資料類型，因此我們測量 float32 和自動混合精度 (Automatic Mixed Precision, AMP) 的加速。我們報告一個不均勻的加權平均加速，0.75 * AMP + 0.25 * float32，因為我們發現 AMP 在實務中更常見。

在這 163 個開源模型中，torch.compile 在 93% 的時間內有效，並且在 NVIDIA A100 GPU 上進行訓練時，模型的運行速度平均快 43%。 在 Float32 精度下，平均快 21%，在 AMP 精度下，平均快 51%。

注意事項： 在桌面級 GPU (例如 NVIDIA 3090) 上，我們測得的加速低於伺服器級 GPU (例如 A100)。 截至今日，我們的預設後端 TorchInductor 支援 CPU 和 NVIDIA Volta 和 Ampere GPU。 它（尚未）支援其他 GPU、xPU 或較舊的 NVIDIA GPU。

試試看： torch.compile 處於開發的早期階段。 從今天開始，您可以在 nightly 二進位檔中試用 torch.compile。 我們預計在 2023 年 3 月初推出第一個穩定的 2.0 版本。

在 PyTorch 2.x 的路線圖中，我們希望在效能和可擴展性方面進一步推動編譯模式。 其中一些工作正在進行中，正如我們今天在會議上討論的那樣。 有些工作尚未開始。 有些工作是我們希望看到的，但我們沒有足夠的資源自己完成。 如果您有興趣做出貢獻，請在本月開始的 「Ask the Engineers: 2.0 Live Q&A Series」 (詳細資訊在本文章末尾) 和/或透過 Github / 論壇與我們聊天。

使用者回饋

以下是一些 PyTorch 使用者對我們新方向的看法

Sylvain Gugger 是 HuggingFace transformers 的主要維護者

「只需添加一行程式碼，PyTorch 2.0 就能在訓練 Transformers 模型時提供 1.5 倍到 2 倍的加速。 這是自引入混合精度訓練以來最令人興奮的事情！」

Ross Wightman 是 TIMM 的主要維護者 (TIMM 是 PyTorch 生態系統中最大的視覺模型中心之一)

「它可以在大多數 TIMM 模型中開箱即用，用於推論和訓練工作負載，無需更改程式碼」

Luca Antiga 是 Lightning AI 的 CTO，也是 PyTorch Lightning 的主要維護者之一

「PyTorch 2.0 體現了深度學習框架的未來。 能夠有效地擷取 PyTorch 程式且無需使用者干預，並開箱即用地獲得大量的裝置加速和程式操作，為 AI 開發人員開啟了一個全新的維度。」

動機

我們對 PyTorch 的哲學一直是將彈性和可修改性放在首位，而將效能放在第二位。 我們努力實現

1. 高效能的 eager 執行
2. Pythonic 的內部結構
3. 良好的分散式、自動微分、資料載入、加速器等抽象

自 2017 年推出 PyTorch 以來，硬體加速器 (例如 GPU) 的運算速度提高了約 15 倍，記憶體存取速度提高了約 2 倍。 因此，為了保持 eager 執行的高效能，我們不得不將 PyTorch 內部結構的很大一部分移到 C++ 中。 將內部結構移到 C++ 中會降低它們的可修改性，並增加程式碼貢獻的進入門檻。

從第一天起，我們就知道 eager 執行的效能限制。 2017 年 7 月，我們開始了第一個開發 PyTorch 編譯器的研究專案。 編譯器需要讓 PyTorch 程式快速執行，但不能以犧牲 PyTorch 體驗為代價。 我們的關鍵標準是保留某些靈活性 - 支援動態形狀和動態程式，研究人員在各個探索階段都會使用這些特性。

技術概觀

多年來，我們在 PyTorch 中建立了幾個編譯器專案。 讓我們將編譯器分解為三個部分

• 圖形獲取 (graph acquisition)
• 圖形簡化 (graph lowering)
• 圖形編譯 (graph compilation)

在建構 PyTorch 編譯器時，圖形獲取是更困難的挑戰。

在過去的 5 年裡，我們建立了 torch.jit.trace、TorchScript、FX tracing、Lazy Tensors。 但它們都沒有給我們我們想要的一切。 有些很靈活但不快，有些很快但不靈活，有些既不快也不靈活。 有些有糟糕的使用者體驗 (例如靜默錯誤)。 雖然 TorchScript 很有希望，但它需要對您的程式碼以及您的程式碼所依賴的程式碼進行大量變更。 這種對程式碼進行大量變更的需求使其成為許多 PyTorch 使用者的絆腳石。

TorchDynamo：可靠且快速地獲取圖形

今年年初，我們開始研究 TorchDynamo，這是一種使用 PEP-0523 中引入的 CPython 功能的方法，稱為 Frame Evaluation API。 我們採用了一種資料驅動的方法來驗證其在圖形擷取方面的有效性。 我們使用了 7,000 多個以 PyTorch 編寫的 Github 專案作為我們的驗證集。 雖然 TorchScript 和其他工具甚至在 50% 的時間內都難以獲取圖形，而且通常有很大的開銷，但 TorchDynamo 在 99% 的時間內正確、安全且開銷可忽略不計地獲取了圖形 - 無需對原始程式碼進行任何變更。 這時我們知道，我們終於突破了多年來在靈活性和速度方面一直困擾我們的障礙。

TorchInductor：使用 define-by-run IR 的快速程式碼生成

對於 PyTorch 2.0 的新編譯器後端，我們從使用者編寫高效能自定義核心的方式中獲得了靈感：越來越多地使用 Triton 語言。 我們還希望編譯器後端使用與 PyTorch eager 類似的抽象，並且具有足夠的通用性，可以支援 PyTorch 中的各種功能。 TorchInductor 使用 Pythonic define-by-run 迴圈層級 IR，以自動將 PyTorch 模型對應到 GPU 上生成的 Triton 程式碼和 CPU 上的 C++/OpenMP。 TorchInductor 的核心迴圈層級 IR 僅包含約 50 個運算符，並且以 Python 實作，使其易於修改和擴展。

AOTAutograd：重用 Autograd 用於提前 (ahead-of-time) 圖形

對於 PyTorch 2.0，我們知道我們想要加速訓練。 因此，至關重要的是，我們不僅要擷取使用者層級的程式碼，還要擷取反向傳播。 此外，我們知道我們想要重用現有的經過實戰考驗的 PyTorch autograd 系統。 AOTAutograd 利用 PyTorch 的 torch_dispatch 可擴展機制來追蹤我們的 Autograd 引擎，使我們能夠「提前」擷取反向傳遞。 這使我們能夠使用 TorchInductor 加速我們的正向和反向傳遞。

PrimTorch：穩定的 Primitive 運算符

為 PyTorch 撰寫後端程式碼是一項具有挑戰性的任務。PyTorch 擁有超過 1200 個運算子 (operators)，如果考慮每個運算子的各種多載 (overloads)，則超過 2000 個。

因此，撰寫後端或跨領域功能變成一件令人精疲力盡的事情。在 PrimTorch 專案中，我們正致力於定義更小且更穩定的運算子集合。PyTorch 程式可以一致地降級 (lowered) 到這些運算子集合。我們的目標是定義兩個運算子集合：

• Prim ops：大約有 ~250 個運算子，它們相當於低階 (low-level) 運算子。這些運算子適合編譯器使用，因為它們的階層夠低，您需要將它們融合 (fuse) 在一起才能獲得良好的效能。
• ATen ops：大約有 ~750 個標準 (canonical) 運算子，適合直接匯出 (export) 使用。這些運算子適合已經整合到 ATen 層級的後端，或是不需要透過編譯從像 Prim ops 這樣的低階運算子集合中恢復效能的後端。

我們將在下方的「開發者/供應商體驗」章節中更詳細地討論這個主題。

使用者體驗

我們引入一個簡單的函式 torch.compile，它可以包裝您的模型並傳回一個已編譯的模型。

compiled_model = torch.compile(model)

這個 compiled_model 持有您模型的參考，並將 forward 函式編譯成更優化的版本。在編譯模型時，我們提供了一些調整參數：

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  # advanced backend options go here as kwargs
  **kwargs
) -> torch._dynamo.NNOptimizedModule

• mode 指定編譯器在編譯時應該優化什麼。

  • 預設模式是一個預設值，它試圖有效地編譯，而不會花費太長時間編譯或使用額外的記憶體。
  • 其他模式，例如 reduce-overhead，可以大幅減少框架的開銷，但會花費少量的額外記憶體。max-autotune 會長時間編譯，試圖為您提供它可以產生的最快程式碼。
• dynamic 指定是否啟用動態形狀 (Dynamic Shapes) 的程式碼路徑。某些編譯器優化無法應用於動態形狀的程式。明確說明您想要具有動態形狀還是靜態形狀的已編譯程式，將有助於編譯器為您提供更優化的程式碼。
• fullgraph 類似於 Numba 的 nopython。它將整個程式編譯成一個單一圖形，或者提供一個錯誤訊息說明為什麼它無法這樣做。大多數使用者不需要使用這個模式。如果您非常注重效能，那麼您可以嘗試使用它。
• backend 指定要使用的編譯器後端。預設情況下，使用 TorchInductor，但還有其他一些可用的後端。

編譯體驗旨在以預設模式提供最大的好處和最大的靈活性。以下是您在每種模式下獲得的心理模型。

現在，讓我們來看一個完整的範例，編譯一個真實的模型並執行它（使用隨機資料）。

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()

您第一次執行 compiled_model(x) 時，它會編譯模型。因此，需要更長的時間才能執行。後續的執行速度很快。

模式

編譯器有一些預設值，可以通過不同的方式調整編譯後的模型。您可能正在運行一個小型模型，由於框架開銷而速度緩慢。或者，您可能正在運行一個大型模型，該模型幾乎無法放入記憶體中。根據您的需求，您可能想要使用不同的模式。

# API NOT FINAL
# default: optimizes for large models, low compile-time
#          and no extra memory usage
torch.compile(model)

# reduce-overhead: optimizes to reduce the framework overhead
#                and uses some extra memory. Helps speed up small models
torch.compile(model, mode="reduce-overhead")

# max-autotune: optimizes to produce the fastest model,
#               but takes a very long time to compile
torch.compile(model, mode="max-autotune")

讀取和更新屬性

存取模型屬性的方式與在 eager 模式下相同。您可以像通常一樣存取或修改模型的屬性（例如 model.conv1.weight）。就程式碼校正而言，這是完全安全可靠的。TorchDynamo 會在程式碼中插入 guard，以檢查其假設是否成立。如果屬性以某種方式發生更改，那麼 TorchDynamo 會知道在需要時自動重新編譯。

# optimized_model works similar to model, feel free to access its attributes and modify them
optimized_model.conv1.weight.fill_(0.01)

# this change is reflected in model

Hooks

Module 和 Tensor 的 hooks 目前尚未完全正常運作，但隨著我們完成開發，它們最終會正常運作。

序列化

您可以序列化 optimized_model 或 model 的 state-dict。它們指向相同的參數和狀態，因此是等效的。

torch.save(optimized_model.state_dict(), "foo.pt")
# both these lines of code do the same thing
torch.save(model.state_dict(), "foo.pt")

目前您無法序列化 optimized_model。如果您希望直接保存該物件，請改為保存 model。

torch.save(optimized_model, "foo.pt") # Error
torch.save(model, "foo.pt")           # Works

推論和匯出

對於模型推論，在使用 torch.compile 產生已編譯的模型之後，在實際模型服務之前執行一些暖機 (warm-up) 步驟。這有助於減輕初始服務期間的延遲高峰。

此外，我們將引入一個名為 torch.export 的模式，該模式會仔細匯出整個模型和 guard 基礎架構，以用於需要保證和可預測延遲的環境。torch.export 需要對您的程式進行更改，尤其是當您具有資料依賴的控制流程 (data dependent control-flow) 時。

# API Not Final
exported_model = torch._dynamo.export(model, input)
torch.save(exported_model, "foo.pt")

這仍處於早期開發階段。請在 PyTorch 會議上觀看關於 Export Path 的演講，以了解更多詳細資訊。您也可以在本月的「Ask the Engineers: 2.0 Live Q&A Series」中參與此主題的討論（更多詳細資訊在本文章末尾）。

偵錯問題

已編譯模式是不透明且難以偵錯的。您會有以下問題：

• 為什麼我的程式在已編譯模式下崩潰？
• 已編譯模式是否與 eager 模式一樣準確？
• 為什麼我沒有看到加速？

如果已編譯模式產生錯誤或崩潰，或與 eager 模式產生不同的結果（超出機器精度限制），那麼這不太可能是您的程式碼的錯誤。但是，了解哪一部分程式碼是錯誤的原因是有用的。

為了幫助偵錯和可重現性，我們創建了幾個工具和日誌記錄功能，其中一個特別突出：The Minifier。

The minifier 會自動將您看到的問題簡化為一小段程式碼。這個小片段程式碼重現了原始問題，您可以提交包含簡化程式碼的 github issue。這將有助於 PyTorch 團隊輕鬆快速地修復問題。

如果您沒有看到您期望的加速，那麼我們有 torch._dynamo.explain 工具，它可以解釋程式碼的哪些部分導致了我們所謂的「graph breaks」。Graph breaks 通常會阻礙編譯器加速程式碼，減少 graph breaks 的數量可能會加速您的程式碼（達到收益遞減的某個極限）。

您可以在我們的 疑難排解指南 中閱讀有關這些以及更多內容。

動態形狀

在考慮支援 PyTorch 程式碼的通用性所必需的條件時，一個關鍵要求是支援動態形狀，並允許模型採用不同大小的 tensors，而無需每次形狀改變時都進行重新編譯。

截至今天，對動態形狀的支援是有限的，並且正在快速開發中。它將在穩定版本中全面推出。它由 dynamic=True 參數控制，並且我們在一個功能分支（symbolic-shapes）上有更多的進展，我們已經使用 TorchInductor 在具有完整符號形狀的訓練中成功運行了 BERT_pytorch。對於具有動態形狀的推論，我們有更多的覆蓋範圍。例如，讓我們看看一個常見的設定，其中動態形狀很有用 - 使用語言模型生成文本。

我們可以發現，即使形狀從 4 動態變化到 256，編譯模式始終能比 eager 模式提供高達 40% 的效能提升。在不支援動態形狀的情況下，常見的解決方案是將形狀填充到最接近的 2 的冪次方。然而，從下面的圖表可以看出，這會產生大量的效能開銷，並且也會導致更長的編譯時間。此外，填充有時很難正確地完成。

藉由在 PyTorch 2.0 的編譯模式中支援動態形狀，我們可以同時獲得最佳的效能以及易用性。

目前的工作正在快速發展，當我們對基礎架構進行重大改進時，可能會暫時讓某些模型出現效能倒退的情況。關於我們在動態形狀方面進展的最新消息，請見這裡。

分散式

總而言之，torch.distributed 的兩個主要分散式包裝器在編譯模式下運作良好。

DistributedDataParallel (DDP) 和 FullyShardedDataParallel (FSDP) 都能在編譯模式下運作，並且相較於 eager 模式，提供了更好的效能和記憶體利用率，但有一些注意事項和限制。

DistributedDataParallel (DDP)

DDP 依賴於 AllReduce 通訊與反向計算的重疊，以及將較小的、基於每層的 AllReduce 操作分組到「buckets」中以提高效率。 由 TorchDynamo 編譯的 AOTAutograd 函數會阻止通訊重疊 (當與 DDP 天真地組合時)，但通過為每個「bucket」編譯單獨的子圖，並允許通訊操作在子圖外部和子圖之間發生，可以恢復效能。目前，編譯模式下的 DDP 支援也需要 static_graph=False。有關 DDP + TorchDynamo 的方法和結果的更多詳細資訊，請參閱此文章。

FullyShardedDataParallel (FSDP)

FSDP 本身是一個 "beta" PyTorch 功能，並且由於能夠調整哪些子模組被包裝，以及通常有更多的配置選項，因此比 DDP 具有更高的系統複雜性。 如果配置了 use_original_params=True 標誌，FSDP 可以與 TorchDynamo 和 TorchInductor 配合使用，用於各種流行的模型。 目前預期會有與特定模型或配置的相容性問題，但將會積極改進，並且如果提交 github issue，可以優先處理特定模型。

使用者指定一個 auto_wrap_policy 參數，以指示將其模型的哪些子模組一起包裝在用於狀態分片的 FSDP 實例中，或者手動將子模組包裝在 FSDP 實例中。 例如，當每個「transformer block」被包裝在單獨的 FSDP 實例中時，許多 transformer 模型都能很好地工作，因此一次只需要實例化一個 transformer block 的完整狀態。 Dynamo 將在每個 FSDP 實例的邊界插入圖形中斷，以允許在 forward (和 backward) 中的通訊操作在圖形外部並與計算並行發生。

如果 FSDP 在沒有將子模組包裝在單獨的實例中的情況下使用，它會回退到與 DDP 類似的運行方式，但沒有 bucketing。 因此，所有梯度都在一個操作中減少，即使在 Eager 模式下，也可能沒有計算/通訊重疊。 此配置僅經過 TorchDynamo 的功能測試，而未經過效能測試。

開發者/供應商體驗

藉由 PyTorch 2.0，我們希望簡化後端 (編譯器) 的整合體驗。 為了做到這一點，我們專注於減少運算子的數量和簡化語義，這是啟動 PyTorch 後端所需的運算子集合。

以圖形形式來看，PT2 堆疊如下

從圖表的中間開始，AOTAutograd 以提前時間 (ahead-of-time) 的方式動態捕捉 autograd 邏輯，以 FX 圖形格式產生 forward 和 backward 運算子的圖形。

我們提供了一組強化的分解 (即用其他運算子編寫的運算子實現)，可以用來減少後端需要實現的運算子數量。我們還通過選擇性地重寫複雜的 PyTorch 邏輯 (包括 mutations 和 views)，通過稱為functionalization 的過程來簡化 PyTorch 運算子的語義，並保證運算子元數據資訊，例如形狀傳播公式。這項工作正在積極進行中；我們的目標是提供一套primitive和stable的約 250 個運算子，具有簡化的語義，稱為PrimTorch，供應商可以利用 (即選擇加入) 它們來簡化他們的整合。
在減少和簡化運算子集合之後，後端可以選擇在 Dynamo (即中間層，緊接在 AOTAutograd 之後) 或 Inductor (較低層) 進行整合。 我們在下面描述了做出此選擇的一些考慮因素，以及圍繞後端混合的未來工作。

Dynamo 後端

具有現有編譯器堆疊的供應商可能會發現最容易整合為 TorchDynamo 後端，接收以 ATen/Prims IR 表示的 FX Graph。 請注意，對於訓練和推理，整合點將緊接在 AOTAutograd 之後，因為我們目前將分解應用為 AOTAutograd 的一部分，並且如果目標是推理，則僅跳過 backward 特定的步驟。

Inductor 後端

供應商還可以將他們的後端直接整合到 Inductor 中。 Inductor 接收由 AOTAutograd 產生的圖形，該圖形由 ATen/Prim 運算組成，並將它們進一步降低到迴圈級別的 IR。 今天，Inductor 為 pointwise、reduction、scatter/gather 和 window 運算提供到其迴圈級別 IR 的降低。 此外，Inductor 創建融合組、進行索引簡化、維度崩潰並調整迴圈迭代順序，以支持有效的代碼生成。 然後，供應商可以通過提供從迴圈級別 IR 到特定於硬體的代碼的映射來進行整合。 目前，Inductor 有兩個後端：(1) 生成多線程 CPU 代碼的 C++，(2) 生成高性能 GPU 代碼的 Triton。 這些 Inductor 後端可用作其他後端的靈感。

後端混合介面 (即將推出)

我們構建了用於將 FX 圖形劃分為子圖的實用程序，這些子圖包含後端支援的運算子，並積極地執行其餘部分。 這些實用程序可以擴展以支持 "後端混合"，配置要為哪個後端運行哪些部分的圖形。 但是，目前還沒有穩定的介面或合約供後端公開其運算子支援、對運算子模式的偏好等。 這仍然是一項正在進行的工作，我們歡迎早期採用者的回饋。

最後的想法

我們對 PyTorch 2.0 及以後採取的方向感到非常興奮。 到最終 2.0 版本的道路將是坎坷的，但請儘早加入我們的旅程。 如果您有興趣深入研究或貢獻編譯器，請繼續閱讀下面的內容，其中包括有關如何開始 (例如，教程、基準、模型、常見問題) 和詢問工程師：2.0 線上問答系列 (從本月開始) 的更多資訊。 其他資源包括

• 開始使用
• 教程
• 文件
• 開發者討論

使用 PyTorch 2.0 加速 Hugging Face 和 TIMM 模型

作者：Mark Saroufim

torch.compile() 使您可以輕鬆地使用不同的編譯器後端進行實驗，只需單行裝飾器 torch.compile() 即可使 PyTorch 代碼更快。 它可以直接在 nn.Module 上作為 torch.jit.script() 的替代品使用，但不需要您對原始程式碼進行任何更改。 我們希望這一行程式碼的更改能為您提供 30%-2 倍的訓練時間加速，這適用於您已在運行的絕大多數模型。

opt_module = torch.compile(module)

torch.compile 支援任意的 PyTorch 程式碼、控制流程、變動，並且實驗性地支援動態形狀。我們對這項開發感到非常興奮，因此稱它為 PyTorch 2.0。

這次發布對我們來說與眾不同之處在於，我們已經對一些最受歡迎的開源 PyTorch 模型進行了基準測試，並取得了顯著的加速，範圍從 30% 到 2 倍不等 https://github.com/pytorch/torchdynamo/issues/681。

這裡沒有任何花招，我們只是 pip 安裝了流行的函式庫，例如 https://github.com/huggingface/transformers、https://github.com/huggingface/accelerate 和 https://github.com/rwightman/pytorch-image-models，然後在它們上面執行 `torch.compile()`，就是這樣。

同時獲得效能和便利性是很罕見的，但這也是核心團隊對 PyTorch 2.0 感到如此興奮的原因。

需求

針對 GPU (較新一代的 GPU 將看到顯著更好的效能)

pip3 install numpy --pre torch --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

針對 CPU

pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

選用：驗證安裝

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

選用：Docker 安裝

我們也在 PyTorch nightly binaries 中提供了所有必要的依賴項，您可以透過以下方式下載：

docker pull ghcr.io/pytorch/pytorch-nightly

對於臨時的實驗，只需確保您的容器可以存取您的所有 GPU 即可

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

開始使用

請閱讀 Mark Saroufim 的完整部落格文章，他將帶您瀏覽教學課程和真實模型，讓您今天就可以試用 PyTorch 2.0。

我們使用 PyTorch 的目標是建立一個廣度優先的編譯器，它可以加速人們在開源中運行的絕大多數實際模型。 Hugging Face Hub 最終成為我們非常有價值的基準測試工具，確保我們所做的任何最佳化實際上都有助於加速人們想要運行的模型。

部落格教學課程將向您展示如何重現這些加速效果，以便您可以像我們一樣對 PyTorch 2.0 感到興奮。所以請試用 PyTorch 2.0，享受免費的效能提升，如果您沒有看到效能提升，請開啟一個 issue，我們將確保您的模型得到支援 https://github.com/pytorch/torchdynamo/issues

畢竟，除非您的模型實際運行得更快，否則我們不能聲稱我們創造了一個廣度優先的編譯器。

常見問題

1. 什麼是 PT 2.0？
   2.0 是最新的 PyTorch 版本。 PyTorch 2.0 提供相同的 eager-mode 開發體驗，同時透過 torch.compile 新增了編譯模式。這種編譯模式有可能在訓練和推論期間加速您的模型。

2. 為什麼是 2.0 而不是 1.14？
   PyTorch 2.0 就是原本的 1.14。我們發布了大量新功能，我們相信這些功能會改變您有意義地使用 PyTorch 的方式，因此我們將其稱為 2.0。

3. 我該如何安裝 2.0？有其他額外的要求嗎？

   安裝最新的 nightly 版本

   CUDA 11.8
   

    pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118
   

   CUDA 11.7

    pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117
   

   CPU

    pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu
   

4. 2.0 的程式碼是否與 1.X 向下相容？
   是的，使用 2.0 不會要求您修改 PyTorch 工作流程。單行程式碼 model = torch.compile(model) 可以最佳化您的模型以使用 2.0 堆疊，並與其餘的 PyTorch 程式碼順利運行。這是完全可選的，您不需要使用新的編譯器。

5. 2.0 預設啟用嗎？
   2.0 是發布版本的名稱。 torch.compile 是在 2.0 中發布的功能，您需要明確使用 torch.compile。

6. 我該如何將我的 PT1.X 程式碼遷移到 PT2.0？
   您的程式碼應該可以按原樣工作，而不需要任何遷移。如果您想使用 2.0 中引入的新編譯模式功能，您可以從使用一行程式碼最佳化您的模型開始：model = torch.compile(model)。
   雖然加速主要在訓練期間觀察到，但如果您的模型運行速度比 eager mode 快，您也可以將其用於推論。

    import torch
         
    def train(model, dataloader):
      model = torch.compile(model)
      for batch in dataloader:
        run_epoch(model, batch)
   
    def infer(model, input):
      model = torch.compile(model)
      return model(\*\*input)
   

7. 為什麼我應該使用 PT2.0 而不是 PT 1.X？
   請參閱問題 (2) 的答案。

8. 當運行 PyTorch 2.0 時，我的程式碼有什麼不同？
   PyTorch 2.0 開箱即用與 PyTorch 1.x 相同，您的模型以 eager-mode 運行，即每一行 Python 程式碼都依序執行。
   在 2.0 中，如果您將您的模型包裝在 model = torch.compile(model) 中，您的模型在執行前會經過 3 個步驟：
   1. 圖獲取 (Graph acquisition)：首先，模型被重寫為子圖塊。可以被 TorchDynamo 編譯的子圖被“扁平化”，而其他子圖（可能包含控制流程程式碼或其他不受支援的 Python 結構）將回退到 Eager-Mode。
   2. 圖降低 (Graph lowering)：所有的 PyTorch 操作都被分解為其特定於所選後端的組成核心。
   3. 圖編譯 (Graph compilation)：其中核心呼叫其相應的低階特定於設備的操作。
9. PT2.0 為 PT 新增了哪些新組件？
   • TorchDynamo 從 Python bytecode 產生 FX Graphs。它使用 guards 維護 eager-mode 功能，以確保產生的圖是有效的 (閱讀更多)
   • AOTAutograd 生成與 TorchDynamo 捕獲的前向圖對應的後向圖 (閱讀更多)。
   • PrimTorch 將複雜的 PyTorch 操作分解為更簡單和更基本的 ops (閱讀更多)。
   • [後端] 後端與 TorchDynamo 集成，將圖編譯為可以在加速器上運行的 IR。例如，TorchInductor 將圖編譯為 Triton 以進行 GPU 執行，或編譯為 OpenMP 以進行 CPU 執行 (閱讀更多)。

10. 2.0 目前支援哪些編譯器後端？
    預設且最完整的後端是 TorchInductor，但 TorchDynamo 有一個不斷增長的後端列表，可以透過呼叫 torchdynamo.list_backends() 來找到。

11. 分散式訓練如何與 2.0 協同工作？
    在編譯模式下，DDP 和 FSDP 的運行速度比 FP32 中的 Eager-Mode 快達 15%，在 AMP 精度中快達 80%。 PT2.0 進行了一些額外的最佳化，以確保 DDP 的通訊-計算重疊與 Dynamo 的部分圖創建良好地協同工作。確保您以 static_graph=False 運行 DDP。更多細節請參閱此處。

12. 我該如何了解更多關於 PT2.0 的開發資訊？
    PyTorch 開發者論壇是直接從構建它們的開發者那裡了解 2.0 組件的最佳場所。

13. 我的程式碼在使用 2.0 的編譯模式後運行速度變慢了！
    效能下降最可能的原因是圖中斷太多。例如，您模型的前向傳播中的一個看似無害的 print 語句會觸發圖中斷。我們有方法來診斷這些問題 - 閱讀更多此處。

14. 我先前運行的程式碼在使用 2.0 的編譯模式後崩潰了！我該如何對其進行除錯？
    以下是一些分類程式碼可能失敗的位置並打印有用日誌的技巧：https://pytorch.dev.org.tw/docs/stable/torch.compiler_faq.html#why-is-my-code-crashing。

諮詢工程師：2.0 線上問答系列

我們將舉辦一系列的線上問答環節，讓社群能與專家進行更深入的提問與對話。請隨時回來查看整年度的完整主題行事曆。如果您無法參加：1) 問答環節將會錄影以供日後觀看；2) 您可以每週五太平洋標準時間上午 10 點參加我們的開發基礎設施辦公室時間：https://github.com/pytorch/pytorch/wiki/Dev-Infra-Office-Hours。

請點擊此處查看日期、時間、說明和連結。

免責聲明：請勿在加入線上會議和提交問題時分享您的個人資訊、姓氏或公司。

觀看 PyTorch Conference 的演講

• TorchDynamo
• TorchInductor
• 動態形狀
• 匯出路徑