如果你在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在推特上官宣了這個消息。

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

「雖然所有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等平臺上發布的問題。