國產GPU制造商摩爾線程近期宣布,其面向PyTorch深度學習框架的MUSA擴展庫——Torch-MUSA,已成功升級到v2.0.0版本。這一重大更新標志著摩爾線程在加速AI計算領域邁出了重要一步。
通過Torch-MUSA v2.0.0,用戶只需簡單地將設備指定為torch.device("musa"),便能輕松地將現有的PyTorch模型遷移至MUSA架構的GPU上運行,無需對代碼進行大規模修改。更令人振奮的是,Torch-MUSA已全面開源,開發者可以通過GitHub輕松獲取其源代碼。
此次升級的核心亮點在于,Torch-MUSA v2.0.0在國產GPU上首次實現了對FP8數據類型的全面支持。FP8作為一種低精度格式,在支持該格式的GPU上,大語言模型的訓練采用FP8混合精度可以顯著提升GPU的算力,并大幅度降低顯存的占用。
摩爾線程的新一代MUSA Compute Capability 3.1計算架構全功能GPU原生支持FP8計算,這為Torch-MUSA v2.0.0實現FP8矩陣乘法和分布式通信優化提供了堅實的基礎。借助這一底層架構的優勢,Torch-MUSA v2.0.0能夠充分發揮FP8的計算效能,從而顯著提高大語言模型的訓練和推理效率。
除了對FP8的支持外,Torch-MUSA v2.0.0還在MUSA計算平臺上引入了多項創新功能,旨在進一步提升深度學習任務的執行效率。其中,新增的虛擬內存管理技術能夠有效解決GPU內存碎片化問題,降低模型訓練過程中的峰值內存占用,尤其適用于FSDP、DeepSpeed和Megatron-LM等主流的大模型訓練框架。
Torch-MUSA v2.0.0還引入了MUSA Graph技術,該技術將多個MUSA內核整合到一個圖中,通過減少CPU的調度次數來降低啟動開銷,從而提升計算效率。同時,該技術還與CUDA Graph接口實現了高效兼容。
為了進一步提升性能,Torch-MUSA v2.0.0還為torch.compile增加了Triton后端支持。這意味著開發者可以直接使用PyTorch的原生接口,從而獲得更高效的性能表現。
Torch-MUSA v2.0.0在全面支持PyTorch 2.2.0的基礎上,還新增了對PyTorch 2.5.0的支持。這使得開發者能夠在基于MUSA Compute Capability 3.1計算架構的全功能GPU上無縫運行新版本的PyTorch,從而享受更多的功能和性能優化。
隨著AI技術的不斷發展,摩爾線程表示,Torch-MUSA未來將繼續跟進PyTorch的版本更新,計劃支持更高版本的PyTorch,以滿足開發者對性能和功能不斷提升的需求。
Torch-MUSA v2.0.0的發布,不僅展示了摩爾線程在GPU技術研發方面的實力,也為AI開發者提供了更加高效、靈活的深度學習解決方案。
隨著摩爾線程在GPU領域的不斷探索和創新,我們有理由相信,未來的Torch-MUSA將為AI計算帶來更多的驚喜和突破。