Towards a Standardized Representation for Deep Learning Collective Algorithms

Authors: jinsun  yoo
william  won
meghan  cowan
nan  jiang
benjamin  klenk
srinivas  sridharan
tushar  krishna
Venue: iscaconf_org_ISCA_2024_Workshop_MLArchSys
Published: 2024-05-06 00:09:10

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 algorithm and running collective communication, which acts as a bottleneck to distributed machine learning. However, the lack of a standardized collective algorithm representation has hindered interoperability between the workload representation, collective algorithm producers, and consumers. The trend of collective algorithm producers and consumers using their own representation has pushed away from co-optimizing collective communications and the rest of the workload. Additionally, tool-specific conversions and modifications have to be made for each pair of tools producing and consuming collective algorithms. 

In this 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, which is being used to represent distributed machine learning workloads. Such a common representation enables 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 MSCCLang domain-specific language through ASTRA-sim distributed machine learning simulator using various network configurations.

Main paper: 
 i. introduction: 
 Recent trends in enormous machine learning (ML) models, such as recommendation models [1] or Large Language Models (LLMs) [2] , [3] , have made it impractical to execute them on a single Neural Processing Unit (NPU, such as GPU, TPU, or custom ASIC) [4] . Consequently, ML execution has evolved to distribute the job across multiple NPUs [5] .Within distributed ML, each of the participating NPUs completes a portion of the overall compute task. They periodically transfer and synchronize their intermediate compute results (e.g., weight or input gradients) in accordance with a predefined schedule [6] . Such traffic patterns are collective in nature. Therefore, collective communication primitives, such as All-Reduce or All-Gather, have been the key building blocks of distributed ML platforms [7] .The communication of intermediate data has become a bottleneck to the overall distributed ML execution, and recent studies have tried to optimize this [8] - [10] . Within the scope of collectives, Collective Communication Libraries (CCLs) such as NVIDIA's NCCL [11] provide implementations of several predefined collective algorithms (e.g., Ring [12] and Double Trace (Chakra ET) as a common representation for both distributed ML workload and collective algorithm. Downstream tools receive both workload and collective algorithms represented using a common Chakra ET format. Sample Chakra ETs of the workload (left) and that of the algorithm of a single collective (right) is also shown.Binary Tree [13] ). Several works have tried to further optimize collective algorithms [14] , [15] , by automatically synthesizing topology-aware algorithms [16] , [17] , or allowing users to easily define and test their own [18] .The outcome of such optimizations is not limited to actual execution on real clusters. For example, distributed ML simulators [19] - [21] allow users to swiftly compare the performance of different collective algorithms or workload parallelization strategies without running them on actual hardware. This enables them to easily test new optimizations for current systems and additionally allows the co-design and evaluation of futuristic distributed ML platforms.Unfortunately, these collective optimizations have been done separately from distributed ML workload optimization, due to a lack of common representation. Specifically, distributed ML workload information does not contain details on how each collective operation is implemented, and the downstream tool must fill in the blank information on its own. As such, downstream tools have their own implementation of collective algorithms and choose from these options as necessary. To the best of our knowledge, there is no common format that represents both the distributed ML workloads and collective algorithms. Because of this, collective communications and other operators within a workload (such as compute or memory access) are also optimized separately. Additionally, the absence of a common representation between upstream producers and downstream consumers limits interoperability. In order to evaluate custom collective algorithms with a desired simulator, users have to manually program the algorithm following simulator-specific rules. This adds a layer of development burden to the search for collective algorithm optimizations.In order to bridge this gap, we propose a common representation format for both workload and collective algorithms across different tools. We draw attention to the Chakra Execution Trace (Chakra ET) [22] format. Chakra ET is a standard graph-based representation of execution traces for distributed ML workloads, and is currently being developed by MLCommons [23] . In this work, we propose to extend Chakra ET to also encode collective algorithms. Fig. 1 summarizes the standardization proposal, where (i) both the distributed ML workload and collective algorithm (produced by the upstream tools) are represented in a common Chakra ET format and (ii) downstream tools can directly ingest a single Chakra ET representation for simulation or execution. We develop a proof-of-concept workflow allowing users to write custom collective algorithms using the MSCCLang Domain Specific Language [18] , represent them in Chakra ET format, and compare the custom algorithms using the ASTRA-sim distributed ML system simulator [19] . Similar workflows with different tools can be easily built leveraging this standard representation format.There are three key benefits to employing Chakra ET as a common collective algorithm representation. First, we envision to co-optimize collective communication with other operations in the workload via leveraging a streamlined workflow encapsulating both the workload and collective algorithms. Second, the decoupling between upstream and downstream tools helps reduce the effort to develop and implement collective algorithms for each downstream tool. For example, distributed ML simulators like ASTRA-sim [19] could ingest the common representation instead of users having to implement the collective algorithm within its codebase, which requires simulator-specific knowledge. Finally, interoperability across different downstream tools is another benefit, where users can leverage a single representation to test the same collective algorithm on both simulators like ASTRAsim and real systems like MSCCL-Runtime [18] .To summarize, our contributions are as follows:• We motivate and propose a standard workflow that uses a common collective algorithm representation to bridge distributed ML workload information, upstream collective algorithm producers, and downstream tools. • As proof-of-concept, we provide a case study showing interoperability with MSCCLang and the ASTRA-sim simulator.Chakra ET [22] aims to provide a standard graph-based representation to capture the trace of distributed ML workload execution. Chakra ET represents a distributed workload using a directed acyclic graph whose vertex denotes ML operations and edge indicates inter-operation dependency. Each node can be of type COMP, MEM, or COMM COLL (either collective communication or single COMM SEND and COMM RECV). Chakra ET can be fed into and consumed by distinct downstream tools, such as distributed ML simulators or benchmarking tools. These downstream tools leverage Chakra ET by traversing through the graph and issuing the operations whose dependencies are resolved and are ready to be dispatched. Chakra ET graphs can be collected in multiple ways, for instance, through profiling actual PyTorch executions or via synthetic generations.Note that a collective communication node in Chakra ET merely indicates that a collective communication has taken place. However, it does not encode the actual collective algorithm. In other words, Chakra ET itself is oblivious to exactly how the messages are orchestrated and transferred. Consequently, the downstream tools exploit their internal, toolspecific implementations of collective algorithms.Given the significance of collective communication in distributed ML, numerous studies have focused on optimizing collective algorithms. Primarily, they have pursued two main threads: (i) developing domain-specific languages (DSLs) to enable users to define their own collective algorithms and (ii) implementing synthesizers to autonomously generate them.One notable example is MSCCLang [18] , which introduces a Python-based DSL for collective algorithms. This enables users to easily construct NCCL-based collective algorithms. MSCCLang compiles these algorithms into an XML-based representation (MSCCL-IR), which is then executed on real clusters via the NCCL-based MSCCL-Runtime. Meanwhile, synthesizers like TACCL [17] and TACOS [24] generate collective algorithms tailored to the network topology. TACCL employs Integer Linear Programming (ILP) to identify optimal collective algorithms, while TACOS utilizes a Time-expanded Network (TEN) approach.Downstream tools receive distributed ML workloads or collective algorithm representations to execute meaningful tasks. Collective communication runtimes and distributed ML simulators are notable instances of such downstream tools. For instance, the MSCCL-Runtime orchestrates a collective communication by taking the collective algorithm in MSCCL-IR format (compiled from MSCCLang) and executing it on real GPU clusters via an NCCL-based runtime. Conversely, ASTRA-sim is a notable example of simulation infrastructure. Its Workload Layer can receive and simulate a distributed ML workload represented in Chakra ET format, while its System Layer implements collective algorithms to fill in the gaps of the workload Chakra ET. The Network Layer captures requests from these layers and simulating the actual network transfers over a network simulator of user choice.Currently, upstream collective algorithm producers employ unique representations to describe their results. For example, MSCCLang utilizes a NCCL-based, low-level XML interpretation, while TACOS relies on its own TEN representation.This lack of standardization often leads downstream tools to rely on their unique internal implementation. Consequently, the format and pipeline that downstream tools use to fetch collective information diverge from those used to inject workload information. As a result, users are constrained to optimizing either collective operations or other workload operations, not both.Moreover, the absence of a standard format means that executing an upstream producer's algorithm with a specific downstream tool requires users to comprehend the internal details of both tools and implement the algorithm themselves. This task is not only highly prohibitive but also implies that it must be repeated for every pair of upstream and downstream tools. Consequently, upstream and downstream processes become disjointed and lack plug-and-play functionality.We standardize the collective algorithm representation by utilizing the Chakra ET format that is already employed for distributed ML workloads. It readily offers mechanisms to capture point-to-point message transfers between arbitrary NPUs as well as compute operations necessary for collectives.By representing both distributed ML workloads and collective algorithms in Chakra ET, we cleanly resolve the issue of separation among workloads, upstream, and downstream tools. Fig. 1 depicts the proposed standardized workflow. Both distributed ML workloads and collective algorithms are in Chakra ET format and passed on to the downstream. The downstream tool then traverses the workload ET and executes operations as their dependencies are resolved. During the process, since collective communication nodes lack the exact mechanism to execute collectives, the downstream tools must decide the algorithm. Previously, they selected a collective algorithm from a range of native implementations or custom algorithms tailored to their specifications. However, with the capability to receive any collective algorithm in Chakra ET, the downstream tools can easily expand the collective communication node with the provided algorithm. Users can rapidlyAll-Gather By representing collective algorithms with the same format as the workload, we elevate the send and receive messages of a collective algorithm to the same level as other operators in the workload. This opens up opportunities for co-optimizing collective communication and other operators such as compute. For example, it is much simpler to test a scheduling feature that reorders a compute node and a send node (assuming the reorder respects inter-node dependencies) as the compute and send node now use a common Chakra ET format.Note that real-world systems such as MSCCL-Runtime do not take workload information as input. Even in this case, having the Chakra ET as a common collective algorithm representation helps bridge the gap with upstream producers. For example, a user may want to simulate multiple prospective algorithms using ASTRA-sim, then validate the bestperforming candidates by actually using the MSCCL-Runtime. The user in this workflow can reuse the same Chakra ET format across both tools without modification. This is achieved by the fact that the common representation abstracts away the details of the downstream tools.Table I lists the types of Chakra ET operator nodes that are used to represent collective algorithms and their definition. Leveraging these nodes, Chakra ET can represent point-topoint network transfer between two NPUs as well as compute operations. Arbitrary collective algorithms can then be described as a combination of network message transfers and reductions.To meet the standardization requirement, upstream tools need to convert resulting collective algorithms from their default representation to the common Chakra ET format when producing collective algorithms. We highlight that implementing this conversion is a one-time task such that, once developed, can be reused across multiple downstream frameworks. NetworkAs a case study and proof-of-concept, we construct an endto-end workflow using MSCCLang and ASTRA-sim. Here, we describe the extensions made to the two tools for the case study. We expect similar extensions would apply to other tools.We developed a converter that bridges the MSCCL-IR format into Chakra ET. The converter creates a vertex (i.e., Chakra ET node) for each operation in the MSCCL-IR format. The converter then creates edges by extracting inter-operator dependency information encoded in the MSCCL-IR format. To showcase the update, we described a 1D Ring algorithm of All-Reduce using MSCCLang and compiled the result in the standard Chakra ET format. At each run, the user will choose which algorithm to use for each collective. Whenever a collective communication node is issued, ASTRA-sim will run the corresponding algorithm code. We extend ASTRA-sim by adding an input parameter for the user to refer to the collective algorithm in Chakra ET format. The simulator will parse the provided Chakra ET and simulate the point-to-point sends and receives following the dependencies, rather than using its own implementation. This process is depicted in dashed boxes in Fig. 3 . Since ASTRA-sim readily supports the execution of ML workloads in Chakra ET, reusing the components in the workload layer has made it easy to add the extension for the common collective representations as well.We showcase our case study of bridging MSCCLang and ASTRA-sim as described in Sec. IV. We use an All-Reduce collective followed by an All-Gather collective for the workload and use MSCCLang to generate a 1D Ring algorithm for both collectives. Fig. 2 ET represented as a graph. Note how we are able to represent both the workload and the collective algorithm using the same Chakra ET format. The resulting Chakra ET representation is then provided to the ASTRA-sim simulator to be tested across different topologies with varying physical connectivity. We use the Analytical network simulator to model the message transfer. The topologies consist of 64 NPUs with varying connectivity. We observe the workload duration as we differ the size of the collective. Fig. 4 shows the simulation results, using the slowdown of the different topologies compared to a 1D Ring. Note that such an experiment was made possible thanks to the streamlined workflow via standardization of the collective algorithm representation. It is natural that other topologies show slowdowns as we use a simple Ring algorithm for both All-Reduce and All-Gather. While this evaluation showcases a simplistic workload, it is possible to expand the evaluation to include complex workloads and collective algorithms. We leave this to future work.This paper proposes a standardized common representation for collective algorithms. We reuse the Chakra ET format, which already captures distributed ML workload traces, as the collective representation. Representing both workloads and collective algorithms with the same format will allow us to explore the co-optimization of workload operators and collective communication operators. We showcase such a common representation with a case study using collective algorithms produced by MSCCLang on the ASTRA-sim simulator.Our work opens up several future research directions. One important future work is to leverage the proposed standard representation to explore collective optimizations in the context of an actual workload. We anticipate that a common representation across a workload and a collective algorithm will allow us to study the co-optimization of compute and communication operations. This allows researchers to further study the overlapping of compute and communication operations. Another potential direction would be to expand other tools to produce and consume Chakra ET format, to further expand the ecosystem and scope of our proposed workflow. 
 
 Appendix:  


Decision: Congratulations! We are pleased to inform you that your paper has been accepted for presentation at MLArchSys 2024. We look forward to your participation at the workshop. Further details regarding the schedule and format will be provided soon. See you at the workshop!

Reviews: review: title: title: Review: Towards a Standardized Representation for Deep Learning Collective Algorithms
strength_weakness: strength_weakness: None.
questions: questions: As you paper mentions, there are 3 primary reasons for a standardized representation:

(1) Reduce the effort to develop collective algorithms for downstream tools.

(2) Interoperability across different downstream tools.

(3) Co-optimize collective communication with other operations in the workload.

(1) and (2) are valuable and there is no debating that they are necessary to enable more streamlined and effective systems research on this topic. 

However, (3) was what I found to be the most intriguing and compelling as a researcher. This potential benefit was mentioned multiple times in the paper and I kept waiting as a reader for the actual demonstration of this co-optimization piece. Eventually I finished the evaluation and found that it will be part of the future work. I’m excited to see how these results turn out as it was something I was wanting to see in the case study. I think adding this piece will strengthen the contributions needed to take the next step in this work towards a research conference submission. Overall, I thought the paper was well-written, making it easy to read and follow.
title: title: Towards a Standardized Representation for Deep Learning Collective Algorithms
strength_weakness: strength_weakness: NA
questions: questions: ## Summary
* This paper tries to address the problem that the lack of a unified representation for compute operations
and collectives in distributed ML serving use cases leads to suboptimal implementations.
* The paper proposes to extend and use the Chakra execution trace to bridge this gap.
* The paper provides a proof of concept case study where the same representation is used by collective optimization
algorithm and ML partitioning algorithm to co-optimize execution

## Comments
Thank you for submitting your work to MlArchSys. I liked the motivation for using the same intermediate
representation, however I think the paper needs some more evaluation to help readers understand the magnitude
of the contributions. Please try to address the following in any subsequent drafts:
1. Please explain why chakra ET is vital. For example, each tools can use its own representation but the optimization can
be done iteratively and collaboratively to attain an optimized execution configuration. Please highlight
what a common representation enables as compared to this collaborative approach.
2. Please showcase some salient points of the chakra ET representation that makes it useful for both the upstream and downstream tools.
title: title: Official Review
strength_weakness: strength_weakness: * The evaluation section is limited to a specific proof-of-concept and does not cover a wide range of real-world scenarios or more complex workloads, which may limit the generalizability of the findings.
* The paper may be challenging for readers who are not already well-versed in distributed machine learning and the specifics of collective algorithms, potentially limiting its accessibility and impact.
* The paper does not fully address the effort and challenges involved in extending other existing upstream and downstream tools to support the proposed Chakra ET standard, which may pose practical barriers to adoption.
questions: questions: The quality of the work is high, with a well-defined problem statement and a comprehensive proposed solution. The authors provide a clear explanation of the current challenges in distributed machine learning regarding collective algorithms and their representation. The use of Chakra ET as a standardized representation is well-motivated, and the proof-of-concept implementation demonstrates the practicality of the proposed approach.

The clarity of the paper is generally good, with well-organized sections that guide the reader through the problem, solution, methodology, and evaluation. However, some technical sections could benefit from additional explanations and visual aids to make the concepts more accessible to a broader audience. For instance, the detailed process of converting collective algorithms to Chakra ET format could be accompanied by more illustrative examples.

The originality of the work is significant, as it addresses a critical gap in the field of distributed machine learning. The idea of using a standardized execution trace format to unify the representation of collective algorithms and workloads is novel and has the potential to impact both academia and industry positively.

The significance of the work is high. By providing a standardized representation, the authors enable more seamless integration and interoperability between different tools and platforms used in distributed machine learning. This can lead to more efficient development, testing, and optimization of collective algorithms, ultimately advancing the state-of-the-art in distributed ML systems.