今日(ri),寒武紀(ji)開(kai)源(yuan)了PyTorch設備后端擴展插件Torch-MLU,并實現了寒武紀(ji)硬(ying)件對(dui)于(yu)PyTorch的原生支(zhi)持,充分提升了開(kai)發(fa)者(zhe)的使用體(ti)驗(yan)和集成效率。
早在2018年,寒武紀就開始了與PyTorch框架的集成適配工作。自PyTorch 1.3版本開始就支持寒武紀MLU系列智能加速(su)卡作為PyTorch的加速后端,使開發者能夠將原(yuan)先基于GPU的深度學習網絡高效遷移(yi)到(dao)基于寒武紀MLU的PyTorch環(huan)境(jing)中。
隨著(zhu)PyTorch社區(qu)發布PyTorch 2.4版本(ben),配合新(xin)版本(ben)(≥ V1.22)的(de)Torch-MLU插件(jian),寒(han)武紀(ji)硬件(jian)實現了(le)對(dui)PyTorch的(de)原生支持。開發者無需再安裝(zhuang)寒(han)武紀(ji)定制的(de)Cambricon PyTorch,而是可以(yi)基于社區(qu)的(de)PyTorch安裝(zhuang)Torch-MLU插件(jian),就可將(jiang)寒(han)武紀(ji)硬件(jian)的(de)算力接入原生PyTorch。
與此同時(shi),Torch-MLU的(de)(de)(de)開(kai)(kai)源使寒(han)武紀能夠更(geng)快地響應(ying)開(kai)(kai)發(fa)者的(de)(de)(de)反饋和問題,搭建(jian)了高效的(de)(de)(de)溝(gou)通橋(qiao)梁,促進了與全(quan)球開(kai)(kai)發(fa)者的(de)(de)(de)協作和知識(shi)共(gong)(gong)享(xiang)。寒(han)武紀通過持續貢獻開(kai)(kai)源社區的(de)(de)(de)實際行動,展示著(zhu)擁(yong)抱開(kai)(kai)源、支(zhi)持AI生(sheng)態共(gong)(gong)建(jian)的(de)(de)(de)決心。
↓GitEE倉庫(點擊文(wen)末“閱讀原(yuan)文(wen)”可直(zhi)接跳(tiao)轉(zhuan))↓
//gitee.com/cambricon/torch_mlu
*具體版本與分支信息請參考倉庫中的README文檔。
一、?Torch-MLU的技術演進
1、設備后端接入PyTorch框架
PyTorch初期沒有(you)提供第三方(fang)后(hou)端統(tong)一接入方(fang)案。新設備(bei)后(hou)端如果想(xiang)接入PyTorch,往往需要對(dui)PyTorch的核心(xin)(xin)代碼進行修改。比如開發者(zhe)需要在核心(xin)(xin)模塊Dispatcher中加入對(dui)應后(hou)端的Dispatch Key,以及對(dui)其他(ta)模塊中關于(yu)設備(bei)后(hou)端做一些(xie)適配性修改。設備(bei)后(hou)端擴(kuo)展(zhan)插件需要配合修改過的PyTorch才能夠正常工作(zuo)。

▲侵入式設(she)備后(hou)端(duan)支持
2023年8月PyTorch社區發布了PyTorch 2.1版本,重磅推出了基(ji)于PrivateUse1 Dispatch Key的新后端統(tong)一(yi)接(jie)入(ru)(ru)參考設(she)計(ji)(ji)。PrivateUse1設(she)計(ji)(ji)方案通(tong)過PrivateUse1抽象層將設(she)備廠商私有屬性統(tong)一(yi)封裝(zhuang),完善了PyTorch的設(she)備擴展(zhan)機制,為(wei)新硬件非侵(qin)入(ru)(ru)性接(jie)入(ru)(ru)PyTorch提供了途徑。

▲設備后端原生支持
通(tong)過這些改進(jin),PyTorch不(bu)僅增強了(le)其(qi)作為深度學習框(kuang)架的靈活性和可擴展性,還為開發者和設備廠商提(ti)供了(le)更順暢和高(gao)效的合作環境。
2、寒武紀持續貢獻PyTorch社區
PrivateUse1接入方案的(de)(de)推(tui)出為新后端統一接入提(ti)供了啟(qi)發性的(de)(de)思路,但是在初(chu)期仍有一些待完善的(de)(de)特性和模塊,需(xu)要(yao)設備廠商的(de)(de)持續投入和貢獻。
PrivateUse1支(zhi)持(chi)設(she)備在(zai)Eager模(mo)式(shi)下(xia)的基(ji)本的運行和(he)接入,但并不(bu)能完善支(zhi)持(chi)PyTorch的一(yi)些特性。比如Profiler模(mo)塊,初期(qi)僅(jin)支(zhi)持(chi)GPU后端(duan),不(bu)支(zhi)持(chi)外(wai)部自定(ding)義后端(duan),基(ji)于(yu)PrivateUse1注冊(ce)的設(she)備后端(duan)無法(fa)使用PyTorch原(yuan)生的Profiler。
PrivateUse1路徑的限制給開發者帶來了更(geng)多挑戰,使(shi)(shi)他們在進(jin)行設備(bei)集(ji)成時(shi)面臨更(geng)多障礙,影響(xiang)了整體使(shi)(shi)用體驗。
寒武紀在PrivateUse1接入方案的基礎上優化了非CUDA設備的接入體驗,在今年向PyTorch社區提交了數十個Patch,涉及Profiler、Compile、Graph Capture、Autograd、Allocator、Storage, FSDP、Sparse等眾多模塊(kuai),打通了這些(xie)模塊(kuai)與(yu)PrivateUse1的集(ji)成(cheng)路(lu)徑,進(jin)一步(bu)完善了PrivateUse1機制。未來,第(di)三方設備的廠(chang)商(shang)和(he)開發者就能夠充分利用這些(xie)已打通的集(ji)成(cheng)路(lu)徑,更(geng)加(jia)高效和(he)便捷(jie)地接入PyTorch。
3、安裝Torch-MLU插件,保留原生體驗
經過寒武(wu)紀對PyTorch社區的PrivateUse1機制(zhi)進行(xing)完(wan)善(shan)后,自PyTorch 2.4版本起(qi),寒武(wu)紀硬件實現了(le)PyTorch原生(sheng)支持。
開發者可(ke)直(zhi)接基(ji)于社區PyTorch發布(bu)包安裝Torch-MLU插件,即可(ke)將寒武紀MLU的算(suan)力(li)無縫(feng)接入(ru),同時保留原生PyTorch的使用體驗。
二、開源范圍
當前開源的Torch-MLU倉庫支持(chi)的社(she)區版本:
- PyTorch 2.1
- PyTorch 2.3
- PyTorch 2.4(2024年7月25日最新發布)
未來會(hui)隨社區版(ban)本持續更新。
三、擁抱開源社區,持續生態貢獻
寒武紀長期以來秉承開放、合作、共享的理念,積極參與開源社區的建設,在多個重要開源項目中貢獻代碼,如PyTorch、TensorFlow、Huggingface、Transformers、vLLM、Deepspeed等大模型訓練推理應用中(zhong)的核心組件。
近期(qi),寒武紀開源了(le)Triton-Linalg AI編譯(yi)器(qi)前(qian)端(duan),開發(fa)者(zhe)或者(zhe)硬件廠(chang)商可以以極低的(de)(de)開發(fa)成本,快速集成支(zhi)持Triton語言特性的(de)(de)后端(duan)指令(ling)集,并(bing)對接(jie)上層AI應用(yong)。
此次開(kai)源Torch-MLU插件,也是希望未來能更好地理解和(he)更快(kuai)速地解決開(kai)發者(zhe)的(de)問題,同(tong)時為寒(han)武紀(ji)深(shen)度學習框架與開(kai)發者(zhe)之間建(jian)立直接的(de)交(jiao)流渠(qu)道。寒(han)武紀(ji)堅信,推動人工智(zhi)能領(ling)域未來發展的(de)關鍵(jian)在于開(kai)放性、合(he)作精神和(he)共(gong)享資源。通過開(kai)源項目,開(kai)發者(zhe)們可以(yi)團結一致,通過協作和(he)配合(he),使得(de)技術成果得(de)到(dao)廣泛應(ying)用。
未(wei)來(lai),寒武(wu)紀(ji)將(jiang)持續貢獻開源社區,為AI開源生態(tai)注入新動力。同時通過(guo)技術的開源,賦予開發(fa)者(zhe)更多(duo)能力,共同促進(jin)繁榮和多(duo)元的人工智能世界。