MetaDist: An Infrastructure for Automatic Parallelism via ShardCombine Algorithm

Please see the Strengths and Weaknesses.

(Q1) From my humble opinion, this work is not well motivated and the rationality behind the proposed methods is not well explained.

• First, in the introduction section, the authors discussed two main challenges, i.e., (i) lack of ecological compatibility, and (ii) difficulty in development, maintenance, and benchmarking. However, they seems more of an implementation issue rather than a research topic.
• Second, there is no discussion about the limitations of existing methods and how the MetaDist representation and the ShardCombine algorithm tackles them (see questions below). For me, they are just put forward straightforwardly.

(Q2) In essence, MetaDist can be viewed as a kind of IR in a distributed manner. Therefore, I believe it is necessary to compare MetaDist related works such as pONNX [1] and Unity [2]. For me, the ShardSpec and CombineSpec can be broken down into the partition, combine, replicate, and reduce operators in Unity, and the latter can provide more fine-grained representations.

[1] Wang et al. Parallel Training via Computation Graph Transformation.
[2] Unger et al. Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization.

(Q3) MetaDist focuses on intra-op parallelism, so it is unknown whether could it support pipeline parallelism, which is frequently used in large-scale distributed deep learning.

(Q4) Section 3.3 states that “The ShardCombine algorithm is an exploration algorithm that utilizes heuristic information to shard input data, and attempts to re-combine local results into global results using the TryCombine function”. What is the heuristic information here? Furthermore, how the heuristic information helps your algorithm design?

(Q5) Section 3.4.1 states that “The objective of the problem is to minimize the communication cost”. It is confusing why computation cost is not taken into account here. Furthermore, it is unclear how the ILP differs from the one in Alpa.

(Q6) The experiments were conducted on a single GPU server with eight GPUs. I am afraid the exploration space of parallelism would be extremely limited and could not evaluate the performance thoroughly.

(Q7) Regarding the experimental results, there are a few issues:

• The major competitors are Alpa and FairScale, however, there should be frameworks that are more widely used for large model training in practice, such as DeepSpeed and Megatron-LM. I believe comparisons are necessary.
• The authors stated that “ZeRO-3 is difficult to scale efficiently because of its communication overhead”. When memory is not scarce, ZeRO-2 may be a better option than ZeRO-3 and pure DataParallel. Please also compare with ZeRO-2.
• When training on a single GPU, there should be no performance differences, but they are reported in Figure 8. Please elaborate.
• There are several cases that MetaDist-BS runs out of memory. Does it indicate that your searching algorithm does not take memory limit as the constraint?

Please check the Weaknesses for detailed questions to be answered.

• Are there some examples when the cache mechanism lacks effectiveness and is effective?
• Are there artifacts available for reproducibility purposes?
• Why does Beam Search + ResNet + JAX in Figure 7 so much longer than any other combination? Is it simply because JAX has more fine-grained operators?
• Are there results on at least a 1 billion scale model? do the four different model sizes refer to GPT2, 3.5, etc?

1. In section 3.2, what is callable primitive operator stans for in op_function, invars, outvars, and spmd_rules? 2)In figure 2, what is the MetaIR? There is a MetaIR on the left side and another MetaIR in the middle.
2. In Section 3.3 Sharedim(id=j)(Sj) is used to determine the parallelism strategies. How many such strategy could be applied? I notice the paper mentioned S1 S2 S3 in the rest of the paper. For Si, how many ways to partition the tensor?
3. What is the definition of score function? Is it just the train time? If two parallelisms could be applied, how to evaluate which one is better?
4. How is the communication bandwidth affect the result?
5. The treatment and discussion of prior work on automatic parallelism and compiler approaches for automatic parallelization only covers papers from 2022 and 2023 when there are pioneering approaches even before that. Here are some examples of compiler approaches to automatic parallelization: "A load balancing inspired optimization framework for exascale multicore systems: A complex networks approach." In 2017 IEEE/ACM International Conference on Computer-Aided Design (ICCAD), pp. 217-224. IEEE, 2017. "Self-optimizing and self-programming computing systems: A combined compiler, complex networks, and machine learning approach." IEEE transactions on very large scale integration (VLSI) systems 27, no. 6 (2019): 1416-1427. "A distributed graph-theoretic framework for automatic parallelization in multi-core systems." Proceedings of Machine Learning and Systems 3 (2021): 550-568. "Plasticity-on-chip design: Exploiting self-similarity for data communications." IEEE Transactions on Computers 70, no. 6 (2021): 950-962. The authors have to check the literature and carefully contrast all existing related approaches to automatic parallelization and compiler approaches especially when these prior approaches were developed to be computer architecture independent and adaptable to general settings.