ASTRA-sim: Enabling Software-Hardware Co-Design Exploration for Distributed Machine Learning Platforms
Update in Progress
As we prepare for the tutorial, this page will be updated periodically. Please stay tuned!
Overview
In this tutorial, we will educate the research community about the challenges in the emerging domain of distributed machine learning, demonstrate the capabilities of ASTRA-sim with examples and discuss ongoing development efforts.
NEW – In this tutorial, we will (1) introduce the latest features added to the ASTRA-sim framework, and (2) have invited presentations on the usage of ASTRA-sim in various use cases.
Date/Location
- Jun 27, 2026 (full-day)
- Raleigh Convention Center Info
Description
As Artificial Intelligence (AI) models are scaling at an unprecedented rate, Machine Learning (ML) execution heavily relies on Distributed ML over customized neural accelerator (e.g., GPU or TPU)-based High-Performance Computing (HPC) platforms connected via high-speed interconnects (e.g., NVLinks). Examples today include NVIDIA’s HGX, Google’s Cloud TPU, and Meta’s Research Supercluster. Distributed Deep Neural Network (DNN) execution involves a complex interplay between the DNN model architecture, parallelization strategy, scheduling strategy, collective communication algorithm, network topology, remote memory accesses, and the accelerator endpoint. As innovation in AI/ML models continues to grow at an accelerated rate, there is a need for a comprehensive methodology to understand and navigate this complex intertwined co-design space to (i) architect future platforms, (ii) develop novel parallelism schemes to support efficient training of future DNN models, and (iii) develop novel fabrics for AI systems. As an ongoing collaboration between Georgia Tech and several companies, we have been jointly developing (1) a comprehensive methodology to capture arbitrary distributed ML workloads, named Chakra Execution Trace and (ii) a detailed cycle-accurate distributed AI simulator called ASTRA-sim.
ASTRA-sim models the co-design space of distributed ML described above and schedules the compute-communication interactions over plug-and-play computation, network, and remote memory simulators. It enables a systematic study of bottleneck detection and futuristic system evaluation at the software and hardware levels for scaling distributed ML. ASTRA-sim leverages the Chakra format to describe arbitrary distributed ML workloads. It uses a Google TPU-like simulator as its computation model and provides a suite of network models (analytical network, Garnet, and ns-3) for the choice of simulation speed and fidelity.
Target Audience
Any researcher with the interest in full-stack, large-scale AI/ML simulation.
Organizers
To be updated!
Schedule
To be updated!
Additional Resources
MLCommons Chakra Working Group
ASTRA-sim2.0 Paper
William Won, Taekyung Heo, Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Tushar Krishna, “ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale,” In Proc. of the IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS ‘23), 2023.
Chakra Paper
Srinivas Sridharan, Taekyung Heo, Louis Feng, Zhaodong Wang, Matt Bergeron, Wenyin Fu, Shengbao Zheng, Brian Coutinho, Saeed Rashidi, Changhai Man, and Tushar Krishna, “Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces,” In arXiv:2305.14516 [cs.LG], 2023.
CollectiveAPI Paper
Jinsun Yoo, William Won, Meghan Cowan, Nan Jiang, Benjamin Klenk, Srinivas Sridharan†, and Tushar Krishna, “Towards a Standardized Representation for Deep Learning Collective Algorithms,” In Proc. of the 2024 IEEE Symposium on High-Performance Interconnects (HOTI ‘24), 2024.