# ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale

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

## Key facts

- Cited by 47 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 80 references that are in the map.
- First posted on arXiv 24 Mar 2023.

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

- [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)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [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)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [A Survey on Distributed Machine Learning](https://mlsanity.com/paper/1912.09789/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [Scaling Distributed Machine Learning with In-Network Aggregation](https://mlsanity.com/paper/1903.06701/) (2019)
- [Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms](https://mlsanity.com/paper/2007.00156/) (2020)
- [Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models](https://mlsanity.com/paper/2110.04478/) (2021)
- [Flare: Flexible In-Network Allreduce](https://mlsanity.com/paper/2106.15565/) (2021)
- [Impact of RoCE Congestion Control Policies on Distributed Training of DNNs](https://mlsanity.com/paper/2207.10898/) (2022)
- [DeepFlow: A Cross-Stack Pathfinding Framework for Distributed AI Systems](https://mlsanity.com/paper/2211.03309/) (2022)
- [Clio: A Hardware-Software Co-Designed Disaggregated Memory System](https://mlsanity.com/paper/2108.03492/) (2021)
- [COMET: A Comprehensive Cluster Design Methodology for Distributed Deep Learning Training](https://mlsanity.com/paper/2211.16648/) (2022)

## Built on by (most cited first, 30 of 47)

- [LLMServingSim: A HW/SW Co-Simulation Infrastructure for LLM Inference Serving at Scale](https://mlsanity.com/paper/2408.05499/) (2024)
- [ServeGen: Workload Characterization and Generation of Large Language Model Serving in Production](https://mlsanity.com/paper/2505.09999/) (2025)
- [vTrain: A Simulation Framework for Evaluating Cost-effective and Compute-optimal Large Language Model Training](https://mlsanity.com/paper/2312.12391/) (2023)
- [Scaling Down to Scale Up: A Cost-Benefit Analysis of Replacing OpenAI's LLM with Open Source SLMs in Production](https://mlsanity.com/paper/2312.14972/) (2023)
- [MAD Max Beyond Single-Node: Enabling Large Machine Learning Model Acceleration on Distributed Systems](https://mlsanity.com/paper/2310.02784/) (2023)
- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [Demystifying AI Platform Design for Distributed Inference of Next-Generation LLM models](https://mlsanity.com/paper/2406.01698/) (2024)
- [Performance Modeling and Workload Analysis of Distributed Large Language Model Training and Inference](https://mlsanity.com/paper/2407.14645/) (2024)
- [LLMServingSim 2.0: A Unified Simulator for Heterogeneous and Disaggregated LLM Serving Infrastructure](https://mlsanity.com/paper/2602.23036/) (2026)
- [Lumos: Efficient Performance Modeling and Estimation for Large-scale LLM Training](https://mlsanity.com/paper/2504.09307/) (2025)
- [100K or 100 Days: Trade-offs when Pre-Training with Academic Resources](https://mlsanity.com/paper/2410.23261/) (2024)
- [Patterns behind Chaos: Forecasting Data Movement for Efficient Large-Scale MoE LLM Inference](https://mlsanity.com/paper/2510.05497/) (2025)
- [Characterizing the Efficiency of Distributed Training: A Power, Performance, and Thermal Perspective](https://mlsanity.com/paper/2509.10371/) (2025)
- [TACOS: Topology-Aware Collective Algorithm Synthesizer for Distributed Machine Learning](https://mlsanity.com/paper/2304.05301/) (2023)
- [DuetServe: Harmonizing Prefill and Decode for LLM Serving via Adaptive GPU Multiplexing](https://mlsanity.com/paper/2511.04791/) (2025)
- [Frontier: Towards Comprehensive and Accurate LLM Inference Simulation](https://mlsanity.com/paper/2605.21312/) (2026)
- [Towards Easy and Realistic Network Infrastructure Testing for Large-scale Machine Learning](https://mlsanity.com/paper/2504.20854/) (2025)
- [GPU Cluster Scheduling for Network-Sensitive Deep Learning](https://mlsanity.com/paper/2401.16492/) (2024)
- [Echo: Simulating Distributed Training At Scale](https://mlsanity.com/paper/2412.12487/) (2024)
- [CUCo: An Agentic Framework for Compute and Communication Co-design](https://mlsanity.com/paper/2603.02376/) (2026)
- [KernelSight-LM: A Kernel-Level LLM Inference Simulator](https://mlsanity.com/paper/2606.28565/) (2026)
- [m4: A Learned Flow-level Network Simulator](https://mlsanity.com/paper/2503.01770/) (2025)
- [Maya: Optimizing Deep Learning Training Workloads using GPU Runtime Emulation](https://mlsanity.com/paper/2503.20191/) (2025)
- [Efficient AllReduce with Stragglers](https://mlsanity.com/paper/2505.23523/) (2025)
- [Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs](https://mlsanity.com/paper/2511.10480/) (2025)
- [TinyTorch: Building Machine Learning Systems from First Principles](https://mlsanity.com/paper/2601.19107/) (2026)
- [AgentServeSim: Serving-System Simulation and Policy Search for LLM Agent Programs](https://mlsanity.com/paper/2606.09613/) (2026)
- [On the Burstiness of Distributed Machine Learning Traffic](https://mlsanity.com/paper/2401.00329/) (2023)
- [Towards a Flexible and High-Fidelity Approach to Distributed DNN Training Emulation](https://mlsanity.com/paper/2405.02969/) (2024)
- [Phantora: Maximizing Code Reuse in Simulation-based Machine Learning System Performance Estimation](https://mlsanity.com/paper/2505.01616/) (2025)

## Abstract (arXiv)

As deep learning models and input data are scaling at an unprecedented rate, it is inevitable to move towards distributed training platforms to fit the model and increase training throughput. State-of-the-art approaches and techniques, such as wafer-scale nodes, multi-dimensional network topologies, disaggregated memory systems, and parallelization strategies, have been actively adopted by emerging distributed training systems. This results in a complex SW/HW co-design stack of distributed training, necessitating a modeling/simulation infrastructure for design-space exploration. In this paper, we extend the open-source ASTRA-sim infrastructure and endow it with the capabilities to model state-of-the-art and emerging distributed training models and platforms. More specifically, (i) we enable ASTRA-sim to support arbitrary model parallelization strategies via a graph-based training-loop implementation, (ii) we implement a parameterizable multi-dimensional heterogeneous topology generation infrastructure with analytical performance estimates enabling simulating target systems at scale, and (iii) we enhance the memory system modeling to support accurate modeling of in-network collective communication and disaggregated memory systems. With such capabilities, we run comprehensive case studies targeting emerging distributed models and platforms. This infrastructure lets system designers swiftly traverse the complex co-design stack and give meaningful insights when designing and deploying distributed training platforms at scale.