出色的机器学习 (ML) 研究需要出色的系统。随着当今使用的算法和硬件的日益复杂以及运行规模的不断扩大，执行日常任务所需的软件的复杂性也只会增加。在这篇文章中，我们概述了 Google 过去一年在 ML 系统方面取得的众多进展，这些进展使我们能够支持复杂模型的服务和训练，同时降低最终用户的实施复杂性。这篇博文还重点介绍了我们利用 ML 本身来帮助改进和设计下一代系统堆栈的研究。
机器学习的分布式系统
今年，我们在改进系统方面取得了重大进展，以更好地支持 ML 和一般科学计算中的大规模计算。自诞生之日起，Google TPU 硬件就一直以扩展为设计理念，每年我们都努力进一步突破界限。今年，我们为大型模型设计了最先进的服务技术，改进了张量程序的自动分区，并重新设计了我们库的 API，以确保所有这些开发成果都能被广大用户所使用。
我们今年最大的效率改进之一是CollectiveEinsum 策略，用于评估神经网络核心的大规模矩阵乘法运算。与之前流行的将通信与设备本地计算分开的SPMD分区策略不同，此方法使用快速TPU ICI 链路将它们重叠，从而将性能提高高达 1.38 倍。该算法也是我们有效扩展 Transformer 推理工作的关键组成部分，它提出了在延迟和硬件利用率之间权衡的多种策略，在吞吐量优化的配置中达到了 76% 的最先进的模型FLOP利用率 (MFU)。
CollectiveEinsum 策略中提出的具有双向层内模型并行性的 AllGather-Einsum 的图示。顶部：非重叠执行的图示。底部： CollectiveEinsum 技术的图示。
我们还将 SPMD 样式分区作为一流概念集成到TensorFlow（带有DTensor 扩展）和JAX（带有重新设计的数组类型）中。在这两个库中，程序员认为完整的张量只需附加声明性布局注释即可透明地分片到多个设备上。事实上，这两种方法都与为单设备计算编写的现有代码兼容，现在可以扩展到多设备程序，通常无需任何代码修改！
将 SPMD 分区集成到我们的 ML 框架核心意味着能够推断和优化阵列程序映射到更大设备集的方式对于性能至关重要。在过去，这推动了GSPMD的开发，这是该领域的一个重要里程碑。然而，GSPMD 严重依赖启发式方法，有时仍需要手动做出非平凡的决策，这通常会导致性能不佳。为了使分区推理完全自动化，我们与外部同事合作开发了Alpa，这是一个完全自动化的系统，可探索运算符级（模型）并行性和较大子计算之间的流水线并行性策略。它成功地匹配了 Transformers 等流行模型的手动调整性能，但也能够成功扩展其他模型，例如卷积网络和专家混合模型，这些模型通常会导致现有的自动化方法陷入困境。
Alpa概述。操作符间确定将子图分配给子网格的最佳方式。操作符内传递为每个管道阶段找到最佳的操作符内并行计划。最后，运行时编排生成一个静态计划来对计算和通信进行排序。
同样，最近发布的Pathways系统在通常的 TPU 运行时之上添加了额外的虚拟化层——加速器由长寿命进程管理，而不是直接分配给用户。然后，单个最终用户可以连接到任意数量的 Pathways 控制的设备并编写他们的程序，就好像所有设备都直接连接到他们的进程一样，即使实际上它们甚至可能跨越多个数据中心。得益于 Pathways：(1) 可以缩短作业启动时间，(2) 更容易实现容错，(3)多租户成为一种可行的选择，可以同时执行多个作业，从而提高硬件利用率。Pathways 能够轻松实现跨多个 TPU pod 的计算，这一点至关重要，因为它让我们避免了未来的扩展瓶颈。
路径概览。左上 ：以有向无环图表示的分布式计算。右上： 资源管理器为每个编译函数（例如 A、B 和 C）分配加速器网格的虚拟切片。底部 ： 用于分组调度计算的集中式调度程序，然后由每个分片执行器进行调度。（详情请参阅论文。）
另一个值得注意的版本是TensorStore，这是一个用于多维数组存储的新库。TensorStore 对于训练具有多控制器运行时的大型语言模型 (LLM) 特别有用，其中每个进程仅管理所有参数的子集，所有参数都必须整理到一致的检查点中。TensorStore 提供数据库级保证 ( ACID )，可高效、并发地将多维数组序列化到许多存储后端（例如，Google Cloud Storage、各种文件系统、HTTP 服务器），并已成功用于计算密集型工作负载，例如PaLM和人类皮层和果蝇大脑的重建。
苍蝇大脑重建，使用 TensorStore可以轻松访问和操作其底层数据。
顶部
机器学习的编程语言
我们的技术基础设施的稳健性和正确性对于机器学习工作至关重要，这就是为什么我们仍然致力于确保它建立在可靠的技术和理论基础上，并以编程语言和编译器构建的前沿研究为后盾。
我们继续投资开源MLIR编译器基础架构，构建更易于控制、可组合和模块化的编译器堆栈。此外，稀疏线性代数的代码生成也取得了很大进展，现在可以从几乎相同的 MLIR 程序生成密集和稀疏代码。最后，我们还继续开发IREE编译器，为数据中心的强大计算机和智能手机等 移动设备做好准备。
在理论层面，我们探索了形式化和验证代码生成技术的方法。我们还发布了一种用于实现和形式化自动微分(AD) 系统的新方法，这是 ML 库的核心。我们将反向模式 AD 算法分解为三个独立的程序转换，这显著简化了验证过程，凸显了 JAX 实现的独特功能。
利用抽象解释和程序合成 等编程语言技术，我们成功减少了执行神经架构搜索(NAS)所需的资源数量。这项努力，𝛼NAS，促成了更高效的模型的发现，同时又不降低准确性。
在过去的一年中，我们在 JAX 生态系统中发布了许多新的开源库，Rax和T5X只是其中两个例子。随着围绕jax2tf 的持续努力，现在可以使用 TensorFlow Lite 在移动设备上部署 JAX 模型，并使用 TensorFlow.js 在 Web 上部署 JAX 模型。
顶部
硬件加速器和机器学习
机器学习的硬件设计
使用定制硬件（例如 TPU 和 GPU）在性能提升和能效方面都表现出了巨大的优势（从而减少了碳足迹）。在最近的 MLPerf 竞赛中，我们在TPU v4的五个基准测试中创下了新的性能记录，平均比下一个最快的提交速度高出 1.42 倍。但是，为了跟上最近的进展，我们还在为特定的流行模型开发定制的硬件架构。
在所有五个已发布的基准测试（MLPerf 2.0 ）中， TPU都比最快的非 Google 提交（NVIDIA 本地）表现出显著的加速。条形越高越好。条形内的数字代表每次提交使用的芯片/加速器的数量。
但是，构建新的硬件加速器会产生高昂的初始成本，并且需要大量的开发和部署时间。为了使单工作负载加速器可行，必须缩短设计周期。 全栈搜索技术(FAST) 通过引入硬件加速器搜索框架解决了这个问题，该框架可以同时优化数据路径、调度和重要的编译器决策。FAST 引入了一个近似模板，能够描述不同类型的架构和多功能内存层次结构，从而使加速器与TPU v3 相比，将单位热设计功率的单工作负载性能（已知与单位总拥有成本的性能高度相关）提高了 3.7 倍。这表明单工作负载加速器对于中型数据中心部署来说是实用的。
用于硬件设计的机器学习
为了尽可能地实现芯片设计过程的自动化，我们不断推动 ML 在硬件设计各个阶段的能力，包括高级架构探索、验证以及布局和布线。
我们最近开源了一个名为Circuit Training的分布式 RL 基础设施，以及我们最近的《Nature》论文中描述的电路环境。我们在生产中使用了这套基础设施来为最新一代 TPU 芯片进行宏布局。在架构探索方面，PRIME引入了一种基于 ML 的方法来搜索硬件设计空间，该方法仅使用现有数据（例如来自传统加速器设计工作的数据），而无需任何进一步的硬件模拟。即使目标应用程序集发生变化，这种方法也无需运行耗时的模拟。与最先进的模拟驱动方法相比，PRIME 的性能提高了约 1.2 到 1.5 倍，同时将模拟时间缩短了 93% 到 99%。AutoApprox通过将每个神经网络层映射到适当的近似级别，自动生成近似的低功耗深度学习加速器，而不会造成任何准确性损失。
PRIME使用记录的加速器数据（包括可行和不可行加速器）来训练保守模型，该模型用于在满足设计约束的同时设计加速器。PRIME 设计的加速器的延迟最多可减少 1.5 倍，同时将所需的硬件模拟时间减少多达 99%。
硬件相关的模型设计
虽然NAS在准确度和效率方面表现出了发现最先进模型的强大能力，但它仍然受到硬件知识缺乏的限制。平台感知 NAS通过将硬件架构知识纳入 NAS 搜索空间的设计来解决这一差距。由此产生的 EfficientNet-X 模型比 TPU v3 和 GPU v100 上的 EfficientNet 快 1.5 到 2 倍，准确度相似。平台感知 NAS 和 EfficientNet-X 都已部署到生产中，显示出各种生产视觉模型的准确度显着提高，效率提高高达约 40%。NaaS通过同时搜索神经网络架构和硬件架构，走得更远。在 Edge TPU 上使用这种方法， NaaS发现的视觉模型在相同准确度下能效提高 2 倍。
TPU/GPU 上的平台感知 NAS概述，重点介绍搜索空间和搜索目标。
顶部
用于在受限搜索空间中导航的机器学习
除了改变硬件和工作负载以提高效率外，我们还可以优化中间层，包括将工作负载映射到多个设备上的分区器，以及将工作负载转换为硬件可以理解的低级表示的编译器。在前几年，我们展示了如何应用机器学习来找到更好的设备布局和编译器决策。在过去的一年里，我们进一步探索了这个方向，发现许多优化搜索空间受到严重限制，其中有效的解决方案非常稀疏。
为了应对这一挑战，我们 开发了几种技术，使学习模型能够有效地导航受限的搜索空间。Telamalloc采用 ML 模型和启发式方法的组合来在有多个选项可用时做出决策，并利用约束求解器推断进一步的依赖决策。与生产整数线性规划方法相比，Telamalloc 加快了 Edge TPU 编译器中的内存分配过程，并支持原本无法支持的重要真实模型。
“一种在多芯片模块上划分机器学习模型的可迁移方法”提出了一种略有不同的方法。它应用强化学习 (RL) 在一个步骤中提出决策，并要求约束求解器调整所提出的解决方案以使其有效。对于基于 Edge TPU 的多芯片网格上的 BERT 模型，与非学习搜索策略相比，这种方法使用更少的时间预算发现模型在设备之间的更好分布。
顶部
用于大规模生产系统的机器学习
我们还部署了机器学习来提高生产中运行的各种大型系统的效率。我们最近发布了MLGO ，这是第一个在LLVM基础架构中系统地集成机器学习技术的工业级通用框架。MLGO 可以用强化学习策略取代 LLVM 中的启发式方法来做出优化决策。在对一组内部大型应用程序进行测试时，我们发现经过训练的策略在优化内联决策时可以将二进制大小减少 3%–7%，在优化寄存器分配决策时可以将吞吐量提高 0.3% ~1.5%。在我们的生产 ML 编译器中，最近部署了 XLA（几年前发布的学习成本模型），用于指导为顶级 ML 工作负载选择最佳的 TPU 内核图块大小，从而节省了我们数据中心内约 2% 的总 TPU 计算时间。我们最近还用一种新的混合算法取代了 YouTube 缓存替换算法中现有的启发式算法，该算法将简单的启发式算法与学习模型相结合，将峰值时的字节缺失率提高了约 9%。
函数内联期间MLGO的说明。“#bbs”、“#users”和“callsite height”是调用者-被调用者对特征的示例。
顶部
人工智能与可持续发展
鉴于全球气候变化危机，人们对机器学习对环境的影响感到担忧是可以理解的。在最近的一篇论文中，我们表明，通过遵循最佳实践，机器学习从业者可以将训练产生的二氧化碳当量排放量 (CO2e) 减少几个数量级。我们将这些实践称为“4M”
模型。第一步是选择最有效的 ML 模型架构。例如，Primer在相同硬件上的运行速度比四年前开发的 流行Transformer快约 4 倍，同时获得相同的质量分数。
机器。第二种做法是使用最节能的计算机。例如，当 Transformer 模型于 2017 年首次发布时，流行的 GPU 是 Nvidia P100。使用针对 ML 训练优化的最新处理器（例如 TPU v4），可将每瓦性能提高约 15 倍。
机械化。用于培训的计算机需要安置在数据中心。大型云数据中心的能源效率通常比典型的小型本地数据中心高出约 1.4 倍。
地图。我们调查中最大的惊喜是选择最佳位置对能源供应清洁度的影响。此外，在云端，位置是四个因素中最容易改变的。典型位置和精心挑选的位置之间的差异可能约为 9 倍，即使在同一个国家/地区也是如此。
在此示例中，将 4M 相乘，按照使用 2017 年的 GPU 对原始 Transformer 模型进行训练的最佳实践，可在四年内将二氧化碳当量减少 4x × 15x × 1.4x × 9x 或约 750 倍。
我们将继续探索这一领域，并于 2023 年发布进一步的研究，展示如何通过仔细选择机器、机械化和训练地点，将当前模型训练的二氧化碳当量减少 20 倍。
总结
随着机器学习领域的发展，我们继续投资开发高性能、节能且易于使用的系统和基础设施，以便快速探索新想法。同时，我们继续探索机器学习在提高复杂系统性能和自动化系统设计中劳动密集型任务方面的能力。
评论