Paper: Elastic Load Balancing for Dynamic LLMs (ICLR-2024)

Reviewer comment (Writing):
The paper is poorly written and hard to understand, and it makes some exaggerated statements about its contribution. For example, "Research on dynamic models will not deliver practical impact unless there is a platform from which those models can be made efficient." -- This is untrue from the machine learning algorithm's perspective. As long as such an algorithm shows statistical efficiency or better generalization performance, it has a significant practical impact with or without a platform/framework.

Author response:
The paper makes some exaggerated statements about its contribution. For example, "Research on dynamic models will not deliver practical impact unless there is a platform from which those models can be made efficient." -- This is untrue from the machine learning algorithm's perspective. As long as such an algorithm shows statistical efficiency or better generalization performance, it has a significant practical impact w/wo a platform/framework, right?

We argue that it if there is no effective high-throughput solutions (like the one in this paper) that allow for researching the throughput effectiveness of methods like pruning/freezing techniques at scale (and we emphasize "at scale"), researchers would and are limited in studying those methods on a small scale, which is a problem. For instance, without an effective methods for studying dynamic GPT3-sized models on up to 1000 GPUs (as in this paper), the true potential of dynamic models on large models can not be realized. That being said, we can tune down this claim in the paper in the revised PDF to not give wrong impressions to the reader.

Reviewer comment (Writing):
The technique section is confusing and not self-explained; for example, in Lemma 1 and 2, the term "bubble" is referred to without a formal definition.

Author response:
The technique session is confusing and not self-explained; for example, in Lemma 1 and 2, the term "bubble" is referred to without a formal definition.

Bubbles are formally defined in Supp. material Section A.

Reviewer comment (Experiments):
The experimental section is poorly organized, with a lack of formalization of the central hypothesis of the evaluation.

Author response:
The experimental section is poorly organized; there is even a lack of formalization of the central hypothesis of the evaluation.

We will revise the structure of that section in the revised PDF, and also add a summary of the key findings.

Reviewer comment (Experiments):
There is an issue with the baseline selection: Megatron is designed to support standard non-sparse LLM training, and DeepSpeed is similar where the additional effort is made for MOE. Those are not strong baselines for such training tasks. Systems like PEFT from HuggingFace include some relevant functionality, so more advanced baselines should be considered for evaluation.

Author response:
There is also some issue w.r.t the baseline selection, Megatron is designed to support standard non-sparse LLM training, DeepSpeed is similar where the additional effort is made for MOE. Those are not strong baselines for such training tasks. As far as I know, systems like PEFT from huggingface include some relevant functionality. Some more advanced baselines should be considered for evaluation.

PEFT, and other baselines, are not applying true sparse computation (i.e. they just mask the pruned values to zero and use dense MatMult), and hence the real imbalance that comes from using sparse models would not manifest. Hence comparing to PEFT would be akin to comparing to a dense computational baseline (similar to the case of Megatron and DeepSpeed). Section C in the supp. material describes how we had to implement bindings for really sparse operations.

Reviewer comment (Evaluation):
In figure 4, why is the throughput of the model with different number of GPUs roughly the same. Shouldn’t we expecting the throughput will increase accordingly?

Author response:
>In figure 4, why is the throughput of the model with different number of GPUs roughly the same. Shouldn’t we expecting the throughput will increase accordingly?

The results in that figure are for weak scaling (as shown on the label above the figure). The amount of work increases with the number of GPUs, hence no performance improvement is to be expected.

Reviewer comment (Experiments):
In section 4.1, how is the prune region and iteration being selected. What is the affect of changing the settings?

Author response:
>In section 4.1, how is the prune region and iteration being selected. What is the affect of changing the settings?

They follow the pruning/freezing strategy; load re-balancing is done every time the model changes.

Reviewer comment (Reproducibility):
In the description of Diffusion by Param and Diffusion by Time, how do they iteratively minimize load variances among accelerators? Do they communicate or they select parameters based on the variance of after gradient decent?

Author response:
>In the description of Diffusion by Param and Diffusion by Time, how do they iteratively minimize load variances among accelerators? Do they communicate or they select parameters based on the variance of after gradient decent?

Diffusion happens in a decentralized way where each GPU passes some of its load to the neighbor if it has more work that its neighbor. This is repeated iteratively until each GPU has the same amount of work. This decentralized approach is particularly beneficial when scaling to a large number of GPUs, since no collective communication between GPUs would be required.

Reviewer comment (Presentation):
In figure 4, why only partition by param is included?

Author response:
>In figure 4, why only partition by param is included?

In that figure we report "partition by param" only since it is the highest performing baseline.

Reviewer comment (Evaluation):
In figure 5(right), it seems that there is not a clear relationship between number of GPU and throughput/# GPUs. Why is that? The goal is the increase the throughput after pruning in order to use less computing resources. Moreover, if we multiply # of GPU with throughput/# GPUs, the total throughput of the system will decrease. Why is that?

Author response:
>In figure 5(right), it seems that there is not a clear relationship between number of GPU and throughput/# GPUs. Why is that? The goal is the increase the throughput after pruning in order to use less computing resources. Moreover, if we multiply # of GPU with throughput/# GPUs, the total throughput of the system will decrease. Why is that?

There is a relation: Throughput/ #GPUs tends to go up as the number of GPUs decrease; that is since the more the GPUs, they more pronounced the effect of load imbalance will be, and hence the throughput/ #GPUs will drop. That is akin to how throuhput/GPU is higher in single GPU run than multi-GPU runs since you avoid the communication cost.On the other hand, the aggregate throughput (right y-axis) will drop if you use fewer GPUs.

Reviewer comment (Novelty):
**Limited Novelty**: The paper's technical contribution is limited due to its use of existing techniques in a new context, rather than introducing entirely new concepts. The proposed load balancing solutions seem to involve applications of DeepSpeed's workload partitioning algorithms at the end of each pruning, and the diffusion-based algorithm also employs a similar partitioning strategy, albeit more akin to work stealing.

Author response:
To the authors knowledge, this paper is the first to demonstrate provably practical solutions, at scale, for load balancing dynamic models.  We use DeepSpeed API to just move the parameters between GPUs; DeepSpeed itself had no workload partitioning algorithms that we build upon.

Reviewer comment (Evaluation):
**Lack of Overhead Discussion**: The paper does not adequately discuss the overhead of different strategies, which could have provided an interesting perspective on these solutions. It remains unclear whether the throughput reported in the evaluation takes into account all overheads.

Author response:
The supplementary material (Section D.1 "Overhead of Load Balancing") has a detailed analysis of the overhead (Table 1 in Supp. material)

Reviewer comment (Evaluation):
**Absence of Key Training Metrics in Evaluation**: The evaluation does not include important training metrics such as time-to-accuracy or learning curves. Even though model accuracy is not the primary focus of the proposed solution, including this information would have added to the completeness of the paper. For instance, it might be the case that improvements in throughput during large-scale training do not translate into time-to-accuracy gains because the pruning/freezing techniques reduce training quality. This potential issue could limit the usefulness of such solutions.

Author response:
We argue that it is the other way around: if there is no effective high-throughput solutions (like the one in this paper) that allow for researching the throughput effectiveness of methods like pruning/freezing techniques at scale (and we emphasize "at scale"), researchers would and are limited in studying those methods on a small scale, which is a problem. For instance, without an effective methods for studying dynamic GPT3-sized models on up to 1000 GPUs (as in this paper), the true potential of dynamic models on large models can not be realized.

We will include those numbers to the revised PDF. Note that we will not be able to include results for the largest sized model since training those models to completion can take weeks (and the deadline is in 11 days).