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