Towards a Standardized Representation for Deep Learning Collective Algorithms

TL;DR

This paper proposes a standardized, graph-based representation for collective algorithms in distributed machine learning, aiming to improve interoperability, co-optimization, and reduce engineering efforts across tools.

Contribution

It introduces a common collective algorithm representation based on Chakra Execution Trace, enabling better integration and simulation of collective algorithms across different tools.

Findings

Demonstrated the feasibility of the standardized workflow with a proof-of-concept.

Enabled simulation of collective algorithms across various network configurations.

Improved interoperability between collective algorithm producers and consumers.

Abstract

The explosion of machine learning model size has led to its execution on distributed clusters at a very 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 machine learning. However, different works use their own collective algorithm representation, pushing away from co-optimizing collective communication and the rest of the workload. The lack of a standardized collective algorithm representation has also hindered interoperability between collective algorithm producers and consumers. Additionally, tool-specific conversions and modifications have to be made for each pair of tools producing and consuming collective algorithms which adds to engineering efforts. In this position paper, we propose a standardized workflow leveraging a common collective algorithm representation. Upstream producers and downstream consumers converge to a common representation format based on Chakra Execution Trace, a commonly used graph based representation of distributed machine learning workloads. Such a common representation enables us to view collective communications at the same level as workload operations and decouple producer and consumer tools, enhance interoperability, and relieve the user from the burden of having to focus on downstream implementations. We provide a proof-of-concept of this standardized workflow by simulating collective algorithms generated by the MSCCLang domain-specific language through the ASTRA-sim distributed machine learning simulator using various network configurations.