PyTorch 團隊正選技術路線圖，近百頁文檔披露 2024 下半年發展方向

最近，PyTorch 團隊首次公佈了開發路線圖，由內部技術文檔直接修改而來，披露了這個經典開源庫下一步的發展方向。

如果你在 AI 領域用 Python 開發，想必 PyTorch 一定是你的老朋友之一。2017 年，Meta AI 發佈了這個機器學習和深度學習領域的開源庫，如今已經走到了第 7 個年頭。

根據 Assembly AI 2021 年的統計數據，HuggingFace 上最受歡迎的 top 30 模型都能在 PyTorch 上運行，有 92% 的模型是 PyTorch 專有的，這個佔比讓包括 TensorFlow 在內的一眾競爭對手都望塵莫及。

就在 7 月 10 日，PyTorch 的工程團隊首次公開發佈了他們的路線圖文檔，闡述 2024 年下半年的發展方向。

Meta 共同創始人、領導 PyTorch 團隊的 Soumith Chintala 在Twitter上官宣了這個消息。

他表示，希望公開工程師們的研發動機和目標。

「雖然所有 PyTorch 開發都在 GitHub 上公開，但各個 PyTorch 附屬公司的團隊編寫的實際規劃和路線圖文檔並不公開，因此我們決定做出改變，以提高透明度。」

PyTorch 團隊的技術項目經理 Gott Brath 也在論壇中發表了類似的聲明。

我們一直在考慮，如何分享團隊在 PyTorch 上所做的工作的路線圖。我們每半年進行一次規劃，因此這些是我們針對 PyTorch 中多個關鍵領域的 2024 年 H2 OSS 計劃的一些公開版本。

這些文件基本就是 PyTorch 團隊內部的文檔和工作規劃，刪減掉了一些內容就發佈出來成為路線圖，其中涉及 PyTorch 的如下幾個方面：

– 核心庫與核心性能

– 分佈式

– torchune、Torchrec、TorchVision

– PyTorch Edge

– 數據加載（DataLoading）

– 編譯器核心及部署

– 開發者基礎設施

每個文檔都至少包含三個部分的內容，以 OKR 的思路展開：

– 背景

– Top5 關注領域及目標：目標、關鍵結果、已知或未知風險以及相應緩解措施（最多一頁）

– 提升工程水平的 Top3~5 個方面：BE Pillar 分類、目標、指標 / 狀態 / 具體目標、已知或未知風險以及緩解措施、影響 / 成本、優先級 / 信心程度（最多一頁）

其中 BE Pillar 可以看作 Meta 寫給開發團隊的「五句箴言」，具體內容是：

Better Code, Better Doc, Empowering teams, Modern Code, Better Architecture

「最多一頁」的規定不知道有沒有戳到卷文檔長度的開發人員，畢竟文檔貴精不貴長，將眾多開發需求精簡到一頁的內容不僅節省同事時間，也十分考驗撰寫者的功力。

此外，文檔中也可以看出 Meta 開發團隊的一些優秀思路，比如重視各個模塊團隊的協作、重視和外部合作夥伴的 API 集成和共同開發，重視與開源社區和開發者的互動。

當推出 ExecuTorch 這樣的新代碼庫，或者想要提升 PyTorch 編譯器影響力時，團隊一般都會從兩方面思路入手：一是鉚足力氣提升性能，把目標直接頂到 SOTA；另一方面從深度集成入手，提供更多開箱即用的案例。

或許，這些都是 Meta 多年來在開源領域如魚得水、風生水起的關鍵所在。

以下是各個文檔內容的部分截取和概括。

原文地址：https://dev-discuss.pytorch.org/t/meta-pytorch-team-2024-h2-roadmaps/2226

核心庫與核心性能

文檔中涉及到的核心庫包括 TendorDict、torchao、NN、TorchRL 等。

性能方面，PyTorch 團隊提出了在模型訓練和推理方面實現 SOTA 性能的目標，措施包括引入架構優化技術和高性能 kernel，與整個 PyTorch 技術棧形成搭配組合。

過去一年的時間見證了 GenAI 的快速發展，許多支持研究領域進行開發的外部庫應運而生，但其中很多並不直接依賴 PyTorch，這會威脅到 PyTorch 在科研領域的主導地位。

為了重新跟上節奏，PyTorch 將為量化、稀疏化、MoE 和低精度訓練等常用開發技術提供支持，包括構建模塊和 API（主要集成在 torchao 中），幫助各類 Transformer 架構的模型提升性能。

torchao 庫可以支持研究人員在 PyTorch 框架內自定義高性能的 dtype、layout 和優化技巧，將使用範圍擴展到訓練、推理、調優等各種場景。

此外，核心庫的更新將包括以下方面：

– 推出的自動優化庫 torchao 已經取得了突破性的成功，下一步提升其代碼組織性，並將其中的數值運算與核心庫分開

– 解決 TendorDict 的核心模塊性，支持加載 / 存儲的序列化，並使其在 eager 模式下的運行速度提高 2 倍

– 繼續上半年在內存映射加載（memory mapped load）方面的成功，繼續提升模型加載 / 存儲的性能和安全性

– 將 TorchRL 的開銷降低 50%

– 加入對 NoGIL 的核心支持

– 修復用戶反映的 TORCH_env 變量不起作用的問題

文檔中還提及了要實現對 nn.transformer 模塊的棄用，表示會發佈一系列教程和用例，展示如何使用 torch.compile、sdpa、NJT、FlexAttention、custom_op、torchao 等模塊構建 Transformer。

分佈式

LLM 的預訓練通常橫跨數十個甚至上千個 GPU，而且由於模型的參數規模逐漸增大，推理和微調也很難用單個 GPU 完成。

因此，PyTorch 下一步對「分佈式」的佈局全面涵蓋了訓練、推理、微調這三個環節，提出要達成超大規模分佈式訓練、高內存效率的微調、多主機分佈式推理。

訓練

PyTorch 原生支持的並行模式主要包括以下幾種：

– 完全分片數據並行（full sharded data parallel，FSDP）

– 混合分片數據並行（hybrid sharding data parallel，HSDP）

– 張量並行（tensor parallel，TP）

– 流水線並行（pipeline parallel，PP）

– 序列並行（sequence parallel，SP）

– 上下文並行（context parallel，CP）

PyTorch 希望在 TorchTitan 中將各種並行方式進一步模塊化，讓開發者可以自由組合，根據需要實現 N 維並行。

文檔中特別提到，對 MoE 和多模態這兩種新興的架構需要添加支持，比如專家並行、路由算法的優化。

除了 TorchTitan 本身的更新，分佈式團隊還需要與編譯器團隊進一步緊密合作，更好地與 torch.compile 模塊集成，為大規模分佈式場景帶來額外的性能提升。

微調與推理

微調：聯合 torchtune，將 FSDP2 LoRA / QLoRA 方案投入使用，以及支持模形狀態字典的 NF4 量化

推理：PP 和 DP 已經成為分佈式 API 的核心，下一步需要關注 torchtitan 的分佈式推理，支持大模型 PP + 異步 TP 方式，將給出案例展示

文檔中還提到，會將 HuggingFace 的推理 API 從 PiPPy 遷移到 PyTorch（由 HuggingFace 完成）。

torchtune、TorchRec、TorchVision

torchtune

torchtune 的推出旨在幫助用戶更方便微調 LLM，這也是官方給出的 Llama 模型微調的方案。

torchtune 定義的「微調」範圍非常廣，主要可以概括為三類場景：

– 對特定領域數據集或者下遊任務的模型適應

– 獎勵和偏好建模，比如 RLHF、DPO 等

– 包含蒸餾與量化的訓練過程

下半年的更新將支持為 agent 工作流進行的微調，同時著重關注微調性能的提升。

團隊會與 compile、core、distributed 等模塊進行合作，提供高效率微調，並在 PyTorch 生態內建立有代表性的微調性能基準。

由於 torchtune 也是一個較新的開源庫，因此與開源社區的互動也必不可少。

文檔提出發佈博客文章和教程、舉辦技術講座等方式，提升用戶的理解；並會定義量化指標，衡量 torchturn 在 LLM 生態中的貢獻份額。

除了開源社區，torchtune 還會與至少一個合作夥伴集成，參與到它們的社區中，以促進 torchtune 的使用。

TorchVision

TorchVision 作為 CV 領域內的絕對主宰者，技術也相對成熟，因此路線圖中提出的更新很少。

團隊將繼續在預處理方向努力，在圖像編碼 / 解碼空間中支持更多格式（如 WebP、HEIC）和平台（如 CUDA），並提升 jpeg 格式在 GPU 上的編碼 / 解碼性能。

TorchRec

TorchRec 旨在提供大規模推薦系統中常用的稀疏性和並行性原語，將秋季推出第一個穩定版本 TorchRec 1.0。

Edge

目前，開源庫 ExecuTorch 已經推出了 Alpha 版本，主要依賴 torch.compile 和 torch.export，用於支持移動設備和邊緣設備（如 AR / VR、可穿戴設備）上的模型分析、調試和推理。

下半年，Edge 團隊將推出 xecuTorch 的 Beta 版本，同時為 Meta 的 Llama 系列模型和其他開源模型提供 PyTorch 生態內的解決方案。

關鍵目標中主要涵蓋兩個方向。一是為設備上 AI 提供基礎功能和可靠基礎設施，包括：

– 確保 C++ 和 Python 的 API 穩定性

– 實現一系列核心功能：支持模型壓縮、代理緩存位置管理、數據和程序分離

二是為這個新生的代碼庫保駕護航，培育開源社區內的影響力，同時與 Arm、Apple 和 Qualcomm 等公司保持良好合作關係。

其中社區影響力的目標甚至被量化到，要求代碼在 GitHub 上得到 3k 標星，500 次複製（fork）。有興趣的吃瓜群眾可以去持續關注一下，看看團隊能不能在年底完成這個 OKR。

數據加載

基於 Apache Arrow 格式的 HuggingFace datasets 庫憑藉無內存限制的高速加載 / 存儲，近幾年異軍突起，似乎搶走了 PyTorch 相關功能的風頭。

數據加載的文檔中開篇就提出了雄心壯誌，要讓 TorchData 庫再次偉大，重新確立 PyTorch 在數據加載方面的主宰地位。

要達到這個目標，就需要讓相關功能變得靈活、可擴展、高性能、高內存效率，同時實現傻瓜式操作，支持各種規模的多模態訓練。

具體的更新目標包括以下幾個方面：

– DataLoader 的功能開發和接口都將貫徹 GitHub 優先的原則，DataPipes 和 DataLoader v2 則將被逐步被棄用、刪除

– 確保 TorchTune、TorchTitan、HuggingFace、TorchData 之間的清晰邊界和良好互通性，支持多數據集、多模態數據加載

– HuggingFace 使用 StatefulDataLoader 的 API，確保兼容性，及時更新樣例和測試用例

編譯器核心及部署

PyTorch 的編譯器核心功能經過多年發展已經趨於完善，目前亟待彌補的只是對 LLM 和 GenAI 領域的更深度集成和更多優化支持。

路線圖提出，要將 torch.compile () 函數帶到 LLM 和 GenAI 的使用週期的各個方面（推理、微調、預訓練），讓重要模型在發行時就搭載原生的 PyTorch 編譯。

為了實現這個目標，文檔提出了很多具體措施，比如與 torchtune 與 TorchTitan 團隊合作，提升編譯性能，並在下半年發佈至少兩個高知名度模型的原生 PyTorch 編譯版本。

此外，編譯器可能添加可視化功能，在 non-eager 訓練模式下生成表達前向計算 / 後向傳播過程的模型圖。

用戶支持方面也有諸多規劃，比如提升系統的監控性和可觀察性，幫助戶自行調試編譯問題。關鍵目標還包括建立用戶支持團隊，針對幾個關鍵領域（數據類、上下文管理等），解決開發者在 GitHub 等平台上發佈的問題。

參考資料：

• https://dev-discuss.pytorch.org/t/meta-pytorch-team-2024-h2-roadmaps/2226

• https://x.com/soumithchintala/status/1811060935211049046

• https://www.assemblyai.com/blog/pytorch-vs-tensorflow-in-2023/

本文來自微信公眾號：新智元（ID：AI_era）

廣告聲明：文內含有的對外跳轉鏈接（包括不限於超鏈接、二維碼、口令等形式），用於傳遞更多信息，節省甄選時間，結果僅供參考，IT之家所有文章均包含本聲明。