近日,中國移動研究院撰寫的論文“Efficient Automatic Arrangement Algorithm for Computing In Memory Chips Array”被第三十一屆國際神經(jīng)信息處理大會(ICONIP2024)錄用。
近年來,人工智能尤其是大模型對算力的需求呈爆炸式增長,而經(jīng)典的馮·諾依曼架構(gòu)因存儲與計算分離帶來的數(shù)據(jù)搬運時延及能耗成為算力增長的主要瓶頸。存算一體技術(shù)在存儲原位實現(xiàn)計算功能,可以突破馮·諾依曼架構(gòu)瓶頸,大幅提升算力和能效水平。但是,由于存算一體芯片涉及架構(gòu)、器件、材料、算法的全新創(chuàng)新,面臨模型編排復雜度高、陣列利用率低下的問題,阻礙存算一體芯片的應用拓展。針對上述問題,論文提出一種面向存算一體芯片的新型多陣列AI模型編排算法,可以大幅提升編排效率和芯片陣列利用率。該成果為面向存算一體芯片的AI模型編排部署提供重要參考,對于推進存算一體芯片的工程化和產(chǎn)業(yè)化有重要意義。
圖1 論文首頁
論文將AI模型不同層權(quán)重在存算一體芯片多陣列的優(yōu)化編排問題轉(zhuǎn)化為面向陣列利用率提升的序列到序列排序問題,創(chuàng)新地提出了一種基于改進指針網(wǎng)絡(luò)(Pointer Net)的AI模型分層排序算法,采用長短期記憶網(wǎng)絡(luò)(LSTM)作為指針網(wǎng)絡(luò)的編碼器(Encode)與解碼器(Decode),增加了模型序列記憶能力,同時引入注意力機制,提升算法序列預測的準確性。圖2為算法網(wǎng)絡(luò)結(jié)構(gòu)圖,輸入數(shù)據(jù)為模型各層權(quán)重尺寸與存算陣列尺寸的比例,輸出數(shù)據(jù)為模型該層權(quán)重在陣列中的編排順序。
圖2. 基于序列到序列指針網(wǎng)絡(luò)的AI模型分層排序算法結(jié)構(gòu)圖
論文采用人工蜂群算法(Artifical Bee Colony,ABC)在ONNX開源模型庫搜索匹配不同存算一體芯片陣列特征的網(wǎng)絡(luò)計算層最優(yōu)排序,作為標注數(shù)據(jù)集訓練指針網(wǎng)絡(luò)訓練。在模型編排階段,論文提出的基于序列到序列的模型編排算法計算復雜度由ABC算法O(n3)降為O(n),其中n為AI模型的層數(shù),預測準確率達到ABC算法的96%,在保證準確率的同時大幅降低計算時延。尤其對于層數(shù)較多的復雜AI模型,本文方法發(fā)揮的優(yōu)勢更為明顯,如圖3所示,對于100層以上網(wǎng)絡(luò),編排效率超過傳統(tǒng)ABC算法10倍。
圖3. 論文算法跟人工蜂群算法計算時延對比圖
最后,基于指針網(wǎng)絡(luò)輸出地分層編排順序,綜合AI算子特征和陣列冗余空間形成算子在陣列的最優(yōu)編排策略,有效保證芯片陣列占用率。圖4是Yolov5s模型在一個2陣列(陣列尺寸1024x512)存算一體芯片上的編排結(jié)果,陣列1利用率達94.14%,陣列2達利用率93.97%。
圖4. Yolov5s網(wǎng)絡(luò)模型在存算一體芯片上的編排結(jié)果
下一步,中國移動研究院將繼續(xù)開展存算一體關(guān)鍵技術(shù)研究,圍繞軟件架構(gòu)、優(yōu)化算法、場景應用等環(huán)節(jié)持續(xù)發(fā)力,推動存算一體技術(shù)成熟。