# Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models

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

## Key facts

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [Scaling Distributed Machine Learning with In-Network Aggregation](https://mlsanity.com/paper/1903.06701/) (2019)
- [Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes](https://mlsanity.com/paper/1807.11205/) (2018)
- [Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect](https://mlsanity.com/paper/1903.04611/) (2019)
- [Blink: Fast and Generic Collectives for Distributed ML](https://mlsanity.com/paper/1910.04940/) (2019)
- [Synthesizing Optimal Collective Algorithms](https://mlsanity.com/paper/2008.08708/) (2020)
- [Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms](https://mlsanity.com/paper/2007.00156/) (2020)
- [Flare: Flexible In-Network Allreduce](https://mlsanity.com/paper/2106.15565/) (2021)
- [Massively Distributed SGD: ImageNet/ResNet-50 Training in a Flash](https://mlsanity.com/paper/1811.05233/) (2018)

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

- [ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale](https://mlsanity.com/paper/2303.14006/) (2023)
- [vTrain: A Simulation Framework for Evaluating Cost-effective and Compute-optimal Large Language Model Training](https://mlsanity.com/paper/2312.12391/) (2023)
- [MAD Max Beyond Single-Node: Enabling Large Machine Learning Model Acceleration on Distributed Systems](https://mlsanity.com/paper/2310.02784/) (2023)
- [Lumos: Efficient Performance Modeling and Estimation for Large-scale LLM Training](https://mlsanity.com/paper/2504.09307/) (2025)
- [Impact of RoCE Congestion Control Policies on Distributed Training of DNNs](https://mlsanity.com/paper/2207.10898/) (2022)
- [TACOS: Topology-Aware Collective Algorithm Synthesizer for Distributed Machine Learning](https://mlsanity.com/paper/2304.05301/) (2023)
- [COMET: A Comprehensive Cluster Design Methodology for Distributed Deep Learning Training](https://mlsanity.com/paper/2211.16648/) (2022)
- [GPU Cluster Scheduling for Network-Sensitive Deep Learning](https://mlsanity.com/paper/2401.16492/) (2024)
- [EcoFlow: Efficient Convolutional Dataflows for Low-Power Neural Network Accelerators](https://mlsanity.com/paper/2202.02310/) (2022)
- [FRED: Flexible REduction-Distribution Interconnect and Communication Implementation for Wafer-Scale Distributed Training of DNN Models](https://mlsanity.com/paper/2406.19580/) (2024)
- [Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs](https://mlsanity.com/paper/2511.10480/) (2025)
- [LOCALUT: Harnessing Capacity-Computation Tradeoffs for LUT-Based Inference in DRAM-PIM](https://mlsanity.com/paper/2604.04523/) (2026)
- [Towards a Flexible and High-Fidelity Approach to Distributed DNN Training Emulation](https://mlsanity.com/paper/2405.02969/) (2024)
- [ELK: Exploring the Efficiency of Inter-core Connected AI Chips with Deep Learning Compiler Techniques](https://mlsanity.com/paper/2507.11506/) (2025)
- [NEST: Network- and Memory-Aware Device Placement For Distributed Deep Learning](https://mlsanity.com/paper/2603.06798/) (2026)
- [Rethinking Network Topologies for Cost-Effective Mixture-of-Experts LLM Serving](https://mlsanity.com/paper/2605.00254/) (2026)

## Abstract (arXiv)

Distributed training is a solution to reduce DNN training time by splitting the task across multiple NPUs (e.g., GPU/TPU). However, distributed training adds communication overhead between the NPUs in order to synchronize the gradients and/or activation, depending on the parallelization strategy. In next-generation platforms for training at scale, NPUs will be connected through multi-dimensional networks with diverse, heterogeneous bandwidths. This work identifies a looming challenge of keeping all network dimensions busy and maximizing the network BW within the hybrid environment if we leverage scheduling techniques for collective communication on systems today. We propose Themis, a novel collective scheduling scheme that dynamically schedules collectives (divided into chunks) to balance the communication loads across all dimensions, further improving the network BW utilization. Our results show that on average, Themis can improve the network BW utilization of the single All-Reduce by 1.72X (2.70X max), and improve the end-to-end training iteration performance of real workloads such as ResNet-152, GNMT, DLRM, and Transformer-1T by 1.49X (2.25X max), 1.30X (1.78X max), 1.30X (1.77X max), and 1.25X (1.53X max), respectively.