近日,摩爾線程正式發(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)。