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