Elastic and Balanced End-to-end Training of Dynamic LLMs with DynMo

The paper writing needs to be improved. 1/ Load imbalance is a major motivation for this paper, but there is no formal definition in the paper; 2/ This paper identifies six example cases for load imbalance, but it is unclear why these dynamic model training leads to load imbalance and worsens communication bubbles. More fundamental analysis is needed to understand the problem.

We added a formal definition of load balancing and re-packing as a general problem for dynamic models, and adapted it further to formally define the imbalance arising from each of the example cases. We added the formal definitions to the paper: Section 2.2 (line 168) for load balancing and Section 3.4 (line 611) for re-packing. Note: we added the above to the paper proper to help the reviewers locate the added content. should the paper be accepted, the added formal definitions will be only summarized in the paper proper, and the details would be moved to the supplementary material.

The two algorithms introduced in Section 3.3 are not convincing to achieve the claimed goal for optimal load balancing.

The paper includes a full theoretical proof that the algorithms are guaranteed to converge to the maximum reduction in imbalance in bounded rounds (proofs of Lemma 1 and Lemma 2 in Appendix 3.2). We further explain in detail how the algorithm lead to convergence in Appendix 3. Finally, we show empirical results that the algorithms lead to the required ideal balance by eliminating the stalling in the GPUs (the speedups in Figure 3 show that the idealess we observed in Figure 1 is eliminated).

In addition, the detailed description of this critical algorithm is not included in the main body, but in the appendix. At least some intuitions for why it can achieve load balancing should be discussed in Section 3.

We acknowledge that the readability can be improved by changing the structure to move the more important parts to the paper proper. We do however like to point out that is a structuring issue that could be easily addressed by re organizing the paper, and is not a fundamental issue with the methodology or results.

From a system perspective, re-packing dynamic models to fewer workers requires to reshard the model checkpoint, recompute the graph, and restart the training process, which will incur costly overhead up to several minutes for LLMs. However, this paper claims the overhead of DynMo is negligible, without discussing how it minimizes these incurred system overheads.

There are two things that are separate here: redistributing the model layers to load balance, and re-packing. For registering the model layers, the overhead (inclusive of profiling and migrating layers) is very low (single digit percentage) and hence is not an issue. For re-packing, we iterate that re-packing happens at a much lower frequency than rebalancing. While rebalancing can happen at the granularity of iterations, in practice it takes 1000s of iterations for the LLM to change enough to the point where re-packing would be beneficial (See Figure 4). Accordingly, the re-packing can be aligned to happen with the model restarts from a checkpoint since MTBF in practice is lower than the pace by which the model is re-packed. This makes re-packing much simpler and come for free since we piggyback on the restart process.

I think an implementation section is needed for fidelity.

Appendix D covers the implementation details. In addition the code of DynMo is open sourced, and it includes a detailed documentation.

Unclear algorithm descriptions: The second algorithm is an iterative decentralized diffusion-based algorithm. For the "diffusion", it is not clear how the diffusion is performed. The paper should provide a more intuitive example / explanation of the algorithm.

We added an intuitive explanation to Appendix C.3 (line 1275).

Omission of a Dynamic Baseline Configure: I think Megatron and Deepspeed are also support some naive dynamic baseline configure (manually). The paper should compare with them as a stronger baseline.

Megatron and DeepSpeed do not offer any functionality for dynamically balancing the model, not even with manual control, hence there is no baseline in Megatron or DeepSpeed to compare with (even with manual control). What DeepSpeed does simply offer is an API to dictate how to distribute the layer, which is in fact what we use to implement the redistribution of work load by DynMo.

I highly recommend changing the story to cluster-level scheduler instead of a training framework. It makes sense to automatically reduce some resources when some layer is freeze. Other jobs can better utilize the resource.

Most of the speedup reported in this work is attributed to balancing the load by redistributing the, and not due to the reduced communication when we re-pack: the speedup gain from re-packing is between 4~11% of the entire speedup gain. In other words, even if we do not re-pack at, the speedup gains remain almost the same. Hence we treat re-packing as an additional benefit of DynMo to be efficient by using less GPUs, and not as a way to speed up imbalanced dynamic training. We modified Figure 3 to include speedups when re-packing is used versus without.

The experiment is solid, but if I focus on MoE training case, I want to see some comparison with MoE-tailored system like Tutel https://github.com/microsoft/tutel

We added results that compare with Tutel in Figure 3 (line 727). DynMo significantly outperforms Tutel: 1.18x on Mixtral 8x7b and 1.21x on LLaMA-MoE-3.5B.

We would be glad to answer any questions about the paper.

a) it remains unclear how resources are isolated for safety purpose. For example, if NCCL communicator are shared, how could we make sure main job are safe as 3rd party job can do intentionally use it for other purpose.

First of all: excellent point!

Re-packing happens at a much lower frequency than rebalancing. While rebalancing can happen at the granularity of iterations, in practice it takes 1000s of iterations for the LLM to change enough to the point where re-packing would be beneficial (See Figure 4). Accordingly, the re-packing can be aligned to happen with the model restarts from a checkpoint since MTBF in practice is lower than the pace by which the model is re-packed. This makes re-packing much simpler (for example, we anyway launch an entirely new NCCL communicator upon restarting from a checkpoint).

In addition, if we wish to not-pack upon model restart, and sandbox the sub-communicator, there are several options (beyond the scope of this paper, yet we list since we agree with the reviewer that this is an interesting problem to think about):

• To add an extra layer of sandboxing: NCCL ranks validation to validate ranks within each communicator to ensure they only belong to the designated job.
• Implement access control in the NCCL runtime: If modifying the implementation is an option, one can enforce sandboxing at the runtime level by: a) adding an interception layer to wrap NCCL calls with validation checks to ensure processes cannot reference other communicators, or b) use custom communicator IDs that assign unique identifiers to each communicator and restrict access to these IDs based on process membership.

When repacking jobs into fewer workers, what are things we are serializing? Are they model states and optimizer states?

Yes: model parameters, model states, and optimizer states.

Do all the workers stop at the same time and start repacking? How do we deal with NCCL timeout if some rank are faster?

We schedule the re-packing to happen at the end of iterations since there is already a barrier at the end of iteration. That way we don't need to introduce any additional barriers.

We would be glad to answer any questions about the paper.

We thank the reviewer for the encouraging feedback.

In addition, the method proposed in this paper cannot automatically detect imbalance and must be manually configured by the user before training begins.

No manual configuration is needed from the user: the imbalance is done on regular intervals, i.e. we do not detect the imbalance, we instead relay on the fact that the load balancing overhead is very low, and just apply the load balancing at regular intervals. There could be an opportunity to further improve by detecting the imbalance, but that would improve the performance marginally.

I think the contribution of this paper is valuable and universal, so I suggest that the author can explain how their idea can be applied in an automated framework in the future, or even integrate it into existing mainstream frameworks.

This is indeed a priority for us. DynMo is fully automated, and requires no input from the user, hence can be add easily to existing frameworks. In fact we are actively trying to port DyNMo to use the new Megatron-core layer migration capability, and hence that would allow users who already use Megatron-core to use DynMo by adding a single line to their PyTorch code that activates the load balancing.