谷歌揭秘TPU,AlphaChip发表Nature,深度探讨AI芯片发展。
近日,谷歌 DeepMind 在 Nature 上正式公布了其最新的芯片设计算法 AlphaChip,该方法致力于加速和优化计算机芯片的开发,已经历经多款 TPU 的产品考验,可在短短数小时内完成人类专家需要数周甚至数月的芯片布局设计。
2020 年,谷歌发表了一篇具有里程碑意义的预印本论文「Chip Placement with Deep Reinforcement Learning」,首次向世界展示了其采用新型强化学习方法设计的芯片布局。这一创新使得谷歌能够在 TPU 的芯片设计中引入 AI,实现了超越人类设计师的芯片布局。
到了 2022 年,谷歌进一步开源了该论文中描述的算法代码,使得全球的研究人员都能够利用这一资源对芯片块进行预训练。
如今,这一由 AI 驱动的学习方法已经经历了 TPU v5e、TPU v5p 以及 Trillium 等多代产品的考验,并在谷歌内部取得了显著的成就。更令人瞩目的是,谷歌 DeepMind 团队最近在 Nature 上发表了该方法的附录,更为详尽地阐述了其对芯片设计领域所产生的深远影响。同时,谷歌还开放了一个基于 20 个 TPU 模块预训练的检查点,分享了模型权重,并将其命名为 AlphaChip。
AlphaChip 的问世,不仅预示着 AI 在芯片设计领域的应用将变得更加广泛,也标志着我们正迈向一个由「芯片设计芯片」的全新时代。
AlphaChip:谷歌 DeepMind 如何用 AI 革新芯片设计
作为谷歌 DeepMind 的巅峰之作,AlphaChip 正以其在芯片设计领域的革命性进展,捕获全球科技界的瞩目。
芯片设计是一项位于现代科技之巅的领域,其复杂性在于将无数精密元件通过极其细微的导线巧妙连接。作为首批应用于解决现实世界工程问题的强化学习技术之一,AlphaChip 能够在短短数小时内完成与人类相媲美甚至更优的芯片布局设计,无需耗费数周或数月的人力劳动。这一划时代的进展,为我们打开了超越传统极限的想象之门。
那么,AlphaChip 究竟是如何实现这一壮举的呢?
AlphaChip 的秘诀在于其采用的强化学习原理,将芯片布局设计视为一场游戏。从一个空白的网格出发,AlphaChip 逐步放置每一个电路元件,直至全部就位。随后,根据布局的优劣,系统会给予相应的奖励。
更重要的是,谷歌创新性地提出了一种「基于边」的图神经网络,使得 AlphaChip 能够学习芯片元件之间的相互关系,并将其应用于整个芯片的设计中,从而在每一次设计中实现自我超越。与 AlphaGo 类似,AlphaChip 可以通过「游戏」学习,掌握设计卓越芯片布局的艺术。
在设计 TPU 布局的具体过程中,AlphaChip 首先会在前几代芯片的各类模块上进行预训练,包括芯片上和芯片间的网络模块、内存控制器和数据传输缓冲区等。这一预训练阶段为 AlphaChip 提供了丰富的经验。随后,谷歌利用 AlphaChip 为当前 TPU 模块生成高质量的布局。
与传统方法不同,AlphaChip 通过解决更多的芯片布局任务,不断优化自身,正如人类专家不断通过实践提升技能一样。正如 DeepMind 联合创始人兼 CEO Demis Hassabis 所言,谷歌已经围绕 AlphaChip 建立了一个强大的反馈循环:
* 首先,训练先进的芯片设计模型 (AlphaChip)
* 其次,使用 AlphaChip 设计更优秀的 AI 芯片
* 然后,利用这些 AI 芯片训练更出色的模型
* 最后,利用这些模型再去设计更出色的芯片
如此反复,实现了模型与 AI 芯片的同步升级,Demis Hassabis 表示,「这正是谷歌 TPU 堆栈表现如此好的部分原因」。
与人类专家相比,AlphaChip 不仅放置的模块数量更多,而且布线长度也大大减少。随着每一代新 TPU 的推出,AlphaChip 设计出了更优秀的芯片布局,提供了更完善的整体平面图,从而缩短了设计周期并提升了芯片性能。
谷歌在三代 TPU(v5e、TPU v5p) 中 AlphaChip 设计芯片块的数量与平均线长减少量
谷歌 TPU 的 10 年征程:从 ASIC 的坚持到 AI 设计的革新
作为 TPU 领域的探索者与先行者,纵观谷歌在这一技术线上的发展历程,不仅仅是凭借其敏锐的洞察力,更彰显了其非凡的魄力。
众所周知,在 20 世纪 80 年代,ASIC (Application Specific Integrated Circuit) 以其成本效益高、处理能力强和速度快的特点,赢得了市场的广泛青睐。然而,ASIC 的功能是由定制的掩模工具决定的,这就意味着,客户需要支付昂贵的前期一次性工程 (NRE) 费用。
此时,FPGA (Field Programmable Gate Array) 以其降低预付成本和减少定制数字逻辑风险的优势,走入大众视野,尽管在性能上并非全面超越,却在市场上独树一帜。
当时,业界普遍预测摩尔定律将推动 FPGA 的性能超越 ASIC 的需求。但事实证明,FPGA 作为一种可编程的「万能芯片」,虽然在探索性和小批量产品中表现出色,能够实现比 GPU 更优秀的速度、功耗或成本指标,但仍然无法摆脱「通用性与最优性不可兼得」的规律。一旦 FPGA 为某个专用架构铺平了道路,它就会让位给更为专业的 ASIC。
进入 21 世纪后,AI 技术热潮一浪高过一浪,机器学习、深度学习算法持续迭代,业界对于高性能、低功耗的专用 AI 计算芯片需求上涨,CPU、GPU 等在很多复杂任务上显得力不从心。在此背景下,谷歌在 2013 年做出了一个大胆的决定,选择 ASIC 来构建 TPU 基础设施,围绕 TensorFlow 和 JAX 进行开发。
值得注意的是,自主研发 ASIC 是一个周期长、投入大、门槛高、风险极大的过程。一旦方向选择错误,可能会导致巨大的经济损失。然而,为了探索更具成本效益和节能的机器学习解决方案,谷歌在 2012 年通过深度学习在图像识别上取得突破性进展后,就立即在 2013 年开始研发 TPUv1,并在 2015 年宣布第一代 TPU 芯片 (TPU v1) 在内部上线,这标志着全球首款专为 AI 设计的加速器诞生。
幸运的是,TPU 很快就迎来了一个引人注目的展示机会——2016 年 3 月,AlphaGo Lee 成功击败了世界围棋冠军李世石,作为 AlphaGo 系列的第二代版本,其运行于 Google Cloud,耗用 50 个 TPU 进行计算。
然而,TPU 并没有因此立即在业界取得大规模的成功应用,直到 AlphaChip 芯片布局方法的提出,TPU 才真正走进全新的发展阶段。
谷歌 TPU 发展历程
2020 年,谷歌在预印本论文「Chip Placement with Deep Reinforcement Learning」中展示了 AlphaChip 的能力,它能够从过去的经验中学习并不断改进,通过设计一种能够准确预测各种网表及其布局上的奖励神经体系结构,能够对输入网表生成丰富的特征嵌入。
AlphaChip 将性能优化的条件视为游戏的胜利条件,采用强化学习的方法,通过训练一个智能体,以累计奖励最大化为目标,不断优化芯片布局的能力。他们开启了 1 万局游戏,让 AI 在 1 万个芯片上练习布局布线并收集数据,同时不断学习优化。
最终,他们发现与人类工程师相比,AI 在面积、功率和电线长度方面的表现更优越或媲美手动布局,同时满足设计标准所需的时间要少得多。结果表明,AlphaChip 在不到 6 小时的时间内就可以生成媲美或超过人工的现代加速器网表上的布局,而在同样条件下,现有人类专家可能需要几个星期来完成同样的工作。
在 AlphaChip 的助力下,谷歌对 TPU 的依赖日益增加。2023 年 12 月,谷歌推出了多模态通用大模型 Gemini 的 3 个不同版本,该模型的训练大量使用了 Cloud TPU v5p 芯片。2024 年 5 月,谷歌又发布了第六代 TPU 芯片 Trillium,它能够在单个高带宽、低延迟 Pod 中扩展为多达 256 个 TPU 的集群,相较于前代产品,Trillium 在适配模型训练方面的功能更强。
同时,TPU 芯片也逐渐走出谷歌公司,获得了更广泛的市场认可。2024 年 7 月 30 日,苹果公司在其发布的一篇研究论文中声称,在训练 Apple Intelligence 生态中的人工智能模型 AFM 时,选择了谷歌的两种张量处理单元 (TPU) 云集群。另有数据显示,超过 60% 的生成式 AI 初创公司和近 90% 的生成式 AI 独角兽都在使用 Google Cloud 的 AI 基础设施和 Cloud TPU 服务。
种种迹象表明,谷歌十年磨一剑,TPU 已经走出培育期,开始以卓越的硬件性能为谷歌在 AI 时代进行反哺。而 AlphaChip 所蕴含的「AI 设计 AI 芯片」路径,也为芯片设计领域开辟了新的视野。
AI 革新芯片设计:从谷歌 AlphaChip 到全流程自动化的探索
尽管 AlphaChip 在 AI 设计芯片领域独树一帜,但它并非孤军奋战。AI 技术的触角已经广泛延伸至芯片验证、测试等多个关键环节。
芯片设计的核心任务是优化芯片的功耗 (Power)、性能 (Performance) 和面积 (Area) ,这三个关键指标统称为 PPA,这一挑战也被称作设计空间探索。传统上,这一任务由 EDA 工具完成,但为了达到最佳性能,芯片工程师必须不断手工调整,然后再次交给 EDA 工具进行优化,如此循环往复。这个过程就像在家中摆放家具,不断尝试以实现空间利用最大化和动线最优化,但每次调整都相当于将家具搬出再重新布置,极其耗时耗力。
为了破解这一难题,Synopsys 新思科技在 2020 年推出了 DSO.ai,这是业界首个融合 AI 与 EDA 的芯片设计解决方案。DSO.ai 采用强化学习技术,通过 AI 自动搜索设计空间,寻找最佳平衡点,无需人工干预。这一工具已在多家芯片巨头中得到应用。
例如,微软在使用 DSO.ai 后将芯片模块的功耗降低了 10%-15%,同时保持性能不变;意法半导体将 PPA 探索效率提高了 3 倍以上;存储芯片巨头 SK 海力士则将芯片面积减少了 5%。Synopsys 的数据显示,DSO.ai 已成功助力超过 300 次商业流片,标志着 AI 在真正的芯片设计与生产中发挥了重要作用。
在 AI 辅助芯片验证方面,Synopsys 新思科技发布的技术报告也指出,验证过程占据了整个芯片开发周期的高达 70% 的时间。一次芯片流片的成本高达数亿美元,而现代芯片的复杂性不断增加,验证难度可想而知。为此,Synopsys 新思科技推出了 VSO.ai 工具,利用 AI 优化验证空间,加速覆盖率的收敛速度。
VSO.ai 能够推断出不同的覆盖率类型,与传统的代码覆盖率形成互补,AI 还能从验证经验中学习,不断优化覆盖率目标。此外,Synopsys 新思科技还推出了 TSO.ai 工具,能够帮助芯片开发者筛选出代工厂制造的有缺陷的芯片。
AI 在芯片设计领域的深度介入,引发了一个大胆的设想:我们能否用 AI 设计一颗完整的芯片?事实上,英伟达已经在这一领域进行了尝试。通过深度强化学习代理设计电路,英伟达的 H100 中就有近 13,000 条电路由 AI 设计。中国科学院计算所也利用 AI 在 5 小时内生成了一个名为「启蒙一号」的 RISC-V 处理器芯片,拥有 400 万个逻辑门,性能与 Intel 80486 相当。
总体来看,AI 设计完整芯片的能力依然有限,但这无疑是未来芯片发展的一个重要机遇。随着技术的不断进步,AI 在芯片设计领域的潜力必将得到进一步挖掘和利用,并最终改变整个芯片的设计过程。