选自arxiv，作者：Keno Fischer, Elliot Saba，设备之心编辑部。
Julia 语言发展非常快速，它能够视为同期具备了 Python 的灵活性与 C 的速度，但日前 TensorFlow 和 PyTorch 等框架官方都不支持 Julia 语言。因此呢近期有科研者借助 XLA 底层编译器为 Julia 构建 TPU 支持，她们暗示该办法能够将 Julia 程序编写的 VGG19 模型融合到 TPU 可执行文件中，并调用 TPU 实现有效计算。而 Google.ai 的负责人 Jeff Dean 在推特上亦暗示「Julia + TPUs = fast and easily expressible ML computations!」
1. 引言
过去的几年里推动设备学习技术稳步发展的基本性改变之一是训练和优化设备学习模型的巨大计算力。许多技术都是很年前就已然提出，唯有近几年提高的计算力能够为现实世界的问题供给足够优秀的处理方法。这些计算能力的很大一部分是经过 GPU 获取的，其针对向量的计算能力最初是为图形而设计的，但设备学习模型一般需要执行繁杂的矩阵运算，因此呢 GPU 一样表现出了非常好的性能。
这些办法及 GPU 在现实世界，尤其是在设备学习行业的成功诱发了硬件设计者的一系列创新，她们致力于为设备学习工作负载开发新的加速器。然而，尽管 GPU 很长一段时间都在 CUDA 等软件系统发力，但这些库一般不会扩展到新的非 GPU 加速器，为这些加速器研发软件仍然是一大挑战。
2017 年，谷歌宣布她们将经过云服务向公众供给她们专有的 TPU 设备学习加速器。最初，TPU 的运用局限于按照谷歌 TensorFlow 设备学习框架编写的应用。幸运的是，2018 年 9 月，谷歌经过底层 XLA（Accelerated Linear Algebra）编译器的 IR 开放了 TPU 的拜访权限。这个 IR 是一个通用的优化编译器，用于表达线性代数基元的任意计算，因此呢为运用 TPU 的非 TensorFlow 用户以及非设备学习工作负载供给了良好的基本。
在本文中，咱们介绍了运用这个接口编译通用 Julia 代码的初步工作，它们能够进一步拜访谷歌云的 TPU。这一办法与 TensorFlow（Abadi et al., 2016）采用的办法形成对比，后者无编译 python 代码，而是先用 Python 构建一个计算图，而后再对这个计算图进行编译。它在美学上类似于 JAX（Frostig et al., 2018），JAX 的目的是经过跟踪和 Offload 高级数组运算来 Offload Python 本身编写的计算。然而重要的是，咱们不依赖于跟踪，而是利用 Julia 的静态分析和编译能力来编译全部程序，包含传递到设备端的所有掌控流。
值得一提的是，咱们的办法准许用户在编写模型时充分利用 Julia 语言的表现力。这些表现力重点表现在有些高级特征上，如多重派发、高阶函数和现有库，如微分方程求解器（Rackauckas & Nie，2017）和通用线性代数例程等。因为仅在纯 Julia 代码上运行，因此它亦与 Zygote.jl（Innes, 2018）自动微分工具兼容，该工具能执行自动微分做为高级编译过程。总的来讲，咱们能够编译运用 Flux 设备学习框架编写的完整设备学习模型，将模型的前向、反向传播及训练回路融合成一个可执行文件，并 Offload 到 TPU 中。
论文：Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs
论文链接：
https://arxiv.org/abs/1810.09868
摘要：谷歌的云 TPU 是一种前景广阔的新型设备学习工作负载硬件架构，近年来已然成就了谷歌非常多里程碑式的设备学习突破。如今，谷歌已然在其云平台上为公众供给 TPU，近期又进一步开放，准许非 TensorFlow 前端运用。咱们描述了一种经过这一新 API 及谷歌 XLA 编译器将 Julia 程序的适当部分 Offload 到 TPU 的办法和实现。咱们的办法能够将 Julia 程序编写的 VGG19 模型及其正向传播完全融合到单个 TPU 可执行文件中，以便 Offload 到设备上。咱们的办法与 Julia 代码上现有的基于编译器的自动微分技术很好地结合在一块，因此呢亦能够自动得到 VGG19 反向传播并采用类似的办法将其 Offload 到 TPU。运用咱们的编译器拜访 TPU，咱们能够在 0.23 秒内完成批量为 100 张图像的 VGG19 前向传播，而 CPU 上的原始模型则需要 52.4s。咱们的实现仅需不到 1000 行的 Julia 代码，无需按照 TPU 对核心 Julia 编译器或任何其他 Julia 包进行特有的更改。
5. 将 Julia 语义映射到 XLA
只要 Julia 程序是根据 XLA 基元来编写的，咱们就能将其编译到 XLA。然而，Julia 程序不是按照晦涩难懂的 HLO 操作来编写的，而是按照由 Julia 基本库供给的函数和抽象来编写的。幸运的是，Julia 运用了多重派发，使得按照 HLO 操作来表达标准库的抽象变得容易。下面展示了几个简单的例子：
除了这些简单的操作以外，咱们还供给了高级数组抽象的实现，尤其是 mapreduce 和 broadcast。依据 HLO 操作实现的 broadcast 大大概 20 行代码，为节省空间起见，此处不予展开，但「mapreduce」的实现非常简单：
从上图能够看到将任意 Julia 函数做为静态计算运算的效果。因为 Julia 对泛型抽象的依赖，它只需指定极少许定义，就能覆盖海量 API。详细来讲，从 mapreduce 的定义中，咱们能够自动得到在 base 中所定义运算（如 sum 和 prod）的降维。事实上，获取足够的 API 覆盖来编译 VGG19 模型的前向传播和反向传播需要不到 200 行定义。
5.1 结构映射
咱们做了一个额外的识别。embedded IR 中的任意元组或 immutable 结构被映射至一个 XLA 元组，即 julia 值 1 + 2im（由两个整数结构构成的繁杂数字）将被映射至 XLA 元组 (s64[], s64[])。咱们在 XLA IR 的 Julia 嵌入中保留该结构类型，但很显然 XLA 不认识 julia 类型，因此呢在最后的转换过程中这些类型被转换成适当的元组。类似地，（julia）元组构造函数（以及 immutable 结构的构造函数）变成为了 XLA 的元组构件。元组引用（immutable 结构的字段引用）变成为了 XLA 的元组引用。
5.2 处理掌控流
有一个额外的繁杂问题咱们还没讨论：Julia 供给的命令式掌控流和 XLA 供给的函数式掌控流之间的语义不匹配。为认识决 if/else 掌控流模块，咱们在 Julia 编译器的 SSA IR 中查看 φ 节点，而后将这些节点做为 XLA 函数式掌控流的结果（倘若在同一个合并点存在多个 φ 节点，则咱们构造这些节点的元组）。引起计算流分化的要求变成为了函数式掌控流的要求，二者之间的任意计算都可做为函数调用。循环掌控流类似要求掌控流的构建，咱们识别掌控流图的强连接区域，将其做为循环的主体。
7 结果
7.2 VGG19 前向传播
咱们的第1个繁杂示例是完整 VGG19 前向传播。咱们运用 Metalhead 包中的 VGG19 实现 (Mike Innes & Contributors, 2018)，它利用 Flux (Innes & Contributors, 2017) 框架将熟练的设备学习层（卷积层、全连接层）转换成线性代数运算。但重要的是，Flux 框架中的每一个层只是通常的函数，它们能够反过来调用通常的线性代数运算。因此呢，Flux 中表达的设备学习模型（包含 VGG19）只是通常的 Julia 函数，因此呢能够运用本论文介绍的办法。
咱们的编译器能够完全推断、offload 和融合（fuse）VGG19 的所有前向传播。在 Julia 级别的优化之后，顶层函数的最后 IR 包含 181 个指令（每一个 HloOp 都是具备适当推断的常数静态参数和适当形态推断的动态参数）。每一个级别计算的 HLO operands 总数是 183（多出的两个用于嵌入中隐匿的参数指令），29 个计算一共有 361 个 HLO operands，指令数详情见图 3。因为咱们能够 offload 所有前向传播计算，因此呢 Julia 不参与任何评定过程，从而能够同步执行其他任务（如为下一准许备数据）。另外，得到代码的性能仅受限于 XLA 生成的代码质量，不受限于前端（性能评定见 7.4）。咱们在 ImageNet 验证集上评定了 VGG19 模型，并验证了得到结果与原版 Metalhead 的结果相匹配，从而验证了生成的 XLA 代码准确性。
7.3 VGG19 反向传播
为了获取反向传播，咱们利用基于 Zygote.jl 编译器的 AD 框架 (Innes, 2018)。Zygote 在 Julia 代码上运行，其输出亦是 Julia 函数（适合重新导入 Zygote 以获取更高阶的导数，亦适合编译成针对 TPU 的模型）。如下是一个详细示例：
即模型当前值和特定训练样本（或训练样本批）所对应的导数。咱们运用 sum 做为损失函数的简单替代。意外的是，第 6 章介绍的类型推断修改亦能够加强所有 VGG19 反向传播的类型推断精度。至于前向传播，优化和未优化的指令总数如图 1 所示。反向传播生成的 XLA 指令显著多于前向传播，其最大贡献者之一便是 Zygote 的混合模式广播融合（mixed mode broadcast fusion）——在一个映射内核（map kernel）中同期计算前向传播和反向传播。因为 XLA 日前不支持来自一个映射指令的多个输出，该函数在多个映射指令上重复运行，因此呢后续需要清洗 XLA 的 DCE。通常，咱们的编译过程处理了 XLA 对映射指令的处理，由于在泛型代码中调用 Julia 映射和 broadcast 函数非常广泛。
7.4 在 TPU 上进行评定
图 2：区别批体积对应的 VGG19 前向传播时长。Flux CPU 是 Flux master/Julia master，但不运用 XLA 编译器。PyTorch CPU 是同一 CPU 上的相同 PyTorch 模型。FluXLA CPU 是咱们的科研在 CPU 上的 xrt 实现；FluXLA TPU (total) 是端到端时间，和客户端报告的时间一致（包含 kernel launch 开销和从谷歌云把数据迁移回来，重视由于额外的网络迁移，该测绘结果会显现极重的变动）；FluXLA TPU (compute) 是 TPU 上的总计算时间，和云分析器报告的时间一致（与 FluXLA TPU (total) 区别，该测绘很稳定）。所有 CPU 测绘基于支持 AVX512 的 Intel(R) Xeon(R) Silver 4114 CPU @ 2.20GHz CPU。可获取高达 20 个内核，且 CPU 基准不限于单个内核（即使在实践中，亦不是所有 CPU 基准都运用并行化）。TPU 基准仅限单个 TPU 内核。所有时间最少经过 4 次运行（除了 FluXLA CPU for N=100，由于它没法在 10 分钟内完成一次运行）。
图 3：被编译为 XLA 后，Metalhead.jl VGG19 前向传播和反向传播的指令数分解，上图展示了未优化（Julia 前端之后）和优化指令数（XLA 优化流程之后，与 CPU 后端所用流程类似，但无 HLO fusion）。每一个指令数被进一步拆分为实体计算中的指令（E）和所有计算中的指令数（T）。
发表于 2024-7-28 04:42:57