Layout工程师危，谷歌自动芯片设计AlphaChip问世，开放权重可外部定制

机器学习研究组订阅 · 公众号 · AI · 2024-09-27 19:01

正文

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 物理设计团队生成的位置进行了比较。

Layout工程师危，谷歌自动芯片设计AlphaChip问世，开放权重可外部定制

正文

请到「今天看啥」查看全文