博文

中国半导体十大研究进展候选推荐（2022-022）——支持多粒度稀疏的AI训练芯片

已有 1553 次阅读 2022-9-28 11:34 |系统分类:论文交流

1 工作简介

——支持多粒度稀疏的AI训练芯片

清华大学集成电路学院魏少军、尹首一教授团队在人工智能（AI）训练芯片方向取得突破。该团队设计的AI训练芯片Trainer，以实现场景自适应的高能效模型训练为目标，突破传统AI训练芯片学习机制和电路实现的技术局限性，为高效AI模型训练提供了坚实的硬件基础。2022年5月5日，该研究成果以"Trainer: An Energy-Efficient Edge-Device Training Processor Supporting Dynamic Weight Pruning"为题发表于集成电路领域顶级期刊IEEE Journal of Solid-State Circuits (JSSC)。

近年来，AI作为一项影响深远的颠覆性技术，在机器翻译、人机交互、医学诊断、自动驾驶等多个领域取得了突破性进展。AI技术的成功高度依赖于“算法、算力、数据”三个关键要素。AI算法的参数量和训练数据量爆炸式增长，给AI芯片带来巨大的能耗，严重制约AI技术的持续发展和广泛应用。例如，迄今最强的自然语言处理模型GPT-3具有1750亿参数，其训练所用的数据量高达45TB，需要在微软Azure云平台的1万颗GPU训练30天，消耗1.16×10¹³ J的能量(约为3000个成年人1年的能量消耗)。因此，提升AI模型训练的能效成为实现AI持续发展必须克服的严峻挑战。然而，基于传统训练机制的AI芯片难以解决这一问题。通常，AI模型的训练包含两个阶段。首先，需要基于特定数据集在AI芯片上对模型所有参数进行训练，以达到理想的推理精度。而后，利用模型的冗余性，对较小的参数进行剪枝和再训练，减小模型规模。这种机制需要首先对所有参数训练更新，消耗大量的训练时间和能量。

为了避免对冗余参数训练导致的时间和能量浪费，Trainer采用边更新边剪枝的训练机制，其流程如图1所示。面对不同复杂度的应用场景，Trainer在训练迭代过程中，基于当前训练精度自适应生长或修剪网络连接。在每次迭代过程中，Trainer只使用和更新保留的参数，避免冗余参数的相关计算，从而大幅减少计算和访存开销，高效适应多样化应用场景。

图1. 基于动态权重剪枝的稀疏训练机制。

动态剪枝可有效减少训练计算量，需要设计全新的训练芯片架构从而充分利用其动态权重稀疏特性。Trainer包含三个关键技术，实现对动态权重稀疏的高效利用，其整体架构如图2所示。首先，Trainer包含系统级冗余计算预测单元，通过剖析训练过程中结构化权重稀疏在前馈计算、反向传播和权重更新三个阶段的全局作用，预测并移除训练阶段中的隐式冗余计算。不同于显式冗余计算，隐式冗余计算的输入值、权重值以及输出值均不为0，但对训练无效。其次，Trainer针对不规则的非结构化稀疏权重，采用实时复用检测、乱序稀疏压缩的计算数据流，动态适配权重复用情况，解决不规则权重稀疏导致的数据复用不均衡问题，提高训练过程中的硬件资源利用率。最后，Trainer通过提取BN计算公因子，并基于公因子重组BN公式的方式，解耦BN计算的串行数据依赖，实现并行正反向BN计算，减少训练过程中访存开销。解决卷积层和全连接层运算量随动态剪枝显著减小后，串行BN计算时重复数据访存导致的训练瓶颈。

图2. Trainer芯片的整体架构。

Trainer芯片版图和性能总结如图3所示，采用28 nm CMOS工艺实现，在FP8的计算精度下，Trainer的峰值能效为276.55TFLOPS/W，是NVIDIA A100 GPU的177.3倍。相比于GPU的训练后剪枝的模型进化机制，Trainer可以减少60倍的训练时间和1500倍的训练能量。团队此次研究成果，从训练机制和硬件架构角度为现有AI训练芯片带来了突破，显著增强了芯片面向不同任务时的学习效率，大幅减少芯片训练的时间和能量开销, 为AI训练芯片的演进开拓了新方向。

图3. Trainer芯片的版图与性能总结。

2 作者简介

通讯作者

尹首一，博士，清华大学教授，集成电路学院副院长，国家杰出青年科学基金获得者。

研究方向为可重构计算、人工智能芯片设计。已发表学术论文200余篇，包括IEEE JSSC、TPDS、TCSVT、TVLSI、TCAS-I/II和ISSCC、ISCA、VLSI、DAC、HPCA等集成电路和体系结构领域权威期刊和学术会议。出版《可重构计算》、《人工智能芯片设计》专著2部。曾获国家技术发明二等奖、中国电子学会技术发明一等奖、中国发明专利金奖、教育部技术发明一等奖、江西省科技进步二等奖、中国电子学会优秀科技工作者奖、中国电子信息领域优秀科技论文奖。现任集成电路领域国际会议IEEE DAC、ISCA、MICRO和A-SSCC的技术委员会委员，国际期刊《IEEE Transactions on Circuits and System I: Regular Papers》、《ACM Transactions on Reconfigurable Technology and Systems》及《Integration, the VLSI Journal》的Associate Editor 。

第一作者

王扬，清华大学集成电路学院博士后。

王扬，长期从事低功耗VLSI设计研究，共负责/参与完成可重构AI芯片/数字通信芯片共10余颗。包括多颗CNN推理、训练处理器、以及国际首款Transformer处理器，成果以第一作者发表于IEEE JSSC, IEEE TCAS-I, ISSCC, VLSI等集成电路顶级期刊会议。获得2019年度“未来芯片学者计划”项目支持，作为课题负责人承担2022年度科技创新2030—“新一代人工智能”重大项目。

3 原文传递

详情请点击论文链接：

https://ieeexplore.ieee.org/document/9779311

转载本文请联系原作者获取授权，同时请注明本文来自阎军科学网博客。
链接地址：https://blog.sciencenet.cn/blog-3406013-1357193.html

上一篇：中国半导体十大研究进展候选推荐（2022-021）——面向智能机器人、具有时空弹性的类脑计算芯片：天机X
下一篇：中国半导体十大研究进展候选推荐（2022-023）——具有室温面内反常霍尔效应的异维超晶格结构