Search papers, labs, and topics across Lattice.
This paper enhances the ASTRA-sim simulator by addressing its limitations and introducing new features that enable high-fidelity simulations of distributed machine learning infrastructure. By implementing a cache-line-sized load-store granularity and a detailed GPU execution model, the authors achieve a balance between simulation scalability and fidelity. The updated ASTRA-sim 3.0 allows for comprehensive design space explorations, leading to optimized collective algorithms and improved network and GPU architecture requirements.
ASTRA-sim 3.0 reveals that high-fidelity simulations can significantly optimize distributed ML algorithms and infrastructure design.
Distributed machine learning (ML) is a key paradigm for today's large-scale artificial intelligence applications. As model inference arises as an important use case, faithful modeling of latency-sensitive collective communication has never been more important. Capturing the device architecture and modeling control and data paths at high fidelity is therefore a necessity today. Having a common, detailed representation for distributed ML infrastructure is also crucial. We revisit the promising open-source, community-driven simulator: ASTRA-sim. In this work, we identify limitations of the current ASTRA-sim simulator and augment it with new features. To this end, we enable fine-grained, high-fidelity simulation with a standardized infrastructure representation, opening new design space exploration opportunities. We propose the simulation at cache-line-sized load-store granularity, with a detailed graphics processing unit (GPU) execution model, to balance simulation scalability and fidelity. We also introduce InfraGraph, a standardized representation to capture distributed ML network infrastructure in detail. Using the updated ASTRA-sim 3.0 simulator, we showcase interesting design space explorations for designing optimized collective algorithms, network requirements, and GPU architectures.
