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