Jeff Dean发推：谷歌超硬年终总结「第三弹」来了！大力发展Jax

公众号新闻

2023-02-03 05:02

新智元报道

编辑：Aeneas 好困

【新智元导读】谷歌费心打造的年终总结第三弹，火热出炉了！

刚刚，Jeff Dean发推表示，谷歌重磅打造的超级硬核年终大总结，出第三弹了！

第一弹：「超详超硬Jeff Dean万字总结火热出炉！图解谷歌2022年AIGC、LLM、CV三大领域成就」

第二弹：「谷歌2022年度回顾：让AI更负责任，主要做了4点微小的工作」

伟大的机器学习研究需要伟大的系统。

随着算法和硬件越来越复杂，以及运行规模越来越大，执行日常任务所需的软件的复杂性也在不断增加。

在这篇文章中，研究人员概述了过去一年整个谷歌在ML系统方面取得的众多进展，这些进展使谷歌能够支持复杂模型的服务和训练，同时减轻了终端用户的实施复杂性。

同时，这篇文章还提到了谷歌如何利用ML本身来改进和设计下一代系统堆栈的研究。

机器学习编程语言

对于机器学习的工作，基础架构的稳健性和正确性至关重要。

谷歌一直在努力，确保基础架构建立在可靠的技术和理论基础之上。并且，作为后盾，谷歌一直在做编程语言和构建编译器方面的前沿研究。

谷歌会继续对开源MLIR编译器的基础架构投资，构建更加可控、可组合和模块化的编译器堆栈。

论文地址：https://research.google/pubs/pub49988/

此外，谷歌在稀疏线性代数的代码生成方面也取得了很大进展，现在可以从几乎相同的MLIR程序中生成密集和稀疏的代码。

最后，谷歌还继续开发了IREE编译器，这个编译器既可以在位于数据中心的强大计算机上使用，在可以在智能手机之类的移动设备上使用。

IREE的端到端流程

在更理论的层面，谷歌探索了哪些方法可以形式化（formalize）和验证自己使用的代码生成技术。

谷歌还发布了一种新颖的方法，用于执行和形式化一套自动微分（AD）系统，它正是ML库的核心。

源代码转换

谷歌将反向模式的AD算法分解成三个独立的程序转换，这就变得更简单，更容易验证，从而突出了JAX实现的独特性。

反向模式自动微分作为正向微分、解压缩和转置

利用抽象解释和程序合成等编程语言技术，谷歌成功地减少了进行神经结构搜索（NAS）所需的资源数量。这项𝛼NAS成果，可以让我们在不降低准确性的前提下，发现了更有效的模型。

在用于图像分类的视觉Transformer架构演化过程中由𝛼NAS合成的突变

在过去的一年里，谷歌在JAX生态系统中发布了许多新的开源库，比如Rax和T5X。

随着围绕jax2tf的持续努力，JAX模型现在可以使用TensorFlow Lite部署在移动设备上，并使用TensorFlow.js部署在网络上。

「Plane Strike」中的演示

用于机器学习的分布式系统

2022年，谷歌在更好地支持ML和通用科学计算进行大规模计算方面取得了重大进展。

不仅为大型模型设计了SOTA的服务技术，改进了张量程序的自动分区，而且还重新设计了库的API，以确保所有这些发展能够被广大用户所接受。

其中最大的改进之一，便是用于评估大规模矩阵乘法运算的CollectiveEinsum策略，这是神经网络的核心。

论文地址：https://dl.acm.org/doi/abs/10.1145/3567955.3567959

与之前流行的SPMD分区策略不同，CollectiveEinsum会将通信与设备本地计算分开，并通过快速的TPU ICI链接进行叠加，进而使性能提高了1.38倍。

同时，CollectiveEinsum算法也是谷歌扩展Transformer推理工作的一个关键组成部分。比如，在吞吐量优化的配置中达到SOTA模型76%的FLOPs利用率（MFU）。

此外，谷歌还将SPMD风格的分区概念整合进了TensorFlow（通过DTensor扩展）和JAX（通过重新设计的数组类型）。

在这两个库中，那些程序员看来是完整的张量，可以通过附加声明性的布局注释，在一些设备上透明地进行分片。

事实上，这两种方法不仅和为单设备计算编写的现有代码兼容，并且还可以扩展到多设备程序中，而不需要修改任何代码！

论文地址：https://arxiv.org/abs/2105.04663

然而，GSPMD在很大程度上依赖于启发式方法，也就是有时仍然需要手动做出决定，而这通常会让性能无法达到最优。

为了使分区推理完全自动化，谷歌开发了Alpa——一个它探索了运算器级（模型）并行和较大子计算之间管线并行策略的自动化系统。

Alpa不仅实现了在Transformer等主流模型上与「人工微调」相媲美的性能，同时也能够扩展到其他模型之中，如卷积网络和专家混合模型（MOE）。

与之类似，谷歌最近提出的Pathways系统，在TPU运行时间之上增加了一个额外的虚拟化层——加速器由长期存在的进程管理，而不是直接分配给用户。

然后，单个终端用户可以连接到任意数量的Pathways控制的设备，并编写他们的程序。就像所有的设备都直接连接到他们的进程一样，即使现实中的情况是跨越多个数据中心的。

论文地址：https://arxiv.org/abs/2203.12533

由于Pathways：（1）作业启动时间减少，（2）更容易实现容错，以及（3）使多租户成为一个可行的选择，从而让多个作业可以同时执行，更有效地利用硬件。

更重要的是，Pathways能够轻松实现跨越多个TPU pods的计算，而这可以有效避免未来的扩展瓶颈。

左上：用有向无环图表征的分布式计算；右上：资源管理器为每个编译的函数（如A、B和C）分配虚拟的加速器网格片；下：集中的调度器对计算进行分组调度，然后由每个分片的执行器进行调度

此外，还有一个全新的用于多维阵列存储的库——TensorStore。

TensorStore在训练具有多控制器运行时间的大型语言模型（LLM）时非常实用，其中每个进程只用管理参数的一个子集，而所有的参数则需要被整理成一个一致的检查点。

TensorStore为高效和并发的多维数组序列化提供了数据库级的保证（ACID），并已成功用于计算密集型工作负载，如PaLM和人类皮层和果蝇大脑的重建。

一个苍蝇大脑的重建，其基础数据可以使用TensorStore轻松访问和操作

硬件加速器和机器学习

用于ML的硬件设计

使用定制的硬件（如TPU和GPU），在性能提升和能源效率上会有巨大的优势，还能减少碳足迹。

在最近的MLPerf竞赛中，谷歌在TPU v4上的五项基准测试中创造了新的性能记录，实现了比第二名平均高1.42倍的速度。

不过，为了跟上最近的进展，谷歌也在为特定的流行模型开发定制的硬件架构。

在已公布的五个基准测试（MLPerf 2.0）中，谷歌的TPU都比竞品（NVIDIA on-premises）速度更快。（条形图内的数字代表使用的芯片/加速器的数量）

然而，构建新的硬件加速器会产生很高的初始成本，并且需要大量的开发和部署时间。

为了使单工作负载加速器（single-workload accelerators）可行，必须减少设计周期时间。

全栈加速器搜索技术

而全栈搜索技术（FAST）通过引入一个硬件加速器搜索框架，就解决了这个问题。

这个框架同时优化了数据路径、调度和重要的编译器决策。

FAST引入了一个近似的模板，能够描述不同类型的架构和多功能的内存层次，从而使加速器的单位热设计功率（与单位总成本的性能高度相关）的单工作负载性能比TPU v3提高3.7倍。

这表明，单工作负载加速器对于中等规模的数据中心部署是实用的。

用于硬件设计的机器学习

为了尽可能地实现芯片设计过程的自动化，谷歌在硬件设计的各个阶段，都在推动机器学习的功能，包括高级架构探索、验证以及布局和布线。

方法和训练方案概述

谷歌最近开源了一个名为Circuit Training的分布式RL基础设施，以及一个电路环境，后者谷歌在发于Nature的论文中详细做了介绍。

论文地址：https://www.nature.com/articles/s41586-021-03544-w

谷歌在生产中使用了这个基础设施，为最新一代的TPU芯片生成了宏观布局。

在解决架构探索问题时，PRIME引入了一种基于ML的方法来搜索硬件设计空间，只利用现有的数据（比如来自传统加速器设计工作的数据），而不需要进一步的硬件模拟。

这种方法减轻了运行耗时的模拟的需要，即使在目标应用程序集发生变化时。

PRIME比最先进的模拟驱动方法提高了约1.2-1.5倍的性能，同时减少了93%-99%的模拟时间。

AutoApprox通过将每个神经网络层映射到适当的近似级别，自动生成近似的低功耗深度学习加速器，而没有任何精度损失。

PRIME使用记录的加速器数据（包括可行的和不可行的加速器）来训练模型，其设计的加速器的延迟小了1.5倍，同时减少了99%的硬件模拟时间

依赖于硬件的模型设计

虽然神经架构搜索（NAS）在SOTA模型的发现方面展示出了巨大的能力，但它仍然受到缺乏硬件知识的限制。

而基于平台感知（Platform-aware）的NAS，则可以通过将硬件结构的知识纳入NAS搜索空间的设计中，来解决这一问题。

由此产生的EfficientNet-X模型在TPU v3和GPU v100上的速度分别是EfficientNet的1.5倍-2倍，而精度却相差无几。

目前，平台感知的NAS和EfficientNet-X都已在生产中部署。实践证明，对于各种生产型视觉模型来说，都有明显的精度提升和高达40%的效率提升。

论文地址：https://arxiv.org/abs/2102.05610

NaaS通过共同搜索神经网络架构和硬件架构，更进一步。

测试结果显示，NaaS可以在Edge TPU上发现同等精度但能效提高了2倍的视觉模型。

在TPU/GPU上的平台感知NAS概述

用于大规模生产系统的机器学习

在生产中运行的各种大规模系统上，谷歌也利用机器学习实现了效率的提升。

比如，最近发布的第一个在LLVM基础设施中系统地整合ML技术的工业级通用框架——MLGO，可以用RL策略取代LLVM中的启发式方法来做出优化决策。

测试发现，在优化内联决策时，经过训练的策略可以减少3%-7%的二进制大小，而在优化寄存器分配决策时，可以提高0.3%~1.5%的吞吐量。

论文地址：https://arxiv.org/abs/2101.04808

在生产型ML编译器中，几年前发布的学习成本模型XLA，也被用于指导顶级ML工作负载的TPU内核的最佳瓦片大小的选择，进而在数据中心上节省了2%的TPU总计算时间。

论文地址：https://arxiv.org/abs/2008.01040

此外，谷歌还用新的混合算法取代了YouTube缓存替换算法中现有的启发式算法，该算法结合了简单的启发式算法和学习模型，在峰值时将byte miss提高了9%。

总结一下

谷歌表示，随着机器学习领域的发展，自己将在开发高性能、高能效和易于使用的系统和基础设施上持续投入，进而实现对新想法的快速探索。

同时，谷歌也会继续探索机器学习的能力、提高复杂系统的性能，并使系统设计中的劳动密集型任务自动化。

参考资料：

https://ai.googleblog.com/2023/02/google-research-2022-beyond-ml-computer.html#Theme1

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。

来源: qq

点击查看作者最近其他文章