Accelerating Multimodal Model Training: LoongForge's DP Load‑Balancing Optimization Explained

1. DP Load Imbalance Limits Multimodal Training Efficiency

Current large‑language‑model (LLM) and multimodal model training relies on data parallelism (DP), where each node processes a subset of the data and synchronises gradients via AllReduce. Overall throughput depends not only on per‑GPU compute power but also on the consistency of progress across all DP nodes; any node that lags forces the whole system to wait during gradient synchronization.

1.1 DP gradient‑synchronisation bottleneck – The synchronization step amplifies any per‑node latency, reducing overall training speed.

1.2 Fixed‑length packing drawback – To avoid padding overhead, practitioners often use a fixed‑length packing strategy (e.g., 32K/64K/128K) that concatenates samples of varying original lengths into uniform‑length packs. While this equalises token counts across nodes, it ignores the quadratic complexity of the Transformer attention mechanism, which grows with sequence length. Consequently, two nodes with identical token totals can have vastly different compute costs if one processes longer sequences.

Example: with a 32K pack, DP‑0 receives eight 4K short samples, DP‑1 receives two 16K long samples. Both handle the same number of tokens, but DP‑1’s attention cost is dramatically higher, causing DP‑0 to finish quickly and then wait for DP‑1.

1.3 Multimodal models exacerbate the problem – Image and video modalities introduce additional load variance: differing resolutions, numbers of images per sample, and video frame counts lead to inconsistent compute and memory demands. The visual encoder (ViT) therefore suffers its own load‑imbalance alongside the LLM decoder.

2. Core Optimization: Adaptive Data Reallocation Based on Precise Compute‑Cost Modeling

The solution consists of two tightly coupled stages that are embedded directly into the native training loop, requiring no offline preprocessing.

2.1 Warm‑up Modeling Stage – Building an Accurate Cost Model

During a short warm‑up phase, an online performance probe collects per‑node execution time, sample lengths, and other key metrics without disrupting training. These observations are used to fit a DP‑level cost model comprising four terms:

Quadratic attention cost (captures the O(L²) complexity).

Linear layer and communication cost (approximately O(L)).

Fixed kernel‑launch overhead.

Coefficients x, y, z that depend on model architecture, training configuration, and hardware.

The collected data from n warm‑up iterations yields a system of equations. Because the max‑operator across DP ranks is discontinuous, the formulation replaces max with a differentiable softmax approximation and solves for the non‑negative parameters using least‑squares optimisation.

2.2 Online Adaptive Reallocation Stage – Dynamically Smoothing Node Load

After the cost model is obtained, each iteration evaluates the compute pressure of pending samples on every DP node. Samples are then reassigned across nodes to minimise the maximum total compute cost per DP, effectively flattening the per‑node runtime variance and reducing global waiting time during gradient sync.

The method works out‑of‑the‑box for major multimodal models such as InternVL and Qwen2‑VL series, supports both image and video data, and integrates seamlessly with existing DP, tensor‑parallel, and pipeline‑parallel strategies without code changes—only a command‑line flag is required.

2.3 Four Key Features

Dual‑module balancing for LLM decoder and ViT encoder.

Cross‑micro‑batch load tracking, achieving iteration‑level global balance rather than per‑micro‑batch local optimisation.

Intelligent trigger that skips reallocation when a single sample exceeds the average load or when overall load variance is below a threshold.

Asynchronous pipeline: data‑repacking runs in a dedicated pin_memory thread and overlaps with GPU computation via Gloo all_to_all, hiding any extra system cost.

3. Experimental Validation

Experiments were conducted without overlapping All‑Reduce. As DP scale increased from 32 to 512, raw throughput (TGS) steadily declined, with a sharp drop between DP‑256 and DP‑512 due to amplified load variance.

Enabling LoongForge’s DP load‑balancing mechanism mitigated this trend. Across all scales, throughput improved, with a 3.3% gain at DP‑256 and nearly 10% gain at DP‑512.

These results demonstrate that precise compute‑cost modelling and adaptive sample redistribution effectively eliminate the dominant synchronization bottleneck, substantially increasing training throughput and GPU utilisation, especially in ultra‑large‑scale clusters.