摩爾線程公司近期正式揭曉了其面向PyTorch深度學習框架的MUSA擴展庫新版本——Torch-MUSA v2.0.0。此次發布標志著摩爾線程在AI計算領域邁出了重要一步。
該新版本基于MUSA Compute Capability 3.1計算架構,帶來了多項突破性升級。其中,最引人注目的是原生支持FP8數據類型。FP8作為一種前沿的低精度格式,對于大語言模型(LLM)的訓練至關重要。在支持FP8的GPU上,采用FP8混合精度可以顯著提升GPU的算力,并大幅降低顯存占用。
摩爾線程的全功能GPU,憑借新一代MUSA Compute Capability 3.1計算架構,原生支持FP8計算。這一底層架構的優勢為Torch-MUSA v2.0.0提供了堅實的基礎,使其能夠充分發揮FP8的計算效能,在大語言模型的訓練和推理中展現出更高的效率。
除了對FP8的支持外,Torch-MUSA v2.0.0還引入了多項創新功能,以進一步提升深度學習任務的執行效率。其中包括新增的虛擬內存管理支持,該技術能夠有效緩解GPU內存碎片化問題,并降低模型訓練過程中的峰值內存占用。這對于FSDP、DeepSpeed和Megatron-LM等主流大模型訓練框架來說,無疑是一個巨大的福音。
Torch-MUSA v2.0.0還新增了MUSA Graph支持。MUSA Graph技術將多個MUSA內核整合到一個圖中,通過單次CPU調度大幅減少了啟動開銷,提升了計算效率。同時,該技術還與CUDA Graph接口高效兼容,為開發者提供了更多的選擇和便利。
在性能優化方面,Torch-MUSA v2.0.0也取得了顯著進展。通過多項針對MUSA計算平臺的性能優化,該版本進一步提升了對AI模型和大規模數據處理的支持能力。這使得開發者能夠在基于MUSA Compute Capability 3.1計算架構的全功能GPU上,無縫運行新版本的PyTorch,并享受更高效的性能表現。
Torch-MUSA v2.0.0不僅支持PyTorch 2.2.0,還新增了對PyTorch 2.5.0的支持。這意味著開發者可以更加靈活地選擇適合自己項目的PyTorch版本,無需擔心兼容性問題。
最后,Torch-MUSA v2.0.0已完全開源,開發者可以通過訪問GitHub獲取源代碼。這一舉措無疑將吸引更多的開發者參與到Torch-MUSA的生態建設中來,共同推動AI計算領域的發展。