Automap: Towards Ergonomic Automated Parallelism for ML Models

TL;DR

Automap is an automated partitioning tool that integrates with compilers to simplify and optimize the process of parallelizing large neural network training, reducing the need for expert intervention.

Contribution

It introduces a platform-independent, automated partitioner capable of recovering expert-level strategies like Megatron sharding for transformers.

Findings

Automap successfully automates complex partitioning strategies.

It integrates seamlessly into existing workflows.

It achieves expert-level performance in partitioning neural networks.

Abstract

The rapid rise in demand for training large neural network architectures has brought into focus the need for partitioning strategies, for example by using data, model, or pipeline parallelism. Implementing these methods is increasingly supported through program primitives, but identifying efficient partitioning strategies requires expensive experimentation and expertise. We present the prototype of an automated partitioner that seamlessly integrates into existing compilers and existing user workflows. Our partitioner enables SPMD-style parallelism that encompasses data parallelism and parameter/activation sharding. Through a combination of inductive tactics and search in a platform-independent partitioning IR, automap can recover expert partitioning strategies such as Megatron sharding for transformer layers.