Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs
Not on the map. Why?
Abstract (arXiv)
Optimizing the performance of large language models (LLMs) on large-scale AI training and inference systems requires a scalable and expressive mechanism to model distributed workload execution. Such modeling is essential for pre-deployment system-level optimizations (e.g., parallelization strategies) and hardware design-space explorations. While recent efforts have proposed collecting execution traces from real systems, access to large-scale infrastructure remains limited to major cloud providers. Moreover, traces capturing execution on a specific platform cannot be easily adapted to study alternate software and/or hardware configurations, especially at scale. We introduce STAGE, a framework that synthesizes high-fidelity execution graphs to accurately model distributed AI workloads (including LLMs and MoEs). STAGE supports a comprehensive set of parallelization strategies, allowing users to systematically explore a wide spectrum of model architectures and system configurations. STAGE demonstrates its scalability by synthesizing high-fidelity LLM traces spanning over 128K GPUs, while preserving tensorlevel accuracy in compute, memory, and communication. STAGE is publicy available at https://github.com/astra-sim/stage
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 41 of its 73 references that are in the map. First posted on arXiv 13 Nov 2025.
Builds on
41 of 73 references.
References · 41
- 2017Attention Is All You Needcited by 59,932
- 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
- 2020Language Models are Few-Shot Learnerscited by 32,199
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
- 2024The Llama 3 Herd of Modelscited by 14,192
- 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
- 2020Scaling Laws for Neural Language Modelscited by 6,411
- 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
- 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
- 2023Mistral 7Bcited by 3,065
- 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,014
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
29 more
- 2024Mixtral of Expertscited by 1,785
- 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
- 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,278
- 2024DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelcited by 1,093
- 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
- 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 641
- 2024Jamba: A Hybrid Transformer-Mamba Language Modelcited by 390
- 2022Reducing Activation Recomputation in Large Transformer Modelscited by 300
- 2019MLPerf Inference Benchmarkcited by 195
- 2018Beyond Data and Model Parallelism for Deep Neural Networkscited by 177
- 2024Zamba: A Compact 7B SSM Hybrid Modelcited by 118
- 2018PipeDream: Fast and Efficient Pipeline Parallel DNN Trainingcited by 116
- 2020PyTorch Distributed: Experiences on Accelerating Data Parallel Trainingcited by 87
- 2018Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractionscited by 84
- 2021DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusioncited by 56
- 2023ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scalecited by 47
- 2021Torch.fx: Practical Program Capture and Transformation for Deep Learning in Pythoncited by 38
- 2024LLMServingSim: A HW/SW Co-Simulation Infrastructure for LLM Inference Serving at Scalecited by 30
- 2023vTrain: A Simulation Framework for Evaluating Cost-effective and Compute-optimal Large Language Model Trainingcited by 22
- 2021Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Modelscited by 15
- 2023MAD Max Beyond Single-Node: Enabling Large Machine Learning Model Acceleration on Distributed Systemscited by 14
- 2023Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Tracescited by 13
- 2023Rail-only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameterscited by 13
- 2023Proteus: Simulating the Performance of Distributed DNN Trainingcited by 9
- 2025Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimizationcited by 6
- 2025Lumos: Efficient Performance Modeling and Estimation for Large-scale LLM Trainingcited by 6
- 2022Mystique: Enabling Accurate and Scalable Generation of Production AI Benchmarkscited by 5
- 2025Towards Easy and Realistic Network Infrastructure Testing for Large-scale Machine Learningcited by 3
Built on by
Citing papers per year: 2026: 1.
1 paper.