近日，摩爾線程正式發(fā)布Torch-MUSA v2.0.0版本，這是其面向PyTorch深度學(xué)習(xí)框架的MUSA擴(kuò)展庫的重要升級。新版本基于MUSA Compute Capability 3.1計(jì)算架構(gòu)，支持原生FP8數(shù)據(jù)類型，支持PyTorch 2.5.0，并通過多項(xiàng)針對MUSA計(jì)算平臺的性能優(yōu)化，進(jìn)一步提升了對AI模型和大規(guī)模數(shù)據(jù)處理的支持能力。
FP8原生支持，國產(chǎn)GPU的技術(shù)突破
作為本次升級的核心亮點(diǎn)，Torch-MUSA v2.0.0率先在國產(chǎn)GPU上實(shí)現(xiàn)了對FP8數(shù)據(jù)類型的完整支持。FP8是當(dāng)前AI計(jì)算的一種前沿低精度格式，在支持原生FP8的GPU上，大語言模型（LLM）訓(xùn)練采用FP8混合精度可大幅提高GPU算力，顯著降低顯存占用。摩爾線程基于新一代 MUSA Compute Capability 3.1 計(jì)算架構(gòu)的全功能 GPU 原生支持 FP8 計(jì)算，這為 Torch-MUSA v2.0.0 實(shí)現(xiàn) FP8 矩陣乘法和分布式通信優(yōu)化提供了堅(jiān)實(shí)的基礎(chǔ)。依托這一底層架構(gòu)優(yōu)勢，Torch-MUSA v2.0.0 能夠充分發(fā)揮 FP8 的計(jì)算效能，顯著提升大語言模型訓(xùn)練和推理的效率。
三大關(guān)鍵優(yōu)化，提升AI計(jì)算效率
Torch-MUSA v2.0.0在MUSA計(jì)算平臺引入多項(xiàng)創(chuàng)新功能，進(jìn)一步提升深度學(xué)習(xí)任務(wù)的執(zhí)行效率：
▼ 新增虛擬內(nèi)存管理支持：
MUSA虛擬內(nèi)存管理技術(shù)能夠有效緩解GPU內(nèi)存碎片化問題，降低模型訓(xùn)練過程中的峰值內(nèi)存占用，特別適用于FSDP、DeepSpeed和Megatron-LM等主流大模型訓(xùn)練框架。
▼ 新增MUSA Graph支持：
MUSA Graph技術(shù)將多個MUSA內(nèi)核整合到一個圖中，通過單次CPU調(diào)度大幅減少啟動開銷，提升計(jì)算效率，同時與CUDA Graph接口高效兼容。
▼ torch.compile增加Triton后端支持：
為torch.compile提供了Triton-MUSA后端支持，開發(fā)者可以直接使用PyTorch原生接口，獲得更高效的性能表現(xiàn)。
支持PyTorch 2.5.0，生態(tài)兼容性更強(qiáng)
Torch-MUSA v2.0.0在完整支持PyTorch 2.2.0的基礎(chǔ)上，新增了對PyTorch 2.5.0的支持，使開發(fā)者能夠在基于MUSA Compute Capability 3.1計(jì)算架構(gòu)的全功能GPU上，無縫運(yùn)行新版本的PyTorch，享受更高效的AI計(jì)算體驗(yàn)。
Torch-MUSA已完全開源，開發(fā)者可通過訪問GitHub獲取源代碼。摩爾線程鼓勵開發(fā)者積極參與該項(xiàng)目的開發(fā)與改進(jìn)，通過提交問題報(bào)告（issue）或代碼修改申請（pull request）等方式，共同推動Torch-MUSA以及MUSA軟件生態(tài)的持續(xù)進(jìn)步與創(chuàng)新。
▼ Torch-MUSA開源地址：
https://github.com/MooreThreads/torch_musa
▼ 功能特性：
在Torch-MUSA中，用戶只需指定torch.device("musa")，即可輕松將現(xiàn)有的PyTorch模型遷移到MUSA架構(gòu)的GPU上運(yùn)行，無需大幅修改代碼。Torch-MUSA完全兼容PyTorch的自動微分和動態(tài)圖機(jī)制，支持多種常用的神經(jīng)網(wǎng)絡(luò)模塊及優(yōu)化算法，并加速了關(guān)鍵深度學(xué)習(xí)算子的計(jì)算。此外，Torch-MUSA還支持多種PyTorch特性，包括DDP、JIT、FSDP、Profiler、Extension等。
▼ 版本迭代：
MUSA Graph技術(shù)將多個MUSA內(nèi)核整合到一個圖中，通過單次CPU調(diào)度大幅減少啟動開銷，提升計(jì)算效率，同時與CUDA Graph接口高效兼容。
v1.1.0：初次發(fā)布，支持PyTorch 2.0，提供基礎(chǔ)張量操作和常見神經(jīng)網(wǎng)絡(luò)層的MUSA加速。
v1.2.0：進(jìn)一步擴(kuò)展算子支持，支持了完整功能的Profiler、MUSA Extension，并增加了Torch-MUSA專有特性如compare_tool、musa_converter，幫助用戶更快的定位模型精度問題。
v1.3.0：支持PyTorch 2.2.0，性能進(jìn)一步提升，支持FSDP，支持更復(fù)雜的模型和更大規(guī)模的數(shù)據(jù)處理。
v2.0.0：在MUSA Compute Capability 3.1計(jì)算架構(gòu)上，原生支持FP8數(shù)據(jù)類型，支持PyTorch 2.5.0，新增MUSA虛擬內(nèi)存管理技術(shù)優(yōu)化大模型訓(xùn)練內(nèi)存效率，通過集成Triton-MUSA后端顯著提升torch.compile編譯效率，支持MUSA Graph技術(shù)。
▼ 未來計(jì)劃：
Torch-MUSA將繼續(xù)跟進(jìn)PyTorch的版本更新，計(jì)劃支持更高版本的PyTorch。摩爾線程期待與廣大開發(fā)者和研究人員共同完善Torch-MUSA的功能，持續(xù)優(yōu)化性能，為基于MUSA架構(gòu)的國產(chǎn)全功能GPU構(gòu)建更強(qiáng)大的深度學(xué)習(xí)生態(tài)。
