[Paper Review] Whale: A Unified Distributed Training Framework.

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

Advanced Neural Network Applications39 references4 citations

TL;DR

Whale is a unified distributed training framework that supports diverse parallel strategies—data, model, pipeline, operator sharding, and hybrid—through a novel IR abstraction and automatic parallelism via a meta-driven cost model. It accelerates BERT-Large training by 2.32× and large-scale image classification by 14.8× over Horovod data parallelism on 64 GPUs, while remaining compatible with TensorFlow with minimal code changes.

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.

Motivation & Objective

• Address the limitations of data parallelism in scaling to large models and datasets.
• Enable efficient training of giant models whose parameters exceed single-GPU memory capacity.
• Provide a unified framework supporting diverse parallel strategies within a single system.
• Automate parallel strategy selection using a meta-driven cost model to optimize training performance.
• Ensure backward compatibility with existing TensorFlow models with minimal code modifications.

Proposed method

• Design Whale IR as a foundational abstraction to express and compose complex distributed training strategies.
• Implement comprehensive parallelism strategies: data parallelism, model parallelism, operator sharding, pipeline parallelism, and hybrid combinations.
• Integrate an automatic parallelism engine that selects optimal strategies using a meta-driven cost model.
• Ensure seamless integration with TensorFlow by requiring only minimal code changes to deploy distributed training.
• Use a unified IR-based compilation pipeline to map high-level models to efficient distributed execution plans.

Experimental results

Research questions

• RQ1 How can a single framework efficiently support a wide range of distributed training strategies, including hybrid combinations?
• RQ2 To what extent can automatic strategy selection based on a cost model improve training efficiency compared to manual or fixed strategy choices?
• RQ3 Can the framework scale to large models like BERT-Large and 100,000-class image classification with significant speedups over existing data-parallel baselines?
• RQ4 How does the framework maintain compatibility with existing deep learning frameworks like TensorFlow without requiring model code changes?
• RQ5 What performance gains are achievable through operator sharding and hybrid strategies in large-scale training workloads?

Key findings

• Whale’s pipeline strategy achieves 2.32× speedup over Horovod data parallelism on 64 GPUs for BERT-Large training.
• Whale’s hybrid strategy, combining operator sharding and data parallelism, delivers 14.8× speedup over Horovod data parallelism on 64 GPUs for a 100,000-class image classification task.
• The framework successfully supports models whose parameter size exceeds single-GPU memory capacity by leveraging model and operator sharding.
• Automatic parallelism via the meta-driven cost model enables optimal strategy selection without manual tuning.
• Whale maintains full compatibility with TensorFlow, allowing distributed training with minimal code changes to existing models.