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