[论文解读] Whale: A Unified Distributed Training Framework.

Ang Wang, Xianyan Jia+4 more|arXiv (Cornell University)|Nov 18, 2020

Advanced Neural Network Applications参考文献 39被引用 4

一句话总结

Whale 是一个统一的分布式训练框架，通过一种新颖的中间表示（IR）抽象和基于元驱动成本模型的自动并行化，支持多种并行策略——数据并行、模型并行、流水线并行、算子分片以及混合策略。在 64 个 GPU 上，与 Horovod 的数据并行相比，Whale 将 BERT-Large 训练速度提升 2.32 倍，将大规模图像分类训练速度提升 14.8 倍，同时保持与 TensorFlow 的高度兼容性，仅需极少的代码修改。

ABSTRACT

Data parallelism (DP) has been a common practice to speed up the training workloads for a long time. However, with the increase of data size and model size, DP has become less optimal for most distributed training workloads. Moreover, it does not work on models whose parameter size cannot fit into a single GPU's device memory. To enable and further improve the industrial-level giant model training, we present Whale, a unified distributed training framework. It provides comprehensive parallel strategies including data parallelism, model parallelism, operator sharding, pipeline, hybrid strategy, and automatic parallel strategy. To express complex training strategies effectively and efficiently in one framework, Whale IR is designed as the basic unit to explore and implement different distributed strategies. Moreover, Whale enables automatic parallelism upon using a meta-driven cost model. Whale is compatible with TensorFlow and can easily distribute training tasks by adding a few code lines without changing user model code. To the best of our knowledge, Whale is the first work that can support various hybrid distributed strategies within one framework. In our experiment of Bert Large model, Whale pipeline strategy is 2.32 times faster than Horovod data parallelism (HDP) on 64 GPUs. In a large-scale image classification task (100,000 classes), Whale hybrid strategy, which consists of operator sharding and DP, is 14.8 times faster than HDP on 64 GPUs.

研究动机与目标

• 解决数据并行在扩展至大模型和大规模数据集时的局限性。
• 实现参数量超过单张 GPU 显存容量的巨型模型的高效训练。
• 提供一个统一的框架，在单一系统内支持多种并行策略。
• 通过基于元驱动成本模型的自动并行策略选择，优化训练性能。
• 确保与现有 TensorFlow 模型的向后兼容性，仅需极少的代码修改。

提出的方法

• 设计 Whale IR 作为基础抽象，用于表达和组合复杂的分布式训练策略。
• 实现全面的并行化策略：数据并行、模型并行、算子分片、流水线并行以及混合组合。
• 集成自动并行化引擎，利用元驱动成本模型选择最优策略。
• 通过仅需极少代码修改即可将分布式训练无缝集成到 TensorFlow 中。
• 使用基于统一 IR 的编译管道，将高层模型映射为高效的分布式执行计划。

实验结果

研究问题

• RQ1 如何通过单一框架高效支持广泛的分布式训练策略，包括混合组合？
• RQ2 基于成本模型的自动策略选择相比手动或固定策略选择，在提升训练效率方面能达到何种程度？
• RQ3 该框架能否在大规模模型（如 BERT-Large 和 100,000 类图像分类）上实现显著的性能加速，超越现有的数据并行基线？
• RQ4 该框架如何在无需修改模型代码的前提下，保持与 TensorFlow 等现有深度学习框架的兼容性？
• RQ5 在大规模训练工作负载中，算子分片和混合策略能带来多大的性能提升？

主要发现

• Whale 的流水线策略在 64 个 GPU 上对 BERT-Large 训练实现 2.32 倍于 Horovod 数据并行的加速。
• Whale 的混合策略（结合算子分片与数据并行）在 64 个 GPU 上对 100,000 类图像分类任务实现 14.8 倍于 Horovod 数据并行的加速。
• 该框架通过利用模型并行和算子分片，成功支持参数量超过单张 GPU 显存容量的模型。
• 基于元驱动成本模型的自动并行化可实现最优策略选择，无需手动调优。
• Whale 与 TensorFlow 完全兼容，使现有模型仅需极少代码修改即可启用分布式训练。