  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.
