如果你在AI領(lǐng)域用Python開(kāi)發(fā)，想必PyTorch一定是你的老朋友之一。2017年，Meta AI發(fā)布了這個(gè)機(jī)器學(xué)習(xí)和深度學(xué)習(xí)領(lǐng)域的開(kāi)源庫(kù)，如今已經(jīng)走到了第7個(gè)年頭。

根據(jù)Assembly AI 2021年的統(tǒng)計(jì)數(shù)據(jù)，HuggingFace上最受歡迎的top 30模型都能在PyTorch上運(yùn)行，有92%的模型是PyTorch專(zhuān)有的，這個(gè)占比讓包括TensorFlow在內(nèi)的一眾競(jìng)爭(zhēng)對(duì)手都望塵莫及。

就在7月10日，PyTorch的工程團(tuán)隊(duì)首次公開(kāi)發(fā)布了他們的路線圖文檔，闡述2024年下半年的發(fā)展方向。

Meta共同創(chuàng)始人、領(lǐng)導(dǎo)PyTorch團(tuán)隊(duì)的Soumith Chintala在推特上官宣了這個(gè)消息。

他表示，希望公開(kāi)工程師們的研發(fā)動(dòng)機(jī)和目標(biāo)。

「雖然所有PyTorch開(kāi)發(fā)都在GitHub上公開(kāi)，但各個(gè)PyTorch附屬公司的團(tuán)隊(duì)編寫(xiě)的實(shí)際規(guī)劃和路線圖文檔并不公開(kāi)，因此我們決定做出改變，以提高透明度。」

PyTorch團(tuán)隊(duì)的技術(shù)項(xiàng)目經(jīng)理Gott Brath也在論壇中發(fā)表了類(lèi)似的聲明。

我們一直在考慮，如何分享團(tuán)隊(duì)在PyTorch上所做的工作的路線圖。我們每半年進(jìn)行一次規(guī)劃，因此這些是我們針對(duì)PyTorch中多個(gè)關(guān)鍵領(lǐng)域的2024年H2 OSS計(jì)劃的一些公開(kāi)版本。

這些文件基本就是PyTorch團(tuán)隊(duì)內(nèi)部的文檔和工作規(guī)劃，刪減掉了一些內(nèi)容就發(fā)布出來(lái)成為路線圖，其中涉及PyTorch的如下幾個(gè)方面：

- 核心庫(kù)與核心性能

- 分布式

- torchune、Torchrec、TorchVision

- PyTorch Edge

- 數(shù)據(jù)加載（DataLoading）

- 編譯器核心及部署

- 開(kāi)發(fā)者基礎(chǔ)設(shè)施

每個(gè)文檔都至少包含三個(gè)部分的內(nèi)容，以O(shè)KR的思路展開(kāi)：

- 背景

- Top5關(guān)注領(lǐng)域及目標(biāo)：目標(biāo)、關(guān)鍵結(jié)果、已知或未知風(fēng)險(xiǎn)以及相應(yīng)緩解措施（最多一頁(yè)）

- 提升工程水平的Top3~5個(gè)方面：BE Pillar分類(lèi)、目標(biāo)、指標(biāo)/狀態(tài)/具體目標(biāo)、已知或未知風(fēng)險(xiǎn)以及緩解措施、影響/成本、優(yōu)先級(jí)/信心程度（最多一頁(yè)）

其中BE Pillar可以看作Meta寫(xiě)給開(kāi)發(fā)團(tuán)隊(duì)的「五句箴言」，具體內(nèi)容是：

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

「最多一頁(yè)」的規(guī)定不知道有沒(méi)有戳到卷文檔長(zhǎng)度的開(kāi)發(fā)人員，畢竟文檔貴精不貴長(zhǎng)，將眾多開(kāi)發(fā)需求精簡(jiǎn)到一頁(yè)的內(nèi)容不僅節(jié)省同事時(shí)間，也十分考驗(yàn)撰寫(xiě)者的功力。

此外，文檔中也可以看出Meta開(kāi)發(fā)團(tuán)隊(duì)的一些優(yōu)秀思路，比如重視各個(gè)模塊團(tuán)隊(duì)的協(xié)作、重視和外部合作伙伴的API集成和共同開(kāi)發(fā)，重視與開(kāi)源社區(qū)和開(kāi)發(fā)者的互動(dòng)。

當(dāng)推出ExecuTorch這樣的新代碼庫(kù)，或者想要提升PyTorch編譯器影響力時(shí)，團(tuán)隊(duì)一般都會(huì)從兩方面思路入手：一是鉚足力氣提升性能，把目標(biāo)直接頂?shù)絊OTA；另一方面從深度集成入手，提供更多開(kāi)箱即用的案例。

或許，這些都是Meta多年來(lái)在開(kāi)源領(lǐng)域如魚(yú)得水、風(fēng)生水起的關(guān)鍵所在。

以下是各個(gè)文檔內(nèi)容的部分截取和概括。

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

核心庫(kù)與核心性能

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

性能方面，PyTorch團(tuán)隊(duì)提出了在模型訓(xùn)練和推理方面實(shí)現(xiàn)SOTA性能的目標(biāo)，措施包括引入架構(gòu)優(yōu)化技術(shù)和高性能kernel，與整個(gè)PyTorch技術(shù)棧形成搭配組合。

過(guò)去一年的時(shí)間見(jiàn)證了GenAI的快速發(fā)展，許多支持研究領(lǐng)域進(jìn)行開(kāi)發(fā)的外部庫(kù)應(yīng)運(yùn)而生，但其中很多并不直接依賴(lài)PyTorch，這會(huì)威脅到PyTorch在科研領(lǐng)域的主導(dǎo)地位。

為了重新跟上節(jié)奏，PyTorch將為量化、稀疏化、MoE和低精度訓(xùn)練等常用開(kāi)發(fā)技術(shù)提供支持，包括構(gòu)建模塊和API（主要集成在torchao中），幫助各類(lèi)Transformer架構(gòu)的模型提升性能。

torchao庫(kù)可以支持研究人員在PyTorch框架內(nèi)自定義高性能的dtype、layout和優(yōu)化技巧，將使用范圍擴(kuò)展到訓(xùn)練、推理、調(diào)優(yōu)等各種場(chǎng)景。

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

- 推出的自動(dòng)優(yōu)化庫(kù)torchao已經(jīng)取得了突破性的成功，下一步提升其代碼組織性，并將其中的數(shù)值運(yùn)算與核心庫(kù)分開(kāi)

- 解決TendorDict的核心模塊性，支持加載/存儲(chǔ)的序列化，并使其在eager模式下的運(yùn)行速度提高2倍

- 繼續(xù)上半年在內(nèi)存映射加載（memory mapped load）方面的成功，繼續(xù)提升模型加載/存儲(chǔ)的性能和安全性

- 將TorchRL的開(kāi)銷(xiāo)降低50%

- 加入對(duì)NoGIL的核心支持

- 修復(fù)用戶(hù)反映的TORCH_env變量不起作用的問(wèn)題

文檔中還提及了要實(shí)現(xiàn)對(duì)nn.transformer模塊的棄用，表示會(huì)發(fā)布一系列教程和用例，展示如何使用torch.compile、sdpa、NJT、FlexAttention、custom_op、torchao等模塊構(gòu)建Transformer。

分布式

LLM的預(yù)訓(xùn)練通常橫跨數(shù)十個(gè)甚至上千個(gè)GPU，而且由于模型的參數(shù)規(guī)模逐漸增大，推理和微調(diào)也很難用單個(gè)GPU完成。

因此，PyTorch下一步對(duì)「分布式」的布局全面涵蓋了訓(xùn)練、推理、微調(diào)這三個(gè)環(huán)節(jié)，提出要達(dá)成超大規(guī)模分布式訓(xùn)練、高內(nèi)存效率的微調(diào)、多主機(jī)分布式推理。

訓(xùn)練

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

- 完全分片數(shù)據(jù)并行（full sharded data parallel，F(xiàn)SDP）

- 混合分片數(shù)據(jù)并行（hybrid sharding data parallel，HSDP）

- 張量并行（tensor parallel，TP）

- 流水線并行（pipeline parallel，PP）

- 序列并行（sequence parallel，SP）

- 上下文并行（context parallel，CP）

PyTorch希望在TorchTitan中將各種并行方式進(jìn)一步模塊化，讓開(kāi)發(fā)者可以自由組合，根據(jù)需要實(shí)現(xiàn)N維并行。

文檔中特別提到，對(duì)MoE和多模態(tài)這兩種新興的架構(gòu)需要添加支持，比如專(zhuān)家并行、路由算法的優(yōu)化。

除了TorchTitan本身的更新，分布式團(tuán)隊(duì)還需要與編譯器團(tuán)隊(duì)進(jìn)一步緊密合作，更好地與torch.compile模塊集成，為大規(guī)模分布式場(chǎng)景帶來(lái)額外的性能提升。

微調(diào)與推理

微調(diào)：聯(lián)合torchtune，將FSDP2 LoRA/QLoRA方案投入使用，以及支持模型狀態(tài)字典的NF4量化

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

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

torchtune、TorchRec、TorchVision

torchtune

torchtune的推出旨在幫助用戶(hù)更方便微調(diào)LLM，這也是官方給出的Llama模型微調(diào)的方案。

torchtune定義的「微調(diào)」范圍非常廣，主要可以概括為三類(lèi)場(chǎng)景：

- 對(duì)特定領(lǐng)域數(shù)據(jù)集或者下游任務(wù)的模型適應(yīng)

- 獎(jiǎng)勵(lì)和偏好建模，比如RLHF、DPO等

- 包含蒸餾與量化的訓(xùn)練過(guò)程

下半年的更新將支持為agent工作流進(jìn)行的微調(diào)，同時(shí)著重關(guān)注微調(diào)性能的提升。

團(tuán)隊(duì)會(huì)與compile、core、distributed等模塊進(jìn)行合作，提供高效率微調(diào)，并在PyTorch生態(tài)內(nèi)建立有代表性的微調(diào)性能基準(zhǔn)。

由于torchtune也是一個(gè)較新的開(kāi)源庫(kù)，因此與開(kāi)源社區(qū)的互動(dòng)也必不可少。

文檔提出發(fā)布博客文章和教程、舉辦技術(shù)講座等方式，提升用戶(hù)的理解；并會(huì)定義量化指標(biāo)，衡量torchturn在LLM生態(tài)中的貢獻(xiàn)份額。

除了開(kāi)源社區(qū)，torchtune還會(huì)與至少一個(gè)合作伙伴集成，參與到它們的社區(qū)中，以促進(jìn)torchtune的使用。

TorchVision

TorchVision作為CV領(lǐng)域內(nèi)的絕對(duì)主宰者，技術(shù)也相對(duì)成熟，因此路線圖中提出的更新很少。

團(tuán)隊(duì)將繼續(xù)在預(yù)處理方向努力，在圖像編碼/解碼空間中支持更多格式（如WebP、HEIC）和平臺(tái)（如CUDA），并提升jpeg格式在GPU上的編碼/解碼性能。

TorchRec

TorchRec旨在提供大規(guī)模推薦系統(tǒng)中常用的稀疏性和并行性原語(yǔ)，將秋季推出第一個(gè)穩(wěn)定版本TorchRec 1.0。

Edge

目前，開(kāi)源庫(kù)ExecuTorch已經(jīng)推出了Alpha版本，主要依賴(lài)torch.compile和torch.export，用于支持移動(dòng)設(shè)備和邊緣設(shè)備（如AR/VR、可穿戴設(shè)備）上的模型分析、調(diào)試和推理。

下半年，Edge團(tuán)隊(duì)將推出xecuTorch的Beta版本，同時(shí)為Meta的Llama系列模型和其他開(kāi)源模型提供PyTorch生態(tài)內(nèi)的解決方案。

關(guān)鍵目標(biāo)中主要涵蓋兩個(gè)方向。一是為設(shè)備上AI提供基礎(chǔ)功能和可靠基礎(chǔ)設(shè)施，包括：

- 確保C++和Python的API穩(wěn)定性

- 實(shí)現(xiàn)一系列核心功能：支持模型壓縮、代理緩存位置管理、數(shù)據(jù)和程序分離

二是為這個(gè)新生的代碼庫(kù)保駕護(hù)航，培育開(kāi)源社區(qū)內(nèi)的影響力，同時(shí)與Arm、Apple 和Qualcomm等公司保持良好合作關(guān)系。

其中社區(qū)影響力的目標(biāo)甚至被量化到，要求代碼在GitHub上得到3k標(biāo)星，500次克?。╢ork）。有興趣的吃瓜群眾可以去持續(xù)關(guān)注一下，看看團(tuán)隊(duì)能不能在年底完成這個(gè)OKR。

數(shù)據(jù)加載

基于Apache Arrow格式的HuggingFace datasets庫(kù)憑借無(wú)內(nèi)存限制的高速加載/存儲(chǔ)，近幾年異軍突起，似乎搶走了PyTorch相關(guān)功能的風(fēng)頭。

數(shù)據(jù)加載的文檔中開(kāi)篇就提出了雄心壯志，要讓TorchData庫(kù)再次偉大，重新確立PyTorch在數(shù)據(jù)加載方面的主宰地位。

要達(dá)到這個(gè)目標(biāo)，就需要讓相關(guān)功能變得靈活、可擴(kuò)展、高性能、高內(nèi)存效率，同時(shí)實(shí)現(xiàn)傻瓜式操作，支持各種規(guī)模的多模態(tài)訓(xùn)練。

具體的更新目標(biāo)包括以下幾個(gè)方面：

- DataLoader的功能開(kāi)發(fā)和接口都將貫徹GitHub優(yōu)先的原則，DataPipes和DataLoader v2則將被逐步被棄用、刪除

- 確保TorchTune、TorchTitan、HuggingFace、TorchData之間的清晰邊界和良好互通性，支持多數(shù)據(jù)集、多模態(tài)數(shù)據(jù)加載

- HuggingFace使用StatefulDataLoader的API，確保兼容性，及時(shí)更新樣例和測(cè)試用例

編譯器核心及部署

PyTorch的編譯器核心功能經(jīng)過(guò)多年發(fā)展已經(jīng)趨于完善，目前亟待彌補(bǔ)的只是對(duì)LLM和GenAI領(lǐng)域的更深度集成和更多優(yōu)化支持。

路線圖提出，要將torch.compile()函數(shù)帶到LLM和GenAI的使用周期的各個(gè)方面（推理、微調(diào)、預(yù)訓(xùn)練），讓重要模型在發(fā)行時(shí)就搭載原生的PyTorch編譯。

為了實(shí)現(xiàn)這個(gè)目標(biāo)，文檔提出了很多具體措施，比如與torchtune與TorchTitan團(tuán)隊(duì)合作，提升編譯性能，并在下半年發(fā)布至少兩個(gè)高知名度模型的原生PyTorch編譯版本。

此外，編譯器可能添加可視化功能，在non-eager訓(xùn)練模式下生成表達(dá)前向計(jì)算/后向傳播過(guò)程的模型圖。

用戶(hù)支持方面也有諸多規(guī)劃，比如提升系統(tǒng)的監(jiān)控性和可觀察性，幫助戶(hù)自行調(diào)試編譯問(wèn)題。關(guān)鍵目標(biāo)還包括建立用戶(hù)支持團(tuán)隊(duì)，針對(duì)幾個(gè)關(guān)鍵領(lǐng)域（數(shù)據(jù)類(lèi)、上下文管理等），解決開(kāi)發(fā)者在GitHub等平臺(tái)上發(fā)布的問(wèn)題。