谷歌利用 AI 设计芯片

摘要:就好像嫌摩尔定律在晶体管密度和成本方面改进的步伐在放缓还不够糟糕,设计芯片的成本以及蚀刻芯片的工厂的成本也在上升。为了保持IT创新阔步前行,任何这些成本得到节省都将大受欢迎。 目前,芯片设计界大有希望的研究

就好像嫌摩尔定律在晶体管密度和成本方面改进的步伐在放缓还不够糟糕,设计芯片的成本以及蚀刻芯片的工厂的成本也在上升。为了保持IT创新阔步前行,任何这些成本得到节省都将大受欢迎。

1.jpg

目前,芯片设计界大有希望的研究领域之一是,使用机器学习技术来切实帮助处理设计过程中的一些任务。芯片设计中使用机器学习也是谷歌研究部门的高级研究员Jeff Dean在本周于旧金山举行的2020年国际固态电路大会上发表主题演讲时畅谈的话题之一,他帮助谷歌这家超大规模提供商发明了许多关键技术。

事实证明,谷歌对计算引擎绝非一时心血来潮。谷歌是全球最大的CPU和GPU消费者之一,也是TPU的设计者,TPU的应用范围从边缘到数据中心,不一而足,既可用于机器学习推理又可以用于机器学习训练。因此,这不仅仅是这个搜索引擎巨头和公共云提供商的科研活动——如果它打算继续推进TPU路线图,如果它像竞争对手AWS那样决定开始设计自己的定制Arm服务器芯片,或决定为其手机及其他消费级设备设计定制Arm芯片,更是具有现实意义。

出于一定的偶然性,谷歌为了跨大量不同类型的计算引擎一直运行机器学习模型所做的一些工作回过头来进入到它为了使ASIC上IP功能块(IP block)的一些放置和布线实现自动化而在做的工作。

虽然谷歌早在2018年5月展示的TPUv3系统的pod可以将1024个张量处理器(核心数量增加一倍,时钟频率提高15%左右)组合起来,以提供每秒106千万亿次浮点运算的16位半精度乘法总性能(带32位累加),使用谷歌自家的非常巧妙的bfloat16数据格式。这些TPUv3芯片都使用32×32环形网格实现交叉耦合,那样它们就能共享数据,而每个TPUv3核心都有自己的HBM2存储器模组。这个TPUv3 pod是巨大的计算聚合体,能进行机器学习训练或机器学习推理,但是未必如谷歌需要制造的那么大。

就一句话,为了规避风险,谷歌采用结合CPU和GPU的混合架构(也许有一天会结合其他加速器)用于处理强化学习工作负载,因此Dean及其在谷歌的同事从事的研究成果也应用于ASIC设计上。

Dean解释道:“趋势之一就是模型越来越大。因此,整个模型在单单一块芯片上不一定装得下。如果你有庞大模型,那么模型并行处理(将模型分散于多块芯片之间)很重要,通过为它提供一堆计算设备来获得良好性能并非易事,如何有效地处理这件事也并不明显。”

由于AI系统具有异构性,这不像拿来用于在大规模并行超级计算机上分派工作的消息传递接口(MPI),稍加改动后部署到TensorFlow之类的机器学习框架上那么简单。但这可能是将机器学习训练工作负载分布在众多计算单元上的一条思路,已有人这么做了。与其他超大规模提供商一样,谷歌往往构建自己的框架、协议和数据存储系统,当然依托其他技术。

由于神经网络模型变得比存储器空间还要庞大,超过了单单一个CPU、GPU或TPU的计算能力,设备放置显得尤为重要;设备放置是指,在合适的时间将合适的神经网络(或体现神经网络的代码的一部分)放在合适的设备上,以便在整个应用中获得最大吞吐量。而这个问题迅速变得严重起来,框架和硬件跟不上步伐。不妨看一看:

2.jpg

参数的数量就在不断增加,并行使用的设备数量也在不断增加。实际上,让128个GPU或128个TPUv3处理器(这就是为什么上图中有512个核心)协同工作是相当了不起的成就,可与早期超级计算机最强大的处理能力相提并论。20多年前,使用MPI的松散耦合的大规模并行超级计算机唱主角,搭载物理共享存储器的联合NUMA服务器成为了HPC的标准。随着越来越多的设备以某种方式组合起来以处理这些模型,谷歌一直在尝试使用强化学习(机器学习的一个特殊分支),搞清楚全体模型在一组CPU和GPU上运行时,在任何特定的时间哪里运行神经网络模型效果最好。在这种情况下,制定初始策略以调派处理任务的神经网络模型,然后结果反馈回到模型中作进一步调整,以便那些模型运行起来越来越高效。

2017年,谷歌训练强化学习模型以完成这项工作(论文详见https://arxiv.org/abs/1706.04972),这是编码器和解码器最终放置的样子;而与专家手动放置的神经网络相比,机器学习模型将工作负载放置在接受测试的系统中的两个CPU和四个GPU上,其最终运行时间缩短19.3%。Dean补充道,利用强化学习将神经网络工作负载放置在计算引擎上“做了有点不合常理的事”才获得这种结果,许多机器学习应用似乎是这种情况。问题在于,它无法获得大量的强化学习计算能力将工作负载放置在设备上,以运行自行训练的神经网络。2018年,谷歌开展了研究以表明如何将计算图扩展到80000多个操作(节点),去年谷歌为有50000多个操作(节点)的数据流图创建了所谓的广义设备放置方案。

Dean说:“随后我们开始考虑使用该方法,而不是用它将软件计算放置在不同的计算设备上,我们开始考虑将它用于ASIC芯片设计中的放置和布线,因为问题看起来有点相似。强化学习确实适用于处理规则清楚的棘手问题(比如国际象棋或围棋);实际上我们开始问自己:我们能否获得强化学习模型成功地做好ASIC芯片布局这项工作?”

据Dean声称,这么做有几个挑战。举例说,国际象棋和围棋都有一个目标,那就是赢得比赛而不是输掉比赛。(它们是同一枚硬币的两面。)就ASIC上放置IP功能块和IP功能块之间的布线而言,没有简单的赢或输,你要关注很多目标,比如面积、时间、拥塞和设计规则等。更困难的是,神经网络模型为IP功能块放置所要管理的潜在状态数量众多,如下图所示:

3.jpg

最后,决定IP功能块放置的实际奖励函数(在EDA工具中运行)需要花大量时间来运行。

Dean解释:“于是我们有了一种架构(我不会透露很多细节),但实际上它设法把构成芯片设计的一堆东西放到晶圆上。”他展示了将IP功能块放置在低功率机器学习加速器芯片上的一些结果(我们假设该芯片是谷歌为其智能手机制造的边缘TPU),一些地方有意模糊处理,那样外人无法了解芯片的细节。“我们有一队专家来放置该IP功能块,他们有几个代理奖励函数,我们评估这些函数成本很低;我们在两秒钟而不是几小时内对它们进行了评估,这点非常重要,因为强化学习中要迭代多次。因此,我们有一个基于机器学习的放置系统;你可以看到,它稍微扩展了逻辑,而不是让逻辑分布在一个呈矩形的地方,这使它能够在拥塞和导线长度方面都有所改进。我们在迄今为止试用的所有不同IP功能块上取得了与人相当或更胜一筹的结果。”

不管怎样,这是针对强化学习网络的低功耗机器学习加速器与放置IP功能块的人员的比较:

4.jpg

而这张表格显示了手动放置和布线与借助机器学习自动放置和布线之间的区别:

5.jpg

最后,这表明了与人相比,强化学习网络如何处理TPU芯片上的IP功能块:

6.jpg

看看这些AI设计的IP功能块与人类设计的笛卡尔功能块相比是多么地合乎自然。太棒了。

谷歌完成这项工作后问了这个问题:我们可以训练快速有效地放置以前从未见过的新设计的普通代理吗?这正是你在设计新芯片时关注的重点。于是,谷歌针对来自TPU架构的四个不同IP功能块测试了这个广义模型,然后还在“Ariane”RISC-V处理器架构上进行了测试。下比是数据对比:

7.jpg

以下是关于在Ariane RISC-V芯片上进行放置和布线的更多数据:

8.jpg

Dean说:“你可以看到其他设计方面的情况实际上显著改善了结果,因此从本质上讲,你可以在12个小时内得到深蓝色条,”他指上面第一张图,然后以上面第二张图继续说事。“这张图显示了导线成本,我们可以看到如果你从头开始训练,实际上它花一点时间才有某种突破性的真知灼见,才能显著降低导线成本,而事先训练的策略只要看看其他设计,很快就能对芯片设计有一番大致的了解。”

就像我们进行大量模拟以便更准确预报天气一样,Dean说芯片设计中这种基于AI的IP功能块放置和布线可用于快速生成许多不同的布局,自有不同的美中不足。而万一需要添加某个功能,基于AI的芯片设计可以快速重新进行布局,不用花好几个月来做。

9.jpg

而最重要的是,这种自动化设计辅助有望从根本上降低设计新芯片的成本。这类成本在急剧增长,而我们看到的数据(感谢IT行业知名人士、Arista Networks主席兼首席技术官Andy Bechtolsheim),采用16纳米工艺的先进芯片设计平均成本高达1.063亿美元,改而10纳米会使成本增加至1.744亿美元,改而7纳米需要2.978亿美元的成本,预计5纳米芯片的成本在5.422亿美元左右。其中近一半的成本用于软件,现在依然如此。因此,我们知道竭力消除哪方面的部分成本,而机器学习可助一臂之力。

问题是,芯片设计软件开发商会不会融入AI,并促使芯片设计领域迎来可以真正称为寒武纪的爆炸式增长,然后实现批量生产?这方面会发生什么、谷歌开展的此类研究将如何提供帮助值得拭目以待。


本文为 角度商业(http://www.jiaodoo.cn)转载作品,作者: 云头条。转载请联系原作者。本文仅代表作者观点,不代表角度商业观点。

发表评论