Jeff Dean发推：谷歌超硬年末总结“第三弹”来了！鼎力大举发展Jax

新智元报道

编辑：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的论文中具体做了先容。

谷歌在生产中使用了这个根本办法，为最新一代的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