近日，中國移動研究院撰寫的論文“Collaborative Training for Compensation of Inference Errors in NOR Flash Computing in memory Chips”被第二十七屆IEEE國際計算機(jī)協(xié)同計算與設(shè)計大會（CSCWD 2024）錄用，中國移動研究院專家受邀出席大會并就論文方案做演講。
近年來人工智能尤其是大模型對算力的需求呈爆炸式增長，而經(jīng)典的馮·諾依曼架構(gòu)因存儲與計算分離帶來的數(shù)據(jù)搬運(yùn)時延及能耗成為算力增長的主要瓶頸。存算一體技術(shù)在存儲原位實現(xiàn)計算功能，可以突破馮·諾依曼架構(gòu)瓶頸，大幅提升算力和能效水平。但由于存算一體器件存在非理想特性，造成數(shù)據(jù)轉(zhuǎn)換誤差和數(shù)據(jù)寫入誤差問題，從而影響計算精度。針對上述問題，論文提出一種面向存算一體芯片的AI模型協(xié)同訓(xùn)練架構(gòu)，通過在模型訓(xùn)練過程中使用存算一體芯片的推理結(jié)果進(jìn)行反向傳播并更新模型權(quán)重，從而提升模型在存算一體芯片推理的魯棒性。該成果為面向存算一體芯片的模型訓(xùn)練提供重要參考，對于推進(jìn)存算一體芯片的工程化和產(chǎn)業(yè)化有重要意義。
在訓(xùn)練架構(gòu)方面，論文提出了一種面向存算一體芯片的新型協(xié)同訓(xùn)練架構(gòu)（圖1），包括數(shù)據(jù)集量化、模型訓(xùn)練與量化、模型片上計算三部分。首先按照芯片計算精度對輸入的訓(xùn)練數(shù)據(jù)集進(jìn)行INT8量化，然后基于量化數(shù)據(jù)集在CPU/GPU上訓(xùn)練出FP32精度的模型，再將模型進(jìn)行量化后部署到存算一體芯片上并基于量化數(shù)據(jù)集進(jìn)行前向計算，最后基于芯片實際計算結(jié)果與真值計算模型損失并反向傳播更新模型權(quán)重，依此訓(xùn)練出來的模型可以極大提升在存算一體芯片上推理計算的魯棒性。
圖1 面向存算一體芯片的AI模型協(xié)同訓(xùn)練架構(gòu)
此外，論文提出了一種基于對稱定比例因子的協(xié)同量化方法，可以將芯片端的INT8計算精度和訓(xùn)練環(huán)境的FP32精度模型有效融合（圖2）。針對芯片數(shù)據(jù)[-128, 127]的精度范圍，在模型訓(xùn)練更新權(quán)重時將權(quán)重區(qū)間限制在[-0.125, 0.125]范圍，從而保持統(tǒng)一的1024權(quán)重比例因子，避免在訓(xùn)練過程中因更新量化因子而引入新誤差，并提升模型訓(xùn)練收斂速度。
圖2對稱定比例因子量化方法
下一步，中國移動研究院將持續(xù)開展存算一體芯片、軟件、算法、應(yīng)用等相關(guān)技術(shù)研究，推進(jìn)存算一體在端、邊、云等應(yīng)用場景的廣泛應(yīng)用落地。 
