[论文解读] Scale MLPerf-0.6 models on Google TPU-v3 Pods

Sameer Kumar, Victor Bitorff+10 more|arXiv (Cornell University)|Sep 21, 2019

Advanced Data Storage Technologies参考文献 11被引用 35

一句话总结

本文展示了将 MLPerf-0.6 训练基准扩展到 Google TPU-v3 Pods（1024 芯片），通过分布式评估、梯度求和加速、模型并行和权重更新分片等优化，在 ResNet-50、SSD 和 Transformer 上达到创纪录的性能。

ABSTRACT

The recent submission of Google TPU-v3 Pods to the industry wide MLPerf v0.6 training benchmark demonstrates the scalability of a suite of industry relevant ML models. MLPerf defines a suite of models, datasets and rules to follow when benchmarking to ensure results are comparable across hardware, frameworks and companies. Using this suite of models, we discuss the optimizations and techniques including choice of optimizer, spatial partitioning and weight update sharding necessary to scale to 1024 TPU chips. Furthermore, we identify properties of models that make scaling them challenging, such as limited data parallelism and unscaled weights. These optimizations contribute to record performance in transformer, Resnet-50 and SSD in the Google MLPerf-0.6 submission.

研究动机与目标

• 证明在 TPU-v3 Pods（1024 芯片）上对 MLPerf-0.6 进行可扩展训练，并在大规模下识别瓶颈。
• 开发并验证提高吞吐量、保持目标精度的优化技术。
• 展示模型并行与编译器辅助优化，以实现对多样化 ML 工作负载的扩展。

提出的方法

• 使用 XLA 编译器降级 TensorFlow 图以优化 TPU-v3 的执行。
• 在计算中使用混合精度（bfloat16），同时对关键非卷积运算保持 32 位以维持精度。
• 应用一系列优化：分布式评估、二维梯度求和，以及权重更新分片，以减少大规模 Pod 上的瓶颈。
• 采用模型并行技术（空间划分和 halo 交换）以提高各核心之间的并发性。
• 对 Transformer 与 GNMT 使用数据并行和就地评估以降低基础设施开销。
• 降低评估开销并优化输入数据管道，以维持高吞吐量。

实验结果

研究问题

• RQ1 当在 Google TPU-v3 Pods（最高 1024 芯片）上训练时，MLPerf-0.6 基准可以如何扩展？
• RQ2 在大规模下保持 MLPerf 精度目标的前提下，达成峰值吞吐量需要哪些优化？
• RQ3 哪些模型特定的并行和优化策略最能实现 ResNet-50、SSD、Mask-RCNN、Transformer 与 GNMT 的大规模扩展？

主要发现

• 在 TPU-v3 Pods 上实现高达 1024 芯片的大规模扩展，并在若干基准测试中取得创纪录的性能。
• 将评估分布化并将训练/评估循环紧密耦合，降低瓶颈并提高可扩展性。
• 采用带流水线数据传输的二维梯度求和可带来显著加速（ResNet-50 超过 1.5x）。
• 模型并行技术（空间划分和权重更新分片）使得较小的每核心工作负载和具有大参数量的 Transformer 得以扩展。
• SSD 在 4 个核心上实现 1.6x 加速，扩展至 2048 TPU 核心成为可能，Mask-RCNN 的扩展受益于前后阶段的多核分布。
• Transformer 与 GNMT 的扩展需要专门的优化，包括就地评估、减少最大序列长度，以及内存/计算优化以在大批量下达到目标精度。