# Towards Universal Performance Modeling for Machine Learning Training on Multi-GPU Platforms

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 41 references that are in the map.
- First posted on arXiv 19 Apr 2024.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [Predicting the Computational Cost of Deep Learning Models](https://mlsanity.com/paper/1811.11880/) (2018)
- [Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect](https://mlsanity.com/paper/1903.04611/) (2019)
- [FBGEMM: Enabling High-Performance Low-Precision Deep Learning Inference](https://mlsanity.com/paper/2101.05615/) (2021)
- [Understanding Capacity-Driven Scale-Out Neural Recommendation Inference](https://mlsanity.com/paper/2011.02084/) (2020)
- [A Runtime-Based Computational Performance Predictor for Deep Neural Network Training](https://mlsanity.com/paper/2102.00527/) (2021)
- [RecShard: Statistical Feature-Based Memory Optimization for Industry-Scale Neural Recommendation](https://mlsanity.com/paper/2201.10095/) (2022)
- [AutoShard: Automated Embedding Table Sharding for Recommender Systems](https://mlsanity.com/paper/2208.06399/) (2022)
- [DreamShard: Generalizable Embedding Table Placement for Recommender Systems](https://mlsanity.com/paper/2210.02023/) (2022)
- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [Characterizing Deep Learning Training Workloads on Alibaba-PAI](https://mlsanity.com/paper/1910.05930/) (2019)
- [Building a Performance Model for Deep Learning Recommendation Model Training on GPUs](https://mlsanity.com/paper/2201.07821/) (2022)
- [Characterizing and Modeling Distributed Training with Transient Cloud GPU Servers](https://mlsanity.com/paper/2004.03072/) (2020)
- [perf4sight: A toolflow to model CNN training performance on Edge GPUs](https://mlsanity.com/paper/2108.05580/) (2021)

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

- [ATTNChecker: Highly-Optimized Fault Tolerant Attention for Large Language Model Training](https://mlsanity.com/paper/2410.11720/) (2024)

## Abstract (arXiv)

Characterizing and predicting the training performance of modern machine learning (ML) workloads on compute systems with compute and communication spread between CPUs, GPUs, and network devices is not only the key to optimization and planning but also a complex goal to achieve. The primary challenges include the complexity of synchronization and load balancing between CPUs and GPUs, the variance in input data distribution, and the use of different communication devices and topologies (e.g., NVLink, PCIe, network cards) that connect multiple compute devices, coupled with the desire for flexible training configurations. Built on top of our prior work for single-GPU platforms, we address these challenges and enable multi-GPU performance modeling by incorporating (1) data-distribution-aware performance models for embedding table lookup, and (2) data movement prediction of communication collectives, into our upgraded performance modeling pipeline equipped with inter-and intra-rank synchronization for ML workloads trained on multi-GPU platforms. Beyond accurately predicting the per-iteration training time of DLRM models with random configurations with a geomean error of 5.21% on two multi-GPU platforms, our prediction pipeline generalizes well to other types of ML workloads, such as Transformer-based NLP models with a geomean error of 3.00%. Moreover, even without actually running ML workloads like DLRMs on the hardware, it is capable of generating insights such as quickly selecting the fastest embedding table sharding configuration (with a success rate of 85%).