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