Auto-Parallelizing Large Models with Rhino: A Systematic Approach on Production AI Platform

TL;DR

Rhino is a system that automatically parallelizes tensor programs for large models on AI platforms, enabling scalable distributed execution without user configuration, and discovering superior strategies for diverse models.

Contribution

It introduces a systematic approach to parallelization that generalizes across applications and explores a comprehensive strategy space with efficient search heuristics.

Findings

Rhino can replicate expert-crafted strategies for existing models.

It discovers novel parallelization strategies surpassing current systems.

Demonstrates scalability to thousands of devices in production environments.

Abstract

We present Rhino, a system for accelerating tensor programs with automatic parallelization on AI platform for real production environment. It transforms a tensor program written for a single device into an equivalent distributed program that is capable of scaling up to thousands of devices with no user configuration. Rhino firstly works on a semantically independent intermediate representation of tensor programs, which facilitates its generalization to unprecedented applications. Additionally, it implements a task-oriented controller and a distributed runtime for optimal performance. Rhino explores on a complete and systematic parallelization strategy space that comprises all the paradigms commonly employed in deep learning (DL), in addition to strided partitioning and pipeline parallelism on non-linear models. Aiming to efficiently search for a near-optimal parallel execution plan, our analysis of production clusters reveals general heuristics to speed up the strategy search. On top of it, two optimization levels are designed to offer users flexible trade-offs between the search time and strategy quality. Our experiments demonstrate that Rhino can not only re-discover the expert-crafted strategies of classic, research and production DL models, but also identify novel parallelization strategies which surpass existing systems for novel models.