PyTorch團(tuán)隊(duì)首發(fā)技術(shù)路線圖，近百頁文檔披露2024下半年發(fā)展方向
如果你在AI領(lǐng)域用Python開發(fā)，想必PyTorch一定是你的老朋友之一。2017年，Meta AI發(fā)布了這個機(jī)器學(xué)習(xí)和深度學(xué)習(xí)領(lǐng)域的開源庫，如今已經(jīng)走到了第7個年頭。
根據(jù)Assembly AI 2021年的統(tǒng)計(jì)數(shù)據(jù)，HuggingFace上最受歡迎的top 30模型都能在PyTorch上運(yùn)行，有92%的模型是PyTorch專有的，這個占比讓包括TensorFlow在內(nèi)的一眾競爭對手都望塵莫及。
就在7月10日，PyTorch的工程團(tuán)隊(duì)首次公開發(fā)布了他們的路線圖文檔，闡述2024年下半年的發(fā)展方向。
Meta共同創(chuàng)始人、領(lǐng)導(dǎo)PyTorch團(tuán)隊(duì)的Soumith Chintala在推特上官宣了這個消息。
他表示，希望公開工程師們的研發(fā)動機(jī)和目標(biāo)。
「雖然所有PyTorch開發(fā)都在GitHub上公開，但各個PyTorch附屬公司的團(tuán)隊(duì)編寫的實(shí)際規(guī)劃和路線圖文檔并不公開，因此我們決定做出改變，以提高透明度。」
PyTorch團(tuán)隊(duì)的技術(shù)項(xiàng)目經(jīng)理Gott Brath也在論壇中發(fā)表了類似的聲明。
我們一直在考慮，如何分享團(tuán)隊(duì)在PyTorch上所做的工作的路線圖。我們每半年進(jìn)行一次規(guī)劃，因此這些是我們針對PyTorch中多個關(guān)鍵領(lǐng)域的2024年H2 OSS計(jì)劃的一些公開版本。
這些文件基本就是PyTorch團(tuán)隊(duì)內(nèi)部的文檔和工作規(guī)劃，刪減掉了一些內(nèi)容就發(fā)布出來成為路線圖，其中涉及PyTorch的如下幾個方面：
- 核心庫與核心性能
- 分布式
- torchune、Torchrec、TorchVision
- PyTorch Edge
- 數(shù)據(jù)加載（DataLoading）
- 編譯器核心及部署
- 開發(fā)者基礎(chǔ)設(shè)施
每個文檔都至少包含三個部分的內(nèi)容，以O(shè)KR的思路展開：
- 背景
- Top5關(guān)注領(lǐng)域及目標(biāo)：目標(biāo)、關(guān)鍵結(jié)果、已知或未知風(fēng)險以及相應(yīng)緩解措施（最多一頁）
- 提升工程水平的Top3~5個方面：BE Pillar分類、目標(biāo)、指標(biāo)/狀態(tài)/具體目標(biāo)、已知或未知風(fēng)險以及緩解措施、影響/成本、優(yōu)先級/信心程度（最多一頁）
其中BE Pillar可以看作Meta寫給開發(fā)團(tuán)隊(duì)的「五句箴言」，具體內(nèi)容是：
Better Code, Better Doc, Empowering teams, Modern Code, Better Architecture
「最多一頁」的規(guī)定不知道有沒有戳到卷文檔長度的開發(fā)人員，畢竟文檔貴精不貴長，將眾多開發(fā)需求精簡到一頁的內(nèi)容不僅節(jié)省同事時間，也十分考驗(yàn)撰寫者的功力。
此外，文檔中也可以看出Meta開發(fā)團(tuán)隊(duì)的一些優(yōu)秀思路，比如重視各個模塊團(tuán)隊(duì)的協(xié)作、重視和外部合作伙伴的API集成和共同開發(fā)，重視與開源社區(qū)和開發(fā)者的互動。
當(dāng)推出ExecuTorch這樣的新代碼庫，或者想要提升PyTorch編譯器影響力時，團(tuán)隊(duì)一般都會從兩方面思路入手：一是鉚足力氣提升性能，把目標(biāo)直接頂?shù)絊OTA；另一方面從深度集成入手，提供更多開箱即用的案例。
或許，這些都是Meta多年來在開源領(lǐng)域如魚得水、風(fēng)生水起的關(guān)鍵所在。
以下是各個文檔內(nèi)容的部分截取和概括。
原文地址：https://dev-discuss.pytorch.org/t/meta-pytorch-team-2024-h2-roadmaps/2226
核心庫與核心性能
文檔中涉及到的核心庫包括TendorDict、torchao、NN、TorchRL等。
性能方面，PyTorch團(tuán)隊(duì)提出了在模型訓(xùn)練和推理方面實(shí)現(xiàn)SOTA性能的目標(biāo)，措施包括引入架構(gòu)優(yōu)化技術(shù)和高性能kernel，與整個PyTorch技術(shù)棧形成搭配組合。
過去一年的時間見證了GenAI的快速發(fā)展，許多支持研究領(lǐng)域進(jìn)行開發(fā)的外部庫應(yīng)運(yùn)而生，但其中很多并不直接依賴PyTorch，這會威脅到PyTorch在科研領(lǐng)域的主導(dǎo)地位。
為了重新跟上節(jié)奏，PyTorch將為量化、稀疏化、MoE和低精度訓(xùn)練等常用開發(fā)技術(shù)提供支持，包括構(gòu)建模塊和API（主要集成在torchao中），幫助各類Transformer架構(gòu)的模型提升性能。
torchao庫可以支持研究人員在PyTorch框架內(nèi)自定義高性能的dtype、layout和優(yōu)化技巧，將使用范圍擴(kuò)展到訓(xùn)練、推理、調(diào)優(yōu)等各種場景。
此外，核心庫的更新將包括以下方面：
- 推出的自動優(yōu)化庫torchao已經(jīng)取得了突破性的成功，下一步提升其代碼組織性，并將其中的數(shù)值運(yùn)算與核心庫分開
- 解決TendorDict的核心模塊性，支持加載/存儲的序列化，并使其在eager模式下的運(yùn)行速度提高2倍
- 繼續(xù)上半年在內(nèi)存映射加載（memory mapped load）方面的成功，繼續(xù)提升模型加載/存儲的性能和安全性
- 將TorchRL的開銷降低50%
- 加入對NoGIL的核心支持
- 修復(fù)用戶反映的TORCH_env變量不起作用的問題
文檔中還提及了要實(shí)現(xiàn)對nn.transformer模塊的棄用，表示會發(fā)布一系列教程和用例，展示如何使用torch.compile、sdpa、NJT、FlexAttention、custom_op、torchao等模塊構(gòu)建Transformer。
分布式
LLM的預(yù)訓(xùn)練通常橫跨數(shù)十個甚至上千個GPU，而且由于模型的參數(shù)規(guī)模逐漸增大，推理和微調(diào)也很難用單個GPU完成。
因此，PyTorch下一步對「分布式」的布局全面涵蓋了訓(xùn)練、推理、微調(diào)這三個環(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)一步模塊化，讓開發(fā)者可以自由組合，根據(jù)需要實(shí)現(xiàn)N維并行。
文檔中特別提到，對MoE和多模態(tài)這兩種新興的架構(gòu)需要添加支持，比如專家并行、路由算法的優(yōu)化。
除了TorchTitan本身的更新，分布式團(tuán)隊(duì)還需要與編譯器團(tuán)隊(duì)進(jìn)一步緊密合作，更好地與torch.compile模塊集成，為大規(guī)模分布式場景帶來額外的性能提升。
微調(diào)與推理
微調(diào)：聯(lián)合torchtune，將FSDP2 LoRA/QLoRA方案投入使用，以及支持模型狀態(tài)字典的NF4量化
推理：PP和DP已經(jīng)成為分布式API的核心，下一步需要關(guān)注torchtitan的分布式推理，支持大模型PP+異步TP方式，將給出案例展示
文檔中還提到，會將HuggingFace的推理API從PiPPy遷移到PyTorch（由HuggingFace完成）。
torchtune、TorchRec、TorchVision
torchtune
torchtune的推出旨在幫助用戶更方便微調(diào)LLM，這也是官方給出的Llama模型微調(diào)的方案。
torchtune定義的「微調(diào)」范圍非常廣，主要可以概括為三類場景：
- 對特定領(lǐng)域數(shù)據(jù)集或者下游任務(wù)的模型適應(yīng)
- 獎勵和偏好建模，比如RLHF、DPO等
- 包含蒸餾與量化的訓(xùn)練過程
下半年的更新將支持為agent工作流進(jìn)行的微調(diào)，同時著重關(guān)注微調(diào)性能的提升。
團(tuán)隊(duì)會與compile、core、distributed等模塊進(jìn)行合作，提供高效率微調(diào)，并在PyTorch生態(tài)內(nèi)建立有代表性的微調(diào)性能基準(zhǔn)。
由于torchtune也是一個較新的開源庫，因此與開源社區(qū)的互動也必不可少。
文檔提出發(fā)布博客文章和教程、舉辦技術(shù)講座等方式，提升用戶的理解；并會定義量化指標(biāo)，衡量torchturn在LLM生態(tài)中的貢獻(xiàn)份額。
除了開源社區(qū)，torchtune還會與至少一個合作伙伴集成，參與到它們的社區(qū)中，以促進(jìn)torchtune的使用。
TorchVision
TorchVision作為CV領(lǐng)域內(nèi)的絕對主宰者，技術(shù)也相對成熟，因此路線圖中提出的更新很少。
團(tuán)隊(duì)將繼續(xù)在預(yù)處理方向努力，在圖像編碼/解碼空間中支持更多格式（如WebP、HEIC）和平臺（如CUDA），并提升jpeg格式在GPU上的編碼/解碼性能。
TorchRec
TorchRec旨在提供大規(guī)模推薦系統(tǒng)中常用的稀疏性和并行性原語，將秋季推出第一個穩(wěn)定版本TorchRec 1.0。
Edge
目前，開源庫ExecuTorch已經(jīng)推出了Alpha版本，主要依賴torch.compile和torch.export，用于支持移動設(shè)備和邊緣設(shè)備（如AR/VR、可穿戴設(shè)備）上的模型分析、調(diào)試和推理。
下半年，Edge團(tuán)隊(duì)將推出xecuTorch的Beta版本，同時為Meta的Llama系列模型和其他開源模型提供PyTorch生態(tài)內(nèi)的解決方案。
關(guān)鍵目標(biāo)中主要涵蓋兩個方向。一是為設(shè)備上AI提供基礎(chǔ)功能和可靠基礎(chǔ)設(shè)施，包括：
- 確保C++和Python的API穩(wěn)定性
- 實(shí)現(xiàn)一系列核心功能：支持模型壓縮、代理緩存位置管理、數(shù)據(jù)和程序分離
二是為這個新生的代碼庫保駕護(hù)航，培育開源社區(qū)內(nèi)的影響力，同時與Arm、Apple 和Qualcomm等公司保持良好合作關(guān)系。
其中社區(qū)影響力的目標(biāo)甚至被量化到，要求代碼在GitHub上得到3k標(biāo)星，500次克?。╢ork）。有興趣的吃瓜群眾可以去持續(xù)關(guān)注一下，看看團(tuán)隊(duì)能不能在年底完成這個OKR。
數(shù)據(jù)加載
基于Apache Arrow格式的HuggingFace datasets庫憑借無內(nèi)存限制的高速加載/存儲，近幾年異軍突起，似乎搶走了PyTorch相關(guān)功能的風(fēng)頭。
數(shù)據(jù)加載的文檔中開篇就提出了雄心壯志，要讓TorchData庫再次偉大，重新確立PyTorch在數(shù)據(jù)加載方面的主宰地位。
要達(dá)到這個目標(biāo)，就需要讓相關(guān)功能變得靈活、可擴(kuò)展、高性能、高內(nèi)存效率，同時實(shí)現(xiàn)傻瓜式操作，支持各種規(guī)模的多模態(tài)訓(xùn)練。
具體的更新目標(biāo)包括以下幾個方面：
- DataLoader的功能開發(fā)和接口都將貫徹GitHub優(yōu)先的原則，DataPipes和DataLoader v2則將被逐步被棄用、刪除
- 確保TorchTune、TorchTitan、HuggingFace、TorchData之間的清晰邊界和良好互通性，支持多數(shù)據(jù)集、多模態(tài)數(shù)據(jù)加載
- HuggingFace使用StatefulDataLoader的API，確保兼容性，及時更新樣例和測試用例
編譯器核心及部署
PyTorch的編譯器核心功能經(jīng)過多年發(fā)展已經(jīng)趨于完善，目前亟待彌補(bǔ)的只是對LLM和GenAI領(lǐng)域的更深度集成和更多優(yōu)化支持。
路線圖提出，要將torch.compile()函數(shù)帶到LLM和GenAI的使用周期的各個方面（推理、微調(diào)、預(yù)訓(xùn)練），讓重要模型在發(fā)行時就搭載原生的PyTorch編譯。
為了實(shí)現(xiàn)這個目標(biāo)，文檔提出了很多具體措施，比如與torchtune與TorchTitan團(tuán)隊(duì)合作，提升編譯性能，并在下半年發(fā)布至少兩個高知名度模型的原生PyTorch編譯版本。
此外，編譯器可能添加可視化功能，在non-eager訓(xùn)練模式下生成表達(dá)前向計(jì)算/后向傳播過程的模型圖。
用戶支持方面也有諸多規(guī)劃，比如提升系統(tǒng)的監(jiān)控性和可觀察性，幫助戶自行調(diào)試編譯問題。關(guān)鍵目標(biāo)還包括建立用戶支持團(tuán)隊(duì)，針對幾個關(guān)鍵領(lǐng)域（數(shù)據(jù)類、上下文管理等），解決開發(fā)者在GitHub等平臺上發(fā)布的問題。
