返回首頁

GitHub Repo

PyTorch 預編譯提案封裝運算核心,補強重播階段的零編譯檢查

作者回報,整組變更讓推薦模型的 16 個分散式程序完成重播,未觸發編譯。提案仍為草稿,跨模型驗證與部署效益尚待確認。

Pytorch Deepdream (https://github.com/gordicaleksa/pytorch-deepdream) by gordicaleksa (https://github.com/gordicaleksa/pytorch-deepdream/commits?author=gordicaleksa) · MIT · Image source
zh-Hant

PyTorch 開發者於 9 月 27 日更新預編譯提案,嘗試把重播所需的運算核心一併封裝。作者回報,整組變更已在一個使用 FSDP2 的大型推薦模型上,讓 16 個分散式程序完成 100 批次且未觸發編譯;截至查核時,PR #198789 仍為草稿,尚非正式版功能。[提案與測試](https://github.com/pytorch/pytorch/pull/198789)

問題在於,載入計算圖並不保證底層核心已備妥。既有官方文件將快取分為計算圖、Triton 編譯結果與自動調校等層次;Mega-Cache 可匯出、載入這些產物,但會檢查 PyTorch、Triton 版本及 CUDA GPU 是否相符。官方也說明,自動調校會實測候選核心並選出最快版本;重新調校本身仍有成本,跨程序沿用快取時,也需核對執行環境是否符合限制。[官方快取文件](https://docs.pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html)

新方案提出三個步驟:`capture_runtime()` 記錄實際啟動的 Triton 核心與 C++ 二進位檔;`finalize_cache()` 固定核心啟動器及快取內容;`prepare_runtime()` 則在載入模型產物前驗證並準備核心。直接使用 Triton JIT 的路徑,還取決於 Triton 是否提供執行期快取匯出介面。[API 設計](https://github.com/pytorch/pytorch/pull/198789)

配套的禁止編譯設計,會讓計算圖編譯、核心快取未命中及自動調校直接拋錯,避免重播流程悄悄補做編譯。這對工程團隊的意義,是將「預先編譯是否完整」變成可檢查的執行條件;不過,最初的配套 PR 已關閉並由後續提案接替,介面仍可能調整。[配套設計](https://github.com/pytorch/pytorch/pull/198767)

從部署角度推論,把選定核心與模型產物綁在一起,有助於檢查不同程序是否沿用同一套運算路徑,並使缺漏提早暴露。這類檢查對需要可預測啟動行為、又會同時啟動多個工作程序的部署尤其相關。不過,現有結果仍是作者對單一模型的測試,不能推廣成所有模型的效能或數值保證。接下來應追蹤上游審查、完整 CI、其他工作負載重現,以及不同輸入形狀能否被快取涵蓋;正式部署前也需量測啟動時間與快取體積,確認節省的編譯成本是否被產物傳輸與載入開銷抵銷。[驗證範圍](https://github.com/pytorch/pytorch/pull/198789)

來源

  1. PyTorch PR #198789:預編譯執行期快取提案
  2. PyTorch PR #198767:禁止編譯與嚴格載入的原始配套設計
  3. Compile Time Caching in torch.compile