出色的机器学习 (ML) 研究需要出色的系统。 随着当今使用的算法和硬件越来越复杂以及它们运行的规模越来越大，执行日常任务所需的软件的复杂性只会增加。 在本文中，我们概述了谷歌过去一年在 ML 系统方面取得的众多进步，这些进步使我们能够支持复杂模型的服务和训练，同时降低最终用户实施的复杂性。

去年，谷歌在改进系统以更好地支持 ML 中的大规模计算和一般科学计算方面取得了重大进展。 谷歌 TPU 硬件从一开始就在设计时考虑到了扩展性，每年都在努力突破界限。 去年，谷歌为大型模型设计了最先进的服务技术，改进了张量程序的自动分区，并重新设计了库的 API，以确保所有这些开发成果都能为广大用户所用。

最大的效率改进之一是用于评估神经网络核心的大规模矩阵乘法运算的 CollectiveEinsum 策略。 与以前流行的将通信与设备本地计算分开的 SPMD 分区策略不同，这种方法使用快速 TPU ICI 链接将它们重叠，从而导致高达 1.38 倍的性能提升。 该算法也是有效扩展 Transformer 推理工作的关键组成部分，它提出了多种策略，可以在延迟和硬件利用率之间进行权衡，达到 76% 的最先进模型 FLOPs 利用率 (MFU) 吞吐量优化的配置。
其次还将 SPMD 样式的分区作为一流概念集成到具有 DTensor 扩展的 TensorFlow 和具有重新设计的数组类型的 JAX 中。 在这两个库中，对于程序员来说似乎是完整的张量可以通过附加声明性布局注释透明地分片到多个设备上。 事实上，这两种方法都与为单设备计算编写的现有代码兼容，这些代码现在可以扩展到多设备程序中，通常无需任何代码修改！

将 SPMD 分区集成到 ML 框架的核心意味着能够推断和优化数组程序映射到更大的设备集的方式对于性能至关重要。 在过去，这推动了 GSPMD 的发展，这是该领域的一个重要里程碑。 然而，GSPMD 严重依赖启发式算法，有时仍需要手动做出重要决策，这通常会导致性能欠佳。 为了使分区推理完全自动化，我们与外部同事合作开发了 Alpa，这是一个全自动系统，用于探索运算符级（模型）并行性和较大子计算之间的管道并行性的策略。 它成功地匹配了 Transformers 等流行模型的手动调整性能，但也能够成功地扩展其他模型，例如卷积网络和专家混合模型，这些模型通常会导致现有的自动化方法陷入困境。
同样，最近发布的 Pathways 系统在通常的 TPU 运行时之上添加了一个额外的虚拟化层——加速器由长期存在的进程管理，而不是直接分配给用户。 然后，单个最终用户可以连接到任意数量的 Pathways 控制的设备并编写他们的程序，就好像所有设备都直接连接到他们的进程一样，即使实际上他们甚至可能跨越多个数据中心。 多亏了 Pathways：(1) 可以减少作业启动时间，(2) 更容易实现容错，(3) 多租户成为一种可行的选择，可以同时执行多个作业，从而更有效地利用硬件。 Pathways 能够轻松实现跨多个 TPU pod 的计算是至关重要的，因为它让我们避免了未来的扩展瓶颈。

另一个值得注意的版本是 TensorStore，这是一个用于多维数组存储的新库。 TensorStore 对于训练具有多控制器运行时的大型语言模型 (LLM) 特别有用，其中每个进程只管理所有参数的一个子集，所有这些都必须整理到一个一致的检查点中。 TensorStore 为高效并发的多维数组序列化提供数据库级保证 (ACID) 到许多存储后端（例如，谷歌云存储、各种文件系统、HTTP 服务器），并已成功用于计算密集型工作负载，例如 PaLM 和重建人类皮层和果蝇大脑。

内容中包含的图片若涉及版权问题，请及时与我们联系删除