Toward a Standardized Representation for Deep Learning Collective Algorithms

The explosion of machine learning (ML) model size has led to its execution on distributed clusters at a large scale. Many works have tried to optimize the process of producing collective algorithms and running collective communications, which act as a bottleneck to distributed ML. However, different works use their own collective algorithm representation, resulting in a fragmented optimization environment. We propose a standardized workflow leveraging a common representation format based on Chakra Execution Trace, a widely used graph-based representation of distributed ML workloads. This common representation enables us to view collective communications at the same level as workload operations and decouple producer and consumer tools, enhancing interoperability and reducing engineering effort. We provide a proof of concept of this standardized workflow by simulating the collective algorithms generated by MSCCLang domain-specific language and the topology-aware collective algorithm synthesizer (TACOS) through the accelerator scaling for training simulator (ASTRA-sim) distributed ML simulator using multiple network configurations.