# Memory-efficient array redistribution through portable collective communication

arXiv: https://arxiv.org/abs/2112.01075 · HTML: https://mlsanity.com/paper/2112.01075/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 34 references that are in the map.
- First posted on arXiv 2 Dec 2021.

## Builds on (most cited first, 9 of 9)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [HeAT -- a Distributed and GPU-accelerated Tensor Framework for Data Analytics](https://mlsanity.com/paper/2007.13552/) (2020)

## Built on by (most cited first, 3 of 3)

- [On Optimizing the Communication of Model Parallelism](https://mlsanity.com/paper/2211.05322/) (2022)
- [SuperScaler: Supporting Flexible DNN Parallelization via a Unified Abstraction](https://mlsanity.com/paper/2301.08984/) (2023)
- [Diffusion-Based Approaches in Medical Image Generation and Analysis](https://mlsanity.com/paper/2412.16860/) (2024)

## Abstract (arXiv)

Modern large-scale deep learning workloads highlight the need for parallel execution across many devices in order to fit model data into hardware accelerator memories. In these settings, array redistribution may be required during a computation, but can also become a bottleneck if not done efficiently. In this paper we address the problem of redistributing multi-dimensional array data in SPMD computations, the most prevalent form of parallelism in deep learning. We present a type-directed approach to synthesizing array redistributions as sequences of MPI-style collective operations. We prove formally that our synthesized redistributions are memory-efficient and perform no excessive data transfers. Array redistribution for SPMD computations using collective operations has also been implemented in the context of the XLA SPMD partitioner, a production-grade tool for partitioning programs across accelerator systems. We evaluate our approach against the XLA implementation and find that our approach delivers a geometric mean speedup of 1.22×, with maximum speedups as a high as 5.7×, while offering provable memory guarantees, making our system particularly appealing for large-scale models.