2020 年,谷歌发表了预印本论文《Chip Placement with Deep Reinforcement Learning》,介绍了其设计芯片布局的新型强化学习方法。后来在 2021 年,谷歌又发表在 Nature 上并开源了出来。
今天,谷歌发表了这篇 Nature 文章的附录,更详细介绍了该方法及其对芯片设计领域的影响。同时,谷歌也开放了一个在 20 个 TPU 模块上预训练的检查点,分享模型权重并命名为「AlphaChip」。
- Nature 附录地址:https://www.nature.com/articles/s41586-024-08032-5
- 预训练检查点地址:https://github.com/google-research/circuit_training/?tab=readme-ov-file#PreTrainedModelCheckpoint
谷歌首席科学家 Jeff Dean 表示,开放预训练 AlphaChip 模型检查点以后,外部用户可以更容易地使用 AlphaChip 来启动自己的芯片设计。
x@JeffDean
计算机芯片推动了 AI 的显著进步,AlphaChip 利用 AI 来加速和优化芯片设计。该方法已被用于设计谷歌自定义 AI 加速器(TPU)最近三代的「超人」芯片布局。
作为首批用于解决现实世界工程问题的强化学习方法,AlphaChip 只需要数小时便能完成媲美或者超越人类的芯片布局,而无需再花费数周或数月人工努力。并且该方法设计的布局已应用于世界各地的芯片,覆盖场景包括数据中心到手机。
谷歌 DeepMind 表示,AlphaChip 已经彻底改变了我们设计微芯片的方式,从帮助设计用于构建 AI 模型的 SOTA TPU 到数据中心 CPU,它的广泛影响已经扩展到了 Alphabet 内外。
谷歌 DeepMind 联合创始人兼 CEO Demis Hassabis 表示,如今我们形成了这样一种反馈回路:训练 SOTA 芯片设计模型(AlphaChip)→使用 AlphaChip 来设计更好的 AI 芯片→使用这些 AI 芯片来训练更好的模型→再设计更好的芯片,这正是谷歌 TPU 堆栈表现如此好的部分原因。
各路网友对谷歌的 AlphaChip 寄予厚望,称「芯片设计芯片的时代来了」,还预言谷歌将赢得未来 AGI 之争。
AlphaChip 是如何工作的?
芯片设计并非易事,部分原因在于计算机芯片由许多相互连接的块组成,这些块具有多层电路元件,所有元件都通过极细的导线连接。此外,芯片还有很多复杂且相互交织的设计约束,设计时必须同时满足所有约束。由于这些复杂性,芯片设计师们在 60 多年来一直在努力实现芯片布局规划过程的自动化。
与 AlphaGo 和 AlphaZero 类似,谷歌构建了 AlphaChip,将芯片布局规划视为一种博弈。
AlphaChip 从空白网格开始,一次放置一个电路元件,直到完成所有元件的放置。然后根据最终布局的质量给予奖励。谷歌提出了一种新颖的「基于边」的图神经网络使 AlphaChip 能够学习互连芯片元件之间的关系,并在整个芯片中进行推广,让 AlphaChip 在其设计的每一个布局中不断进步。
左图:动画展示了 AlphaChip 在没有任何经验的情况下放置开源处理器 Ariane RISC-V CPU。右图:动画展示 AlphaChip 在对 20 个 TPU 相关设计进行练习后放置相同的块。
谷歌借助 AI 设计 AI 加速器芯片
自 2020 年发布以来,谷歌已经采用 AlphaChip 为每一代 Google TPU 生成超级芯片布局。这些芯片使得大规模扩展基于 Google Transformer 架构的 AI 模型成为可能。
TPU 作为谷歌强大的生成式 AI 系统的核心,应用范围从大语言模型(如 Gemini)到图像和视频生成器(Imagen 和 Veo)。这些 TPU 是 Google AI 服务的核心,可通过 Google Cloud 供外部用户使用。
谷歌数据中心存放的基于 Cloud TPU v5p 的超级计算机。
为了设计 TPU 布局,AlphaChip 首先在前几代的各种芯片块上进行练习,例如片上和芯片间网络块、内存控制器和数据传输缓冲区。这个过程称为预训练。然后谷歌在当前的 TPU 块上运行 AlphaChip 以生成高质量的布局。与之前的方法不同,AlphaChip 解决了更多芯片布局任务实例,因此变得更好、更快,就像人类专家所做的那样。
随着每一代新 TPU(包括谷歌最新的 Trillium(第 6 代))的推出,AlphaChip 设计出了更好的芯片布局并提供了更多的整体平面图,从而加快了设计周期并产生了性能更高的芯片。
图中显示了 Google 三代张量处理单元 (TPU)(包括 v5e、v5p 和 Trillium)中 AlphaChip 设计的芯片块的数量。
图中显示了 AlphaChip 在三代 Google 张量处理单元 (TPU) 中的平均线长(wirelength)减少量,并与 TPU 物理设计团队生成的位置进行了比较。
AlphaChip 带来的更广泛影响
AlphaChip 的影响力体现在 Alphabet、研究界和芯片设计行业的应用上。除了设计 TPU 等专用 AI 加速器外,AlphaChip 还为 Alphabet 的其他芯片设计布局,例如 Google Axion 处理器,这是谷歌首款基于 Arm 的通用数据中心 CPU。
外部组织也在采用和构建 AlphaChip。例如,全球顶级芯片设计公司之一联发科扩展了 AlphaChip,以加速其最先进芯片(如三星手机使用的 Dimensity Flagship 5G)的开发,同时提高了功耗、性能和芯片面积。
AlphaChip 引发了芯片设计 AI 工作的爆炸式增长,并已扩展到芯片设计的其他关键阶段,例如逻辑综合和宏选择。
开创芯片新未来
谷歌坚信,AlphaChip 有潜力优化从计算架构到制造的芯片设计周期的每个阶段,并改变智能手机、医疗设备、农业传感器等日常设备中定制硬件的芯片设计。
目前,AlphaChip 的未来版本正在开发中。谷歌期待与社区合作,继续变革自动芯片设计领域,从而在未来迎来速度更快、价格更低、能效更高的芯片。