OmniBal: Towards Fast Instruction-Tuning for Vision-Language Models via  Omniverse Computation Balance
Authors: Yongqiang Yao
Jingru Tan
Feizhao Zhang
Jiahao Hu
Yazhe Niu
JinXin
Bo Li
Pengfei Liu
Ruihao Gong
Dahua Lin
Ningyi Xu
Venue: ICML 2025 poster (2025)
Keywords: VLM
Balance Training
Primary area: deep_learning->large_language_models
TL;DR: In this work, we effectively addressed the issue of imbalanced computation loads in large-scale 3D parallel training of vision-language models by rebalancing across data, model, and memory dimensions.

Abstract: Vision-language instruction-tuning models have recently achieved significant performance improvements. In this work, we discover that large-scale 3D parallel training on those models leads to an imbalanced computation load across different devices. The vision and language parts are inherently heterogeneous:  their data distribution and model architecture differ significantly, which affects distributed training efficiency. To address this issue, we rebalance the computational load from data, model, and memory perspectives, achieving more balanced computation across devices.  Specifically, for the data, instances are grouped into new balanced mini-batches within and across devices. A search-based method is employed for the model to achieve a more balanced partitioning. For memory optimization, we adaptively adjust the re-computation strategy for each partition to utilize the available memory fully. These three perspectives are not independent but are closely connected, forming an omniverse balanced training framework. Extensive experiments are conducted to validate the effectiveness of our method. Compared with the open-source training code of InternVL-Chat, training time is reduced greatly, achieving about 1.8$\times$ speed-up. Our method's efficacy and generalizability are further validated across various models and datasets. Codes will be released at https://github.com/ModelTC/OmniBal.

OmniBal: Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance
Yongqiang Yao
Jingru Tan *2 Feizhao Zhang *3 Jiahao Hu Yazhe Niu Xin Jin
Bo Li
Pengfei Liu'
Ruihao Gong 63
Dahua Lin 34
Ningyi Xu
Abstract
1. Introduction
Large language models (LLMs) have brought possi- bilities to many fields. Multi-modal models, particularly Vision-Language Models (VLMs) (Alayrac et al., 2022; Team et al., 2023a; Reid et al., 2024; Liu et al., 2023a; Bai et al., 2023b; Chen et al., 2023), are advancing rapidly due to their deeper understanding of the world. The train- ing scale of Vision-Language Models (VLMs) continues to expand, with increasingly larger datasets incorporating more text and higher-resolution images. Compared with the ILLVAA-15 (Liu et al., 2023a), the Intern VL.-Chat (Chen et al., 2024) has expanded the dataset size from 665K to 5M and increased image resolution from 336x336 to 3840x2160. At the model level, larger vision encoders are adopted. The InternVL-Chat upgrades the visual encoder from ~300M ViT-L-336px (Radford al., 2021) to ~-B Inter ViT-448xx (Chen et al., 2023).
Vision-language instruction-tunin models have recently achieved significant performance im- provements. In this work, we discover that large- scale 3D parallel training those models leads to imbalanced computation load across dif- ferent devices. The vision and language parts are inherently heterogeneous: their data distribu- tion and model architecture differ significantly, which affects distributed training efficiency. To address this issue, we rebalance the computational load from data, model, and memory perspectives, achieving more balanced computation across de- vices. Specifically, for the data, instances are grouped into new balanced mini-batches within and across devices. A search-based method employed for the model to achieve more bal- anced partitioning. For memory optimization. adaptively adjust the re-computation strategy for each partition to utilize the available mem- ory fully. These three perspectives are not inde- pendent but are closely connected, forming an omniverse balanced training framework. Exten- sive experiments are conducted to validate the effectiveness of our method. Compared with the open-source training code of InternVL-Chall training time is reduced greatly, achieving about 1.8x speed-up. Our method's efficacy and gen- cralizability are further validated across various models and datasets. Codes will be released at http:/giiibboooo///M///////////O
The larger datasets and models result in more time- consuming training process. Therefore, efficient training strategies are essential for the rapid advancement of the field. 3D parallelism (Shoeybi al., 2019; Rajbhandari et al., 2020; microsoft, 2020) is popular framework for large-scale distributed training. which allows data and mod- els to be distributed across multiple devices. Balancing computational load across devices is crucial in 3D paral- lelism by minimizing idle times.
In this work, we find that for instruction-tuning large vision- language models, the heterogeneous nature of data and model structures brings new challenges to 3D parallelism training: (1) Varying input sizes of LLM and VIT cause imbalanced computational loads across training iterations and devices. (2) The heterogeneity between LLM and VIT models leads to inherent differences in the computational load of their transformer blocks. Along with varying in- put sizes, this inevitably results in uneven computational load and computational bubbles. (3) Input size variation and computational imbalance compel to the most aggressive re- re-comutttion (checkpointing) (Li al., 2014) strategy to prevent program crashes, which wastes compu- tational resources. We refer to those issues caused by the heterogeneity in data and model structures large vision-
'Equal contribution 'Shanghai Jiao Tong University "Central South University "SenseTime Research 'The Chinese University of Hong Kong TToggji University Beihhan University. Correspon- dence Ruihao Gong <gongruihao @ buw..cdu.en.. Ningyi Xu <xuningyi@ @ sjtu.edu.cn..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

OmniBal: Towards Fast Instruction-Tunie Tuning for Vision-Language Models via Omniverse Computation Balance
language models as the Computation Imbalance problem, which reduces training efficiency.
Our approach is systematic rather than isolated and incremental, integrating the three interconnected mod- ules of data, model, and memory to ensure computa- tional balance across both inter- and intra-stage levels.
To address this problem, a simple and efficient training framework called Omniverse Balance (OmniBal) is pro- posed to balance computational load across multiple de- vices. This framework systematically balances computation in three bottlenecks, i.e. data, model, and memory, as shown in Figure 1. OmniBal works in these three closely connected aspects. Data lays the groundwork for addressing model imbalances, while data and model form the foundation for solving memory issues. Ultimately, these three aspects collaborate to achieve balanced computation. Data: The balanced dynamic mini-batch method is proposed to group instances as new mini-batches according to text length and number of images. Specifically, an iterative algorithm based on sampling and filtering combines data of different sizes into balanced groups, ensuring stable input sizes; Model: We propose balanced model partitioning to evenly spread the computational load of LLM and VIT across devices. Using a search-based approach, we efficiently find optimal partition strategies within a small search space, enabling adaptation different model architectures and hardware platforms. The balanced dynamic mini-batch method facil- itates balanced model partitioning by ensuring input sizes are consistent in advance. Memory: A balanced adap- tive re-computation method is proposed to optimize the re-computation strategy on each device, maximizing both memory utilization and training speed. We calculate the memory requirements of different models to adjust the re- computation strategy adaptively. Notably, our proposed balanced dynamic mini-batch and model partitioning en- sures balanced computational loads on each device, making memory analysis feasible.
Our method achieves a great speed-up while maintain- ing model performance on the open-source training code. Its efficacy and generalizability are further vali- dated across various models datasets and tasks.
2. Related Works
2.1. Multi-Modal Large Language Model(MILMM)
Large language models, such as ChatGPT (OpenAI, 2023a), GPT-4 (OpenAI, 2023b), Llama series (Touvron et al., 2023a;b; AI, 2024), and Gemini series (Team et al., 2023b; Reid et al., 2024), have seen significant advancements re- cently. They rely on large datasets for training to achieve strong performance, particularly in few-shot and zero-shot scenarios. Typically, they are built on textual data and can only accept text inputs. However, real-world scenar- ios often involve rich multi-modal information, e.g., im- ages. It has driven the development of large vision language models (VLMs). Visual encoders like Vision Transformer (ViT) (Dosovitskiy et al., 2021) usually incorporate vision information. A cross-modal connector is also required to align the vision encoder outputs to the language models. LLaVA (Touvron al., 2023a) uses the simplest MLP, BLIP series (Li et al., 2022; 2023; Dai et al., 2024) uses Q-former, Qwen-VL-Chat (Bai et al., 2023b) uses cross-attention module. VLMs expand large language models' capabilities and application scenarios by instruction-tunine with text and image data. However, introducing multi-modal data and heterogeneous encoders brings challenges to the training.
Extensive experiments are performed on various open- source VLM models at different scales, reducing overall training times significantly. GPU days are reduced for InternVL-Chat-115 (6+20B) from 61.8 to 21.3 under the Megatron--DeepSpeed (microsoft, 2020) backend. Scaling up to serve a great speed-up. from 75.4 to 30.5 GPU days. We conduct thorough generalization experiments, including var- ious datasets, hardware configurations, and multiple model combinations. Consistent and substantial improvements are observed across all experiments, demonstrating the effec- tiveness and versatility of our method.
2.2. Large-Scale Distributed Training
Distributed training is essential for efficiently utilizing mul- tiple GPUs to train large language models. It is achieved through 3D parallelism (Shoeybi al., 2019; Rajbhandari et al., 2020; microsoft, 2020): data, tensor, and pipeline parallelism. Data Parallelism splits the entire dataset into and assigns them to multiple devices, each with a model replica. This approach maximizes the use of GPU power for large datasets. DeepSpeed Zero (Rajbhan- dari et al., 2020) enhances it by reducing weight redundancy. However, it can still be challenged by the memory limits of individual devices when handling huge models. Tensor Par- allelism distributes model's weight matrices across mul- tiple devices, enabling parallel matrix operations (Shoeybi et al., 2019) and reducing per-device memory requirements. This method accelerates computation but requires dense deter- communication, typically restricted to single- node deployments to minimize latency. Pipeline Parallelism
(6+34B), we consistently ob-
In summary, our contributions are reflected in three levels: framework, method, and results.
We are the first to identify and address the computa- tional imbalance problem in large-scale VLM (Vision- Language Model) instruction- -tniing training, propos- ing a systematic solution tailored to this challenge.
2

---

OmniBal: Towards Fast Instruction-Tunie Tuning for Vision-Language Models via Omniverse Computation Balance
Computation Imbalanced Problem VLM Training
Model
What is the right part of the image
Data
VIT
MLP
LLM
Problem
Imbalance Data
Heterogeneous Model
Redundant
Memory
Balanced
Balanced Model
Balanced
Solution
Dynamic Batch
basic
Partition
basic
Adaptive Rep-com
Omniverse Balanced Training Framework
Figure Overview the computation imbalanced problem and our proposed solution the Standard Vision- Language instruction-tume tuninnnnnttnnggggg framework. We consider the bottleneck issues of data, model, and memory, and propose omniverse solution addressing these three aspects, each providing the foundation for the next.
as a sample.
Table 1. Analysis of computation imbalance.
Dimension
Inpat
Forward time
Cost memory
Data Imbalance: LLMs are trained on texts using next- token prediction, allowing consistent input lengths through arbitrary text sub-strings. In contrast, VLMs handle texts and images, requiring data integrity, and preventing arbi- trary truncation. The varying number of images, resolutions, and text lengths result in considerable differences in input sizes across mini-batches. From Table and Figure 2, data imbalance occurs in Inter-Stage and Intra-Stage. To better quantify the impact of dynamic input, we define the DistRa- tio (introduced in Section 4) measure the degree of data imbalance of VIT and LLM.
Inter-Stage
LAK 0.ttoken
23
Intra-Suge-1
1.K 1.2Kloken
136 + 155
7
divides model into segments and assigns them differ- ent devices, creating computation flow like production line. This technique facilitates larger model scaling across nodes. GPipe (Huang et al., 2019) proposes micro-batching to decrease forward bubbles. PreeDream (Narayanan et al., 2019) further proposes one-forward-on wwwrrrrddo--bbbwwweeeee (IFIB) scheme to optimize memory usage. In pipeline parallelism, uneven layer partitioning can cause significant pipeline bub- bles. PipeDream (Narayanan et al., 2019) and AdaPipe (Sun et al., 2024) optimize model partitioning and re- re-computation strategies based on profiling and dynamic programming, respectively. However, these advancements are primarily tested in text-based models and may require adaptation for large vision- language model scenarios.
Data Imbalance Evidence: As shown in Figure at time T-0, the Vision and LLM inputs for DP-0 (Group 0 of data- parallel) and DP-1 (Group I of data-parallel) differ (Img=4, Text=2k vs. Img=9, Text=4k). Different inputs will bring different computational complexities. For example, the total forward time of DP-0 at time T-O 970ms (500 160 + 150 160), and the total forward time of DP-1 time T-0 is 1730ms (900 300 260 270). DP-0 needs wait for DP-1 to complete all training before updating parameters, which will cause DP-0's GPU resources to wait for a long time. This is the Intra-Stage imbalance problem across de- vices. For DP-0, the inputs at times T-O and T-1 on the same device also differ significantly (e.g., Img=4, Text=2k ver- sus Img=20, Text=16k), which is an Intra-Stage imbalance problem over time. This results in substantial variance in forward time and memory usage, creating challenges in determining a globally optimal model partition strategy.
3. Computation Imbalance
In this section, we explore the unique challenges of large- scale distributed training for laniun-language models, fo- cusing on two dimensions: Inter-Stage and Intra-Stage computation imbalance. Inter-Stage means the computa- tion imbalance of different pipeline parallel stages. Intra- Stage indicates the computation imbalance of the same stage across time and devices. Figure 2 shows these compu- tation imbalances more intuitively. And they both include three specific levels: data, model, and memory. To quantify this problem, we used the ViterVL-Chhee 1.2 dataset (Chen et al., 2024) to perform profile statistics shown in Table For the Intra- Strae,tage, we counted the information of Stage
Model Imbalance: LLMs use identical transformer mod- ules with the same computational load. Evenly dividing these layers pipeline parallelism distributes the load ef- fectively. However, VLMs require additional image pre-
3

---

OmniBal: Towards Fast Instruction-Tuning for Vision- Vision-Language Models via Omniverse Computation Balance
VLM Training Steps
for
T-0
Dist Rotto for T-1 VLT LLM 0.81
$130-3
COOLI
Stagen3
LM
LIM
NIT
DP-0
IMG
160eG
158mG
160ss
NMM 20
2000mss
1680ms
1800
TIME
NEM S0G
MEN
NEM
NEW
796
SStare
Inter- Stage
TIME 148mm HEM
STD TIME
40E/es MEN 4, 3G
Soagess
Trimk
N1T
LLP
LLe
DP-1
LLP
IMG HUM 9 TEXT NIM 4K
TINE 980es NEM 686
TIME 300es MEM 6B6
TIME 250es MEM 566
TIME 270 NEM 50G
IMG
TIME 200es
60G
42G
38G
Inter Stage 27trs 6.5G
Inter- Stage
STD TIME
Data Parallel
Inbro:
Intra- Stage?
Intro: Stage3
Intro-Staaga 344es
MEN 11.0G
NEM
TIME 345es MEN 13.76
Groups
Figure 2. The Problem of Computation Imbalance VLM Ciauuciion-Tunigg Training Pipeline. DP-0 and DP-1 represent different Data Parallel processes. T-0 and T-1 represent different training times. TIME and MEM represent forward time and cost memory in the current stage, respectively. STD stands for standard deviation.
processing, necessitating an image encoder. The structural disparity between VIT and LLM results in different compu- tational demands. From Table and Figure 2, the standard deviation of forward time is huge in both Inter-Stage and Intra-Stage, indicating a serious computation imbalance.
to prevent the program from crashing. However, excessive re-computation can slow down the training process. From Table and Figure 2, memory imbalance is reflected in both Itter-Stage and -ttae in the existing training setting.
Memory Imbalance Evidence: As shown in Figure 2, the memory consumption for Stage-1 I at different time points is as follows: 50 GB at T-0, 80 GB at T-1 for DP-0, 68 GB at T-0, and 60 GB at T-1 for DP-1. To prevent program crashes, we must configure the re-computation strategy based on the highest memory usage, which is 80 GB. This approach in- troduces additional computational overhead, categorized as the Intra-Stage imbalance problem (over time). Addition- ally, for DP-1, the memory usage across different stages at both T-0 and T-1 has a high standard deviation (6.5 GB and 8.4 GB, respectively). Using the same re-computation strategy for different stages will also bring computational overhead, according to the highest memory usage. This variation represents the Inter-Stage imbalance problem.
Model Imbalance Evidence: As shown in Figure 2, for DP-0, at T-0 and T-1, the standard deviation (148ms and 402 ms) of the forward time across different stages is significant, leading to serious bubbles in pipeline parallelism and slowing down training speed. This is the Inter-Stage imbalance problem. Furthermore, the computation distribution across stages shows a huge difference between T-0 and T-1. For example, the computation distribution at T- 0 is (500:160:150:10)) (0:0:001000000005222200 16:00.1:000...... while the computation distribution at T-1 is (2000:1600: 1000: 100)= 110.:3.......0.111.111.111.111 1,00. varying computation distribution makes determining the optimal model partition strategy challenging. This includes the Inter-Stage and Intra-Stage problems (over time).
Differences between VLM and LLM training: As men- tioned above, the difference between VLM and LLM arises from the data composition and model structure, resulting in unique Inter-Stage and Intra-Stage challenges. Inter-Stage: Since LLM has a fixed structure, the model can be equally divided, and there no Inter-Stage imbalance for any input. Dynamic input or inconsistent text-image ratios in heteroge- neous VLM will lead to an Inter-Stage imbalance problem. Intra-Stage: For the LLM-Pretrain task, the input is fixed, and there is no Intra-Stage imbalance problem. Dynamic input can be converted into static input by simple packing
Memory Imbalance: LLMs require significant GPU mem- ory due to their large parameter size. When memory is insufficient, re-computation (Li et al., 2014) techniques dis- card some intermediate activation values and recompute them during backward propagation to save memory. VI.M encounters great memory challenges due to the variable scales of data inputs and the heterogeneity between vision and language models. The presence of numerous images or long text inputs can lead to excessive GPU memory us- age, requiring the most aggressive re-computation settings
4

---

OmniBal: Towards Fast Instruction-Tunine for Vision-Language Models via Omniverse Computation Balance
(Kosec al., 2021) reduce the computation imbalance for the LLM-SFT task. However, VLM instruction-tuning training cannot rely on simple packing to ensure fixed in- puts for VIT and LLM, resulting in computation imbalance problems.
RadRo Σ (Imaz - 14) Imar X B
Where bmaz represents the maximum number of tokens in a mini-batch of size B, and t, denotes the number of tokens for sample i within that mini-batch.
4. Method
Distribution Ratio (cross-device); Even after padding, the sizes of mini-batches on different devices may vary, leading to different input scales across devices. The distribution ratio is calculated as follows:
This section presents our cumputation-balanceer framework, OmniBal, for training large vision-language models. To address an imbalanced computational load across devices, we first manage the large data input variations, which is the most fundamental issue in the computation imbalance problem. This enables balanced model partitioning. Finally, the re-computation strategy for each partition is optimized. Appendix A shows our training pipeline.
E. - Ti)
DistRaits
(2)
T Tmas X N
Where N represents the number of devices, Tmaz denotes the maximum number of mini-batch tokens across all de- vices, and T₃ refers to the number of mini-batch tokens on the 8th device. Non-fixed input sizes in VLMs have larger Pad Ratio and Dist Ratio, as shown in Table 3 (row 1). A high Pad Ratio wastes computational resources, while a high Dist Ratio causes device idle time. They significantly impact training throughput efficiency.
Algorithm/ 1 ISF: Iterative Sampling and Filtering
I: ISF: Sampling Stage
2: D = randperm(D), set 9 = =[]
3: for (I,, v/) in DD do
4: g < 9 + (Ii, yi)
5:
I, > Qu or 4 > Q: then
To address this issue, an adaptive grouping strategy that organizes multiple samples, ensuring that both image and text sizes in the resulting groups remain within a relatively fixed range is implemented. We refer to this method as the Balanced Dynamic Mini-Batch. Determining the opti- mal grouping strategy is non-trivial problem, An iterative method is designed using sampling and filtering to group samples. As illustrated in Algorithm 1, it Iterative Sam- pling and Filtering (ISF) involves the following steps:
6:
g 9 (I,; y4)
7: P P 9. set G [()))))))
8: end if
9: end for
10: return: P
11: ISF: Filtering Stage
12: Get P from Sampling Stage
13: for G in p do
14:
if T < Q. and I; < Q2 then
1. Sampling Stage: For current dataset D {( 36) i). we randomly add samples d, consisted of images x4, text yi to current group G. If the total number of images I. Ex.e0 |a:|| or the total text length I, = Ey:eg [yi reaches the predefined maximum number of images Qu or text Q.. we add this group to the candidate set P and create a new group containing (a,, y4) for the subsequent samples. Other- wise, we will continue adding samples to the current group. At the end of the sampling stage, we will have candidate set P = ( = 3..
15: P p - G
16: else
17: remove all (Ii; vk) of g from D
end if
19:
end for
20:
return P.D
4.1. Balanced Dynamic Mini-Bated
2. Filterrin Stage: We first define the target number of im- ages Q', and text Q'. For each group G, in candidate set P, we keep g whose image number J, or text length I, satisfy Lg >= Q or L, >= Qú, and remove all samples (z4, y1) in that group from D. Otherwise, we remove non-satisfied Gi from P. Ultimately, P becomes our target set, and D becomes our updated dataset for the next iteration.
For instruction-tuning VLMs, each training sample contains various images and texts, resulting in non-fixed input sizes. We evaluate data imbalance from two perspectives: within- device samples and cross-device mini-batches.
Pad Ratio (within-device): When combining samples of different sizes into a mini-batch, smaller samples need to be padded to ensure aligned input sizes. The Pad Ratio is calculated as follows:
The sampling and filtering stages are alternately repeated for a maximum of T times. The candidate set is acquired P each time, which includes more valid sample groups G.
5

---

OmniBal: Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance
Meanwhile, we have the updated dataset D consisting of unselected samples, which is used for sampling and filter- ing in the next iteration. To ensure that the mini-batches constructed by the ISF method achieve lower Pad Ratio and Dist ratio, appropriate values for Qu and text Q need to be determined. The optimal values for Qu and Qi vary across different datasets. In practice, A statistical approach described in Section 5.1 is used to determine these values.
cient evaluate the running time for each one. Therefore, two metrics to rank these candidates are designed.
The first metric is the difference in running time between different pipeline stages S,, Smaller differences generally result in fewer bubbles and faster execution. We use the variance of the running times of different pipeline stages to measure this difference.
VAR(fwd_ttmm)) Σ (FWI((S)) FWD(S.))² (4)
4.2. Balanced Model Partitioning
Given the number of layers L in the model and the pipeline parallel size N. our goal is to find an optimal partition strategy P = (P(1) p(2), p(3), 3.... ....,,((-11)) P(N-1)) such that the training speed of the model is maximized. Here, p(1) < p(2) < p(3) <... < P(N-1), and the h partition stage S, consists of layers 4. where p(i-1) < 4 < p()), with p(o) = and p(N) = I +++ 1. For example, given a model with L- = 20 layers and pipeline size N = 4. assume that we have an optimal partition P = (5, 10, 15). The first partition St consists of layers 1, 12,22, 4 since p(0) 1, P(1) = 5.
The second metric the total point-to-poinn communica- tion volume of the partition strategy P.. It depends on P. consisting of (P(1), P(2), p(3),...
SUM(comm) = Σ ACTIV((p))
(5)
Where lpc is the last layer of partition strategy p()) and ACTIV (Ip;) is the activation number of layer Ips, indicating the point-to-point communication volume of p(t), We use the sum of VAR(fwd time) and SUM(comm) as the metric for the partition and rank them to select the top K candidates for speed evaluation.
However, achieving balanced pipeline partitioning for VLMs is more challenging task compared to LLMs. We must consider: (1) Model Heterogeneity: The structural differences between visual and language models make sim- ple parameter-based or layer-based partition strategies in- effective. (2) Communication Overheads: Different par- titioning strategies result in varying communication vol- umes, as the number of activations in each layer can differ significantly in VLMs. (3) Hardware Variability: Differ- ent platforms exhibit varying levels of capability, particu- larly in terms of communication overhead. On platforms with high network bandwidth, communication overhead can be negligible. Based on the above analysis, heuris- tic search algorithm to find the optimal partition is devel- oped. We first identify candidate set of partition strategies (P& = (P p(2)) p(3) ,...... - k = 1,,,3,...
4.3. Balanced Adaptive Re-Computation
Thanks to the balanced dynamic mini-batch and balanced model partition, balanced computational load is main- tained across each pipeline stage. The memory require- ments are now stabilized as the computational demand has been fixed. As a result, we can optimize the re-computation strategy based on actual memory needs, rather than relying on the most aggressive approach to avoid crashes. Reduc- ing the number of re re-computations accelcrates the model's backward pass, leading to faster training speed.
We find that heterogeneous architectures have different memory requirements. For examplc, the vision model in Intern 1.5 requires more GPU memory than the lan- guage model under the same computational load. Therefore, it is necessary to analyze the memory requirements of each layer in the vision and language models individually and adaptively determine the optimal re-computation strategy for each layer.
that possibly contain the optimal one. Then, the optimal partition strategy p"is selected based on the running time:
p = argminp, p. f(P.)
(3)
Here, f(P.) is the average rumning time obtained by training the model for several iterations.
Partition Candidates: We start by profiling each layer's computation time (WD. A greedy algorithm is em- ployed to compute the anchor partition strategy pt, making the computation time of all partition stages S, close. Around pt, A candidate set of partition strategies is created by jit- tering p(1), p(2) ...., 11N within radius of r. When r r= 1 and N- = 4, there are a total of 33 = 27 candidates.
Balanced Adaptive Re-Computation Strategy. In this context, Q and Q. represent the inputs for Vision Trans- former (VIT) and Large Language Model (LLM), respec- tively. M, denotes the remaining GPU memory at the cur- rent stage, while M, and M, indicate the GPU memory saved by each transformer layer of the L.M and VIT when enabling re-computation.
Partition Metrics: When r and N are very large, there will be a vast number of partition candidates, making it ineffi-
Step-1: Given the inputs Qu and Qt, we enable the re- computation strategy across all transformer modules of the
6

---

OmniBal: Towards Fast Eestruction-Tuningg for Vision-Language Models via Omniverse Computation Balance
Table Main Results. We open-source IntemmV.-Chat-1.5 6+20B and 6+34B the models with either DeepSpeed (ZeRO-3)
Megatron- Deeerspee backend. GPU Days are reported in the Invern Vh-hat 1.2 1.2M training dataset to show the speed-up ratio. Models are also evaluated on five commonly used benchmarks.
Model
Balance?
Backend
MMB-EN/CN
CharQA
A12D
MMVet
MME
GPU Days (speed-up)
DeepSpeed
78.2/77.4
86.2
71.3
48.9
1901.2
38.9 (1.00))
DeepSpeed
78.7/77.6
1969.4
25.3 (1.5x)
6+20B
X
Megatron
79.5/77.7
87.3
71.6
45.0
1957.7
61.8 (0.63x)
Megatron
78.6/77.5
70.9
1956.3
21.3 (L.8x)
X
DeepSpeed
80.0/79.2
86.6
73.4
45.9
2015.8
54.3 (1.00x)
DeepSpeed
80.9/79.0
73.3
47.0
2153.6
35.5 (1.53x)
6+34B
Megatron
80.2/79.3
88.9
73.7
44.2
2111.9
75.4 (0.72x)
Megatron
80.1/78.0
89.3
73.5
45.4
2072.7
30.5 (1.80x)
model. At each forward pass, we clear the cache and record each stage's 's remaining memory usage M,..
evaluation metric to estimate the total training time. Specifi- cally, GPU Days are reported based on A100 GPU usage evaluate the speed-up performance.
Step-2: We manually disable re-computation for some lay- ers based on the remaining GPU memory. Subsequently, we record the GPU memory usage Mr for each stage.
5.2. Implementation Details
How to get Qu and Qt- We determine Qu and Qe using dataset statistics. The total text token lengths and image count are used to compute the average tokens per image. Qi is set to the longest text token length, and the text-to- image ratio determines Q It is challenging to maintain an exact number of images and text length, so we relax these conditions allow for approximation. For images, Q₂ Qu. and for text, Q/ = Qt 128 based on results in Table 6. In the Intern VL-Chat- 1.2 dataset, Q2 4K, Qu = 9, with each image processed into 1K tokens for VIT.
Step-3: Based on the memory differences AM, observed between Step-1 and Step-2, along with the re-computation strategy implemented each stage, estimate the memory savings M, and M. for each transformer layer of the VIT and LLM, respectively.
Step-4: Based on the estimated GPU memory savings M, and M,, measured in Step-3, as well as the remaining mem- ory Mp from Step-1, we first estimate the theoretically op- timal re-computaiion strategy for each stage and conduct the training test. If the test runs successfully, we adopt this strategy. If it fails, we incrementally increase the number of re-computation layers by the remaining GPU memory for each stage.
Sampling Overhead. The sampling overhead is minimal. For instance, with 1.2 million samples, sampling can be completed in just a few tens of seconds, adding negligible overhead the overall training process. The time complex- ity of ISF is OCC (N M)), where N is the total number of samples, M is the number of samples per pack, and C' is the number of iterations.
5. Experiments
In this section, the models and datasets are introduced. Then, demonstrate the acceleration compared the current state-of-the-art VLMs. Subsequently, a detailed comparison of each component proposed in our method is presented, highlighting its specific contribution to training acceleration. Finally, extensive experimental analysis conducted.
Partition Overhead. Establishing the anchor requires pro- filing layer-wise forward time and activation values, which takes only steps and incurs minimal overhead. While the anchor may not be optimal, we observe empirically that the true optimum typically lies within < 3 of it.
Optimal Solution Overhead. To refine the anchor, we measure the speed of the top 10-15 nearby partitions (as in Section 4.2), each requiring just steps. This adds only a minor cost relative to the full training process.
5.1. Experimental setup
Model & Dataset setting: We conduct experiments follow- ing the open-source IntermVL-Chat-1.5 setting. Our vision and language models are InternViT-6B and InternLM2-20B, respectively. Two configurations are employed: Intern VL- Chat-11.5 (6+20B) and Intern -iieemm---Chha-----PPi (6+34B). As the Mh-ehtL-Chha--ae dataset is not yet available, utilize the Intern VL-Chatt- dataset, which comprises approxi- mately 1.2 million samples, as an alternative. All other training settings remain unchanged. GPU Days are our
5.3. Main Results
We demonstrate the superiority of OmniBal under various settings in Table 2. Baseline model is InterVV-CChat- (6+20B) (Chen et al., 2024), with DeepSpeed ZeRO-3 back- end. OmniBal reduces GPU days from 38.9 25.3, achiev- ing 1.54x speed-up. Simultaneously, we maintain com-
7

---

OmniBal: Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance
Table Importance of data balance. AVE-BS indicates average batch size.
We report results with Model Balance (MB) and without MB.
Max Max-Seq-Lan Seq-Len
Dist Ratio
GPU Days
Method
AVE-BS
Pad Ratio
Balanced
VIT
LLM
VIT
LLM
w/o MB
MB
baseline
4
20K
16K
0.31
0.34
0.30
61.8
42.2
length-group
4
20K
16K
0.20
0.26
0.13
X
54.0
40.0
device-group
4
20K
16K
0,378
0.125
0.15
X
54,5
43.6
ISF(ours)
4.6
9K
4K
0.02
0.14
51.9
29.0
Table 4. Importance of model balance. VAR indicates variance. SUM(comm) is the summation of commutation volume (MByte)
Method
VAR(param)
VAR(num_layer)
VAR(fwd.time)
A SUM(comm)
GPU days
(1) parameter-base based
0.03
13.4
93.6
+0.0
42.2
(2) layer-based
0.64
1.2
20.1
+8.2
30.6
(3) profile-based
0.85
+16.6
30.9
(4) BMP (ours)
0.83
1.5
12.2
-21.0
29.0
parable performance across commonly used datasets, such as BMMB-EN/CN (Liu et al., 2023c), ChartQA (Masry et al., 2022), A12D (Kembhavi et al., 2016), MMVet (Yu et al., 2023), and MME (Fu et al., 2023). Experiments are conducted, which integrates tensor, pipeline, and data parallelism for larger-scale models. However, directly applying 3D parallelism can slow down training due the heterogeneous nature of VLM models. Table 2 shows that switching to Megatron-DeeppSped in- creased GPU days from 38.9 to 61.8. OmniBal addresses this issue by achieving computational balance across data, model, and memory, reducing GPU days from 61.8 to 21.3. This demonstrates the importance computational balance for effective 3D parallelism. Notably, our method also out- performed DeepSpeed, highlighting the superiority of 3D parallelism when balanced computation is achieved. Results under larger-scale setting (Intern VL-Chat-1 1---Chhat-.--Piuue .5-Plus) are also reported to verify the generalizability of our method. The larger model consistently improves, accelerating the training process while maintaining model performance.
The Importance of Data Balance: In Table 3, inves- tigate the importance of data balance in large-scale dis- tributed training by comparing four methods: (1) Baseline: Randomly combining data into a mini-batch with padding aligned to the longest input within mini-batches (2) Length- Group: Combining samples with similar text and image sizes into a mini-batch to minimize padding within mini- batches. (3) Device-Group Grouping samples with simi- lar input sizes across devices to minimize idle times. (4) Balanced Dynamic bati-b Using ISF to construct bal- anced
with
within mini-batches and across devices.
Table reveals the following: (1) Baseline: is the slowest due to the completely random combination of different-sized samples, leading significant size variation and excessive padding (0.31). Meanwhile, high Dist Ratio ViT (0.34) and LLM (0.30) result in computation disparities between de- vices, severely impacting throughput efficiency. (2) Length- Group: enhances throughput efficiency by pre-grouping samples of similar sizes into batche, thus reducing the internal padding ratio (0.2). Minimizing the number of redundant tokens within mini-batches effectively lowers the GPU days required to 54.0. (3) Device- Proice:-rou:: reduce idle time by ensuring consistent input sizes across devices. It decreases the Dist Ratio of ViT (0.125) and L.LM (0.15). However, it only balances input sizes between devices and neglects the balance within mini- bini-batches. High padding (0.378) wastes computational resources. (4) Our Approach: balances input sizes within mini-batches on each device and across devices simultaneously. It reduces both the Pad Ratio and the Dist Ratio, achieving a padding ratio of O while maintaining lower Dist Ratio of 0.02 and 0.14. While our method balances input sizes, model partitioning still limits training speed. With model balance (MB), GPU days are reduced from 42.2 to 29.0, O, a gain of 13.2, compared to 9.9 without MB (from 61.8 to 51.9). This underscores the importance of holistic balance approach.
5.4. Ablation Analysis
In this section, ablation experiments on each component of our method are conducted, using IntemVL-CCha-L1. as the baseline model with a 3D parallel Ergatron-DeepSperr backend. Table 7 illustrates the impact of each component. The baseline model experiences considerable slowdown in training speed due to computational imbalance, necessi- tating total of 61.8 GPU days. By achieving data balance, GPU days are reduced to 51.9. Data balance allows us to achieve a more balanced model partition, reducing the training time. Finally, optimizing memory with an adaptive re-computation strategy reduces GPU days to 21.3. These re- sults demonstrate that holistic balance encompassing data, model, and memory is crucial for efficient VLM training. Below, we provide a detailed analysis of each component.
8

---

OmniBal:
Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance
Table Importance of memory balance. VRAM, denotes remaining VRAM(G) in pipeline stage S, For the haseline model, the metric
varies as <minimum>> ~ maximum>>
Method
V-Seq-Len
L-Seq-Len
VRAM;
VRAMa
VRAMs
VRAM,
GPU Days
buselinc
4K~20K
1K~16K
13~50.2 58.2
7.3r.40.5
7.3rw40.5
7.3r.40.5
61.8
data & model balance
9K
4K
56.2
32.5
29.0 32.7
memory balance
9K
4K
12.3
21.7
24.7
30.0
21.3
training by adjusting the re-computation strategy to fully utilize the remaining memory, reducing GPU days to 21.3.
Table 6. Ablation results of Qu, Q
Qu
Qt
DistRatio VIT DistRatio LLM
Qu
Q: - 64
0,0159
0,145
5.5. Generalization Capability
Qu
Q: - 128
0.0156
0.144
We study the generalization capability of our method from multiple aspects: (1) Different Datasets like LLava-665K,
Qu
Q: 256
0.018
0.141
Qu
Qt -28
0.0417
0.20
VLternVL-112MM, and LCS558K in Appendix B; (2) Different Models: verify our methods on different vision models like EVA-CLIP (Sun et al., 2023) and different language models like Llama3 (AI, 2024), Yi-34B (NousResearch, (Bai et al., 2023a), as de- tailed in Appendix C; (3) Different High-Resolution Setting: Under various settings, we achieved a speedup of approxi- mately 2.0x, as demonstrated in Appendix D; (4) Different Tasks: Besides SFT tasks, pretraining tasks are also tested, shown in Appendix E, and we observed consistent im- provements across all settings; (5) Different Image Reso- lutions: As shown in Appendix F, our method consistently delivers a highly satisfactory acceleration effect with differ- ent input image resolutions; (6) Convergence of ISF: We show that ISF can converge in few steps on other datasets, in Appendix G: (7) Different Model-series like LLava- in Appendix H; (8) Pre-Processine Strategy: Qwen2-VL (team, 2024) is a recent, highly regarded open-source project that provides strong support for dynamic image input. We adopt the pre-processing strategy of Qwen2-VL to validate our method in Appendix I; (9) Different Hardware like H100, A100 in Appendix J; (10) Larger-Scale Results. To validate our method, we conduct experiments with larger models and 512 GPUs Results are provided in Appendix K. These results underscore the effectiveness and robustness of our method across a wide range of datasets, models, and tasks.
Table 7. Ablation studies of components
halance
model balance
memory halance
GPU Days
2023), and the Qwenl.5-1
61.8
51.9
29.0
V
21.3
The Importance of Model Balance: Table examines bal- anced model partitioning, focusing on partition strategies for pipeline parallelism. For LLM training. common meth- ods include (1) parameter-baae based and (2) baser, (3) profile-based methods such as DreamPipe (Narayanan et al., 2019), which estimate the computation time for each layer, and use this information to partition the model effectively. Additionally, (4) our search-basee Balanced Model Partition method finds the optimal partition strategy from a set of can- didates. As shown in Table 4, (1) Parameter-based and (2) myyer-based methods split the model's parameters layers across devices, achieving low variation in VAR(param) and VAR(num_layeee However, they still show high variation in forward time VAR(fwd time), leading to computational inefficiencies in the pipeline. (3) The profile-based method ensures the optimal VAR(fwd time). However, this parti- tioning occurs before the vision model's token sub- sabmligg operation, increasing communication overhead and affect- ing training speed. (4) Our proposed BMP method explores high-quality partition strategy space to identify the optimal strategy, achieving the best results in 29.0 GPU days.
6. Conclusion
In this work, we effectively addressed the issue of imbal- anced computation loads in large-scale 3D parallel train- ing of vision-language models by rebalancing across data, model, and memory dimensions. Experimental results demonstrate that our method can significantly speed up train- ing on many open-source models. The effectiveness and gen- cralizability of our approach are also validated across vari- ous models, datasets, and hardware platforms. Our method can accelerate the development of this field by enabling more efficient training.
The Importance of Memory Balance: Table 5, we ex- amine the significance of memory balance. In the baseline model, varying input sizes for vision (4K-20K tokens) and language (1K-16K tokens) lead to varying GPU memory usage. Despite aggressive re-computation, the remaining memory can drop to 7.3 GB. ISF and BMP improve train- ing speed by controlling computational load across devices. However, memory demands still varied, e.g., GPUs and having more remaining memory. ISF further speeds up
9

---

OmniBal: Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance
Acknowledgements
P., Lu, T., Qiao, Y., and Dai, J. Internvl: Scaling up vision foundation models and aligning for generic visual- linguistic tasks, arXiv preprint vrXiv:2312.14238. 2023.
This work was supported by the National Key Research and Development Program of China under Grant No. 2023YFB4405102. It was also supported in part by the Natural Science Foundation of Hunan Province under Grant No. 2024JJ6525. We would like to thank Kaihuan Liang, Shihao Bai, and Shuo Wu for their support on this work.
Chen, Z., Wang, W., Tian, H., Ye, S., Gao, Z., Cui, E., Tong, W., Hu, K., Luo, J., Ma, Z., et al. How far we to gpt-4v? closing the gap to commercial multi- modal models with open-source suites. arXiv preprint arXiv:2404. 2tfi:2404..6821, 2024.
Impact Statement
Dai, W., Li, Li, D., Tiong, A. M. H., Zhao, J.. Wang, W., Li, B., Fung, P., and Hoi, S. Instructblip: towards general- purpose vision-language models with instruction tuning. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2024. Curran Associates Inc.
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
AI, M. Introducing meta llama 3: The most capable openly available IIm to date. .htpps:/ // /aai/ain.. blog/meta- lama- 3, 2024.
Alayrac, J.-B., Donahue, J., Luc, P., Miech, A.. Barr, I, Hasson, Y., Lenc, K., Mensch, A., Millican, K., Reynolds, M., Ring, R., Rutherford, E., Cabi, S., Han, T., Gong, Z., Samangooei, S., Monteiro, M., Menick, J. L., Borgeaud, S., Brock, A., Nematzadeh, A., Sharifzadeh, S., Binkowski, M. a., Barreira, R., Vinyals, O., Zisserman, A., and Simonyan, K. Flamingo: a visual language model for few-shot learning. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 23716-23773 Curran Associates, Inc., 2022. URL https: http::::///pooceeiggg.ccoo neurips. cc [raper_filees/ppper//22////////1
Fu, C., Chen, P., Shen, Y., Qin, Y., Zhang, M., Lin, X., Qiu, Z.. Lin, W., Yang. J., Zheng, X., et al. Mme: A compre- hensive evaluation benchmark for multimodal large lan- guage models. arXiv preprint 23..v:2306.13394 2023.
Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Kembhavi, A., Salvato, M., Kolve, E.,, Seo, M., Hajishirzi, H., and Farhadi, A. A diagram is worth a dozen images. 235-251, 2016.
960a1 960a1 (ataaa/c67Bf111777cccc1111111110 pdf.
Po te
Kosec, M., Fu, S., and Krell, M. M. Packing: Towards 2x NLP BERT acceleration. CoRR, 0201110070222.... 2021. URL
Bai, J., Bai, S., Chu, Y., Cui, Z., Dang, K., Deng, X., Fan, Y., Ge, W., Han, Y., Huang, F,, Hui, B., Ji, L., Li, M., Lin, J., Lin, R., Liu, D., Liu, G., Lu, C., Lu, K., Ma, J. Men, R., Ren, X., Ren, X., Tan, C., Tan, S., Tu, J., Wang. P., Wang. S., Wang, W., Wu, S., Xu, B., Xu, J., Yang, A., Yang. H., Yang, J., Yang, S., Yao, Y., Yu, B., Yuan, H., Yuan, Z.,, Zhang, J., Zhang, X., Zhang, Y. Z.anag. 7 Zhou, C., Zhou, J.. Zhou, X., and Zhu, T. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023a.
/tttp:://axxv..rgg/// /
Li, J., Li, D., Xiong, C., and Hoi, S. Blip: Bootstrapping language-iman pre-training for unified vision- understanding and generation. In ICML, pp. 12888- 12900, 2022.
Li, Li, D., Savarese, S., and Hoi, S. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. In ICML, pp. 19730-19742. PMLR, 2023.
Bai, J., Bai, S., Yang, S., Wang, S., Tan, S., Wang, P., Lin, J., Zhou, C., and Zhou, J. Qwen-vl: A versatile vision- language model for understanding, localization, text reading, and beyond. arXiv preprint (tXXiiv338..1966 2023b.
Li, M., Andersen, D. G., Park, J. W., Smola, A. J., Ahmed, A., Josifovski, V., Long, J., Shekita, E. J., and Su, B.-Y. Scaling distributed machine learn- ing with the parameter server. In 11th USENIX Symposium on Operating Systems Design and Imple- mentation (OSDI 14), pp. 583-598, Broomfield, CO.
Chen, Z., WU, J.. Wang. W.. Su, W., Chen, G., Xing, S., Zhong, M., Zhang, Q., Zhu, X., Lu, L., Li, B., Luo,
10

---

OmniBal: Towards Fast Instruction-Tunine Tuning for Vision-Language Models via Omniverse Computation Balance
October 2014, USENIX Association. ISBN 978-1- org/ 931971-16-4. URL stsechoonallssssionss//// sessions/ conference/osdii wtwpeeeeeeeeccccaaaeeeeeeeeeeeee presentation/Is mu.
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-Im: Training multi- billion parameter language models using model paral- lelism. arXiv preprint 1XXi::999000052.. 2019.
Liu, H., Li, C., Li, Y., and Lee, Y. J. Improved basclines with visual instruction tuning, 2023a.
Sun, Q. Fang, Y., Wu, L.., Wang, X., and Cao, Y. Eva- clip: Improved training techniques for clip scale. arXiv preprint arXiv:2303.
2023.
Liu, H., Li, C., Wu, Q., and Lee, J. Visual instruction tuning. NeurlPS, 36, 2023b.
Sun, Z., Cao, H., Wang, Y., Feng. G., Chen, S., Wang. H., and Chen, W. Adapipe: Optimizing pipeline parallelism with adaptive recomputation and partitioning. In Pro- ceedings of the 29th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 3, pp. 86-100, 2024.
Liu, Y., Duan, H., Zhang, Y., Li, B., Zhang. S., Zhao, W., Yuan, Y., Wang, J., He, C., Liu, Z., et al. Mmbench: Is your multi modal model an all-around player? arXiv preprint ar arXiv:2307.0628/ 2023c.
Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023a.
Masry, A., Do, X. L., Tan, J. Q. Joty, S., and Hoque, E. Chartqa: A benchmark for question answering about charts with visual and logical reasoning. In ACL, pp- 2263-2279, 2022.
Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu,
microsoft.
Megaaron--eeespeed.. https://gitte cttp///icorssottt//aaaaattaeeeee
J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: family of highly capable multimodal models.
2020.
arXiv preprint 2.X.x::312.118880 2023b.
Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Za-
team, Q. Qwen2-vl. 2024.
haria, M. Pipedream: generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM symposium on operating systems principles, pp. 1-15, 2019.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.13977 2023a.
Nous https: https::/htggginnee Nous-Hermest 2-22---11 -2-i-344, 2023.
NousResearch.
hermes
yi-34b.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama Open foundation and fine- tuned chat models. arXiv preprint urXiv:2307.09288 2023b.
OpenAl. ChatGPT. chatgpt/,2022 2023a.
OpenAl. Gpt-4 technical report, 2023b.
Yu, W., Yang, Z., Li, L., Wang. J., Lin, K., Liu, Z., Wang, X., and Wang, L.. Mm-vet: Evaluating large multi- modal models for integrated capabilities. arXiv preprint arXiv:2308.0244 2023.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp.
PMLR, 2021.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Per- formance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Reid, M., Savinov, N., Teplyashin, I Lepikhin, D., Lilli- crap, T., Alayrac, J.-b., Soricut, R., Lazaridou, A., Firat, O., Schrittwieser, J., et al. Gemini 1.5: Unlocking multi- modal understanding across millions of tokens context. arXiv preprint arXiv:2403.055.. 2024.

---

OmniBal: Towards Fast EestructionTuunigg for Visiun-Language Models via Omniverse Computation Balance
A. Training-Pipeline of our balanced method
A.I. Training-Pipeliin
Input
Stage-1
Stage-2
Stage-3
Stage-4
Images
ViT
ViT
LLM
LLM
ViT
ViT
LLM
LLM
ViT
LLM
LLM
LLM
Text1
:
Text2
ViT
LLM
LLM
LLM
Balanced Batch
Balanced Partition Balanced Adaptive Re-Computation
The pipeline consists of four stages, labeled Stage-1 I Stage-4, each representing a different stage of pipeline parallelism. Within this structure, "VIT" stands for the Vision Transformer layer, while "LLM"" refers the Transformer layer used for large language models (LLM). Regarding computational execution, the darker-colored sections signify forward passes re-computation. In contrast, the lighter-colored sections denote standard forward pass without re-computation.
Figure 3.
B. Results on different datasets
We consistently achieved low Dist Ratio across the LLava-665K, InternVL-1.2M, and LCS558K datasets, demonstrated in Table 8. Additionally, our approach significantly enhanced training speed.
Table Results on different datasets
Dist Ratio
Dataset
GPU Days
LLava-665K
0.02
0,145
43.3-> 12.4 (3.5x)
InterVLL- 1.2M
0.02
0.14
61.8- 21.3 (2.9x)
LCS-558K
0.001
0.029
23.8 (3.2x)
C. Results on different Model Size
We test various combinations of vision and language models. As shown in Table 9, our approach significantly reduces the required GPU days for model training, achieving nearly a 2x speedup across models of various sizes.
D. Results on Different Dynamic High- Resoluteso Settings
To validate the effectiveness of our method, test it under various high-resolution settings. Our approach consistently demonstrates low Dist Ratio and strong acceleration across all configurations, as shown in Table 10, significantly improving training speed under different settings.
12

---

OmniBal: Towards Fast Instruction-Tuningg for Vision-Language Models via Omniverse Computation Balance
Table Results for different model sizes shown, with TP, PP, and DP representing various distributed training strategies: Tensor
Parallel (TP), Pipeline Parallel (PP), and Data Parallel (DP), respectively. The Manges-LayerNumm (V+L)" column indicates the number
Vision Transformer (V) Language Transformer (L) layers assigned to each stage. Additionally, "Re-computation" column
denotes the number of re-computations enabled in each stage.
Vision-Model
Language-Model
TP PP DP
Stages-Layer-Num (V+L)
Re-computation
GPU Daysispeed up)
1n-6
Llama3-8B
(1,4,8)
[16.17.20.24]
[8,7,10.24]
27.7 13.8(2.0x)
I1ternVL.-6B
Inteml Inteml.M2-20B M2-20B
(2,4,4)
[22.23,24,24]
[0.0.0.0]
61.8 11.8+2214 21.3(2.9x)
Inter
Yi-34B
(4,4,2)
[28,29,24,24]
[3,2,0,0]
75.4 1144++300.50 -3005(2.55x
1-66
Llama3-70B
(4,8,2)
[22,23.13.14111111111..3
[11.12.8,5332.0000
129 52.55(24x)
InternVI.-68
Qwent 50110B
(8,8,1)
[21,22118]18114,44144444]
[6,91,333.,,00.,0
243 - 7
EVA-CLIP-IB
IntemL M2-20B
(2,4,4)
[43.16,15.14]
[0,0,0.0]
23.6 12.2(1.9x)
EVA-CLIP-4B
IntemL M2-20B
(2,4,4)
[39,22.21.20]
[10.8.1.3]
38.1 17.002.2x)
EVA-CLIP
Inteml M2-20B
(2,4,4)
[17,18,23,2]
[5,5,8,10]
41.8 20.3(2.0x)
EVA-CLIP-18B
IntemL
(4,4,4)
[18,19,25,34]
[2,2,0,0]
63.6 + 33.8(1.9x)
Table 10. Results on different dynamic high-resolution settings. "Max-Patch-Num"" indicates the maximum number of patches into which image divided. This parameter controls the granularity of image segmentation, impacting both model performance and computational efficiency. Adjusting the Max- Patch-Num allows for more flexible handling of high-resolution images in the model, optimizing resource usage while maintaining accuracy.
Max Max-Seq-Len
Dist Ratio
Model
Max-Patch-Num
AVE-BS
GPU Days (speed-up)
I.M
7.6
9K
4K
0.06
28.6 + 13.7 x2.1x)
4K
0.14
61.8- (2.9x)
IntemVL-6B-200
2.7
14K
5K
0.019
0.136
147 -+72 (2.05x)
12
1,9
14K
5K
0.03
0.12
209 ->105 (2.0x)
E. Results on Pretrain Setting
evaluate our method in other tasks, such as pre-training tasks.. Pre-training, we train both the Vision Transformer (ViT) and MLP components for models ranging from 6B to 20B. However, for larger models, such as 6B-34B and 6B- 70B, we focus solely on training the MLP component. Across all configurations, observe consistent performance improvements shown in Table 11, particularly with the largest model, where GPU days significantly reduced from 16.8 to 9.6, demonstrating enhanced training efficiency.
Table 11. Results on Pretrain Setting
Dist Ratio
Model
Dataset
Trainable Module
AVE-BS
GPU Days
VIT
LLM
LntemVL-61-2000
LCS-558K
VIT+MLP MLP MLP
0.03
9.9 - (1.66)))))
IntemVL--6B-3
LCS-558K
5.1
0.031
- (1.69x)
IntemVL-6B--0BB
LCS-558K
5.2
0.029
16.8 9.6 (1.75x)
F. Results on Different Resolutions
We further test our method with different image resolution inputs. As shown in Table 12, our method consistently delivers low Dist Ratio and highly satisfactory acceleration results across varying image resolutions, demonstrating its effectiveness in improving training efficiency.
G. Convergence of ISF.
The convergence performance of ISF is evaluated, with the results illustrated in Figure 4. On the LLava-665K dataset (Liu al., 2023a), we observe that the Dist Ratio for both vision and language data dropped significantly after just one iteration.
13

---

OmniBal: Towards Fast CineructionTunnnng for Vision-Language Models via Omniverse Computation Balance
Table 12. Results Different Resolutions
Dist Ratio
Resolution
AVE-BS
GPU Days
VIT
LLM
224
4.8
0.009
0.068
32.0 - 20 (1.6x)
336
3.3
0.005
0.07
62.0 - 33 (1.88x)
448
4.6
0.02
0.14
61.8 + 21.3 (2.9x)
0.35
Intertall 290 .1
6.4
patttd
PYO 274, (90
0.30
Interwel 2M, (tt
0:35
pattnd 1lm
RTD 22,224, (n. 1t0.
0:20
patzed Im
00.00
0.29
LCSSSEK
putzh12 ratoo
paizhh2 Bm
0.20
0.55
0.25
0.29
0.05
0.05
0.00
10
Iteration step
stRO
step
(a)
(c)
Figure 4. The convergence of ISF in various scenarios, including (a) different datasets, (b) different patch sizes, and (c) different image
resolutions.
After five iterations, the Dist Ratio stabilized considerably. practice, we perform ten iterations to ensure stable results, which only take less than one minute. The computational cost is negligible relative to the overall runtime. Additionally, our method is tested on other datasets, V1-1.2M (Chen et al., 2024) and LCS558K (Liu al., 2023b), and observed similar convergence rates.
H. Results on Open-source LLava-1.6
We also validate our method using another popular open-source model, LLava-1.6, with the DeepSpeed backend, as shown in Table 13. For the DeepSpeed backend, we employ only our balanced dynamic mini-batch strategy. In the case of the open-source Lava model, while its ViT component is relatively small and the imbalance occurs primarily at the data level, we still achieved notable overall speedup. Although the speedup ratio is smaller compared to other models, our method delivered 30% improvement in performance.
Table 13. Results on Open-source LI .Lava-1.6 1.6
Dist Ratio
Model
AVE-BS
GPU Days
VIT
LLM
Llava-1 1.67BB
4.54
0.008
0.037
10.2 - 7.7 (1.3x)
Llava-1.6-13B
4.54
0.008
0.037
18 > 13.3 (1.35x)
4.4
0.009
0.0041
42.7 31.3 (1.36x)
Results on Qwen2-VL Pre-Processing Strategy
Qwen2-VL is recent, highly regarded open-source project that provides strong support for dynamic image input. Conse- quently, we adopt the pre-processing strategy of Qwen2-VL VL to validate our method. As shown in Table 14, our approach demonstrates a substantial acceleration effect (approximately 1.9x) when applied to the Qwen2- VL strategy, significantly reducing both the padding ratio and dist ratio.
14

---

OmniBal: Towards Fast Instruction-Tuningg for Vision-Language Models via Omniverse Computation Balance
Table 14. Results Qwen2-VI. Pre-Prrcesssing strategy
Dist Ratio
Model
Dataset
AVE-BS
Pad-Ratio
GPU Days (speed-up)
VIT
LLM
InterVVL--6B-22
IntemVL.-1.2M
0.31
0.408
0.393
40.2 (1.0x)
IntemVL-6B-20E
IntemVL-1.2M
6.6
0.12
0.06
21.0 (1.9x)
J. Different Hardware Results
We test our method on various hardware platforms with different GPUs (e.g., A100, H100) and network bandwidths. The experiments in Table confirmed consistent performance improvements across all platforms.
Table 15. Results Different Hardware. IB indicates network bandwidths.
Dist Ratio
Dataset
Hardware
IB
GPU Days (speed-up)
VIT
LLM
Intern VL-1.2M
A100
4x200G
0.02
0.145
61.8 - 21.3 (2.90x)
Intern VL-1.2M
A100 H100
2x200G
0.02
0.145
64.0 24.8 (2.58x)
Intern VL-1.2M
8x400G
0.02
0.145
32.5 - 12.2 (2.67))
K. Large-Scale Results
To validate the effectiveness of our method, we conduct study using larger-scale models and greater number of GPUs. As shown in the Table 16, our method achieves speedup ratio exceeding 2.0 across varying GPU configurations. Moreover. the results demonstrate that approach maintains more favorable linear speedup (85% - 95%) GPUs increase.
Table 16. Results on Large-Scale models (6 + 70B) and GPUs
Dist Ratio
Dataset
Hardware
GPUs
GPU Days (speed-up)
VIT
LLM
InternVL-1.2M
H100
8x400G
64
0.02
0.139
72.8 29.3 (2.48
Vnter2VL-1.2M
H100
8x400G
128
0.02
0.139
75.2 29.7 (2.53x)
InternVL--1.2M
H100
8x400G
256
0.02
0.139
82.1 + 30.4 (2.70x)
Intern VL-1.2M
H100
8x400G
512
0.02
0.139
85.3 30.9 (2.76x)
15


Reviews:
[{"review_id": "kmOxGQIq2i", "reviewer": "Reviewer_YLFx", "summary": "This paper focuses on the large-scale distributed training of multimodal large language models, and propose a omniverse computation strategy to manage vision-language data distribution and the training memory optimization. \n\nAlthough the studied problem is significant in the development of multimodal large language models, the paper content does not really match the scope of ICML. For instance, the mainly compared work [Rajbhandari et al. 2020] is published at ICHPCNSA. \n\nMoreover, the arguments made in this paper are often very subjectively. For instance, the authors claim that the data distribution affects distributed training efficiency. This statement has no any reference to support, and the experimental results are hard to make a direct connection with this point. Similarly, in the most part of this paper, the description lacks enough references to support, e.g., Sec.3. In the related work, there are also not other works mentioned about the distributed traing of VL models. \n\nBesides, this paper seems to be finished in a harry. The presentation is very poor, making hard to really understand the motivation, methodology and contribution of this paper, although the experimental results seem significant.  The overall writing of this paper is more like a technical report, rather than an academic paper published by ICML.  \n\nOverall, I would like to suggest the authors to spend more time in polishing this paper, and then final a more relevant conference or journal for the submission.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback.\n\n*Q1: ICML scope concerns*\n\n**A1:** We believe our paper is within the scope of ICML. \n\n**According to the ICML 2025 Call for Papers, topics of interest include (but are not limited to):\n\"Machine Learning Systems (improved implementation and scalability, hardware, libraries, distributed methods, etc.).\"**\n\nOur work focuses on large-scale distributed training and memory optimization for multimodal large language models, which fit the ICML scope.\n\n*Q2: in the most part of this paper, the description lacks enough references to support, e.g., Sec.3. In the related work, there are also not other works mentioned about the distributed traing of VL models*\n\n**A2:** \n\nIn Section 2.1 Multi-Modal Large Language Models (MLLMs), we have cited a number of prior works on vision-language model (VLM) training, including **Qwen-VL, Q-Former, and LLaVA**. These papers, while primarily focused on model design and applications, do include brief discussions of distributed training strategies, often relying on backends such as DeepSpeed.\n\nMoreover, throughout the paper, we have provided extensive references to support our claims regarding distributed training and memory optimization, including established systems such as **PipeDream, DeepSpeed ZeRO, Megatron-LM, and GPipe**, among others.\n\n*Q3: Besides, this paper seems to be finished in a harry. The presentation is very poor, making hard to really understand the motivation, methodology and contribution of this paper, although the experimental results seem significant. \"*\n\n**A3:** \n\n While the current draft may need polishing in parts, the work was not done in haste. Other reviewers did not raise similar concerns about clarity.  \n- *Reviewer rnAP comment \"Well written paper with lots of analysis\"*\n- *Reviewer vzDC comment \"The motivation/observation of the imbalanced computation in VLMs is clear.\"*\n\nFeel free for any other comments.", "claims_and_evidence": "See summary.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "See summary.", "essential_references_not_discussed": "Yes", "other_strengths_and_weaknesses": "See summary.", "other_comments_or_suggestions": "See summary.", "questions_for_authors": "See summary.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sLfen6W1eT", "reviewer": "Reviewer_NPbH", "summary": "The paper addresses the causes of computational imbalance in VLM training, including aspects of data, model, and memory, and introduces OmniBal, a training framework designed for improving training efficiency of VLMs. OmniBal is basically comprised of three algorithms, balancing batch sizes, model partitions, and memory usages, respectively. The authors present experiments on various datasets and VLMs, claiming that their framework accelerates the process of VLM training under the metric of GPU days", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.  Figures and tables are shown at https://anonymous.4open.science/r/O-A/O.pdf.\n\n*Q1:  In section 3 the paper mentions the differences between the training of VLMs and LLMs. However, it needs more solid explanations and explicit data to tell: (i) Whether present strategies experimented on LLMs can be transferred to VLMs? If they can, the data of their performance is expected. (ii) As mentioned in the same subsection, \"(simple packing) results in computation imbalance problems (on VLM instruct-tuning)\", then how severe the problems are?*\n\n**A1:** \n\n**Whether present strategies experimented on LLMs can be transferred to VLMs?**\n\nThe **simple packing** strategy used in LLM training can be directly applied to VLMs, its effectiveness is limited due to the structural differences between the two. Figure 1 (available at https://anonymous.4open.science/r/O-A/O.pdf) provides an example illustrating how LLM-style packing manifests in VLM training.\n\n**(simple packing) results in computation imbalance problems (on VLM instruct-tuning) how severe the problems are?**\n\nThe severity of the imbalance is demonstrated through actual analysis and experimental results, as shown in the accompanying table. \n\n|   model   | backend  | Data method | Dist Ratio VIT | Dist Ratio LLM | GPU Days |\n|:--------:|:--------:|:-----------:|:------------:|:------------:|:--------:|\n| 6 + 20B  | Megatron |   random    |     0.34     |     0.30     |   61.8 |\n| 6 + 20B  | Megatron | LLM packing |     0.40     |     0.05     |   48.3    |\n| 6 + 20B  | Megatron |     ISF     |     0.02     |     0.14     |   **21.3**   |\n\n\n*Q2:  Effective as the experiments show, yet the solutions provided in the paper seem trivial. In addition, the details of optimizing re-computation in the Balanced Adaptive Re-Computation Strategy needs further explanation in Appendix B.1.*\n\n**A2:**  Our problem is both  **challenging and important**, and to our knowledge, it has not been systematically studied before. Although our proposed solution is relatively simple, it is highly efficient and easy to transfer and apply.\n\nRegarding the optimization of re-computation in the Balanced Adaptive Re-Computation Strategy, we provided a brief explanation in Appendix B.1, as the underlying idea is relatively straightforward. We appreciate the reviewer’s suggestion and will include a more detailed discussion in a future version of the paper.\n\n*Q3: In Table 4, the proposed method BMP only shows marginal improvement, compared with present methods.*\n\n**A3:** BMP will yield more substantial benefits in more **communication-constrained** settings (larger models). BMP, compared to previous methods, takes into account the imbalance in point-to-point communication across different pipeline stages in VLMs.\n\n*Q4: the experiments are not sufficient. (i) The paper demonstrates the effectiveness of OmniBal, but lacks comparisons between OmniBal as a holistic framework and other works on this issue. (ii) More ablation studies of different combinations of components are expected in Tabel 6, data + memory balance and model + memory balance for instance. Although these modules might perform better with prior ones applied, such experiments are still recommended.*\n\n**A4:**\n\n**Comparison with other works:**\n\nDue to the limited number of studies specifically focused on large-scale VLM training, most existing approaches are designed for pure LLMs and cannot be directly applied.\n\nModel balance optimization in AdaPipe (combined with our data balance method ISF) is aligned with the profile-based method in Table 4 (paper) which achieves inferior results compared to our BMP.\n\n**Additional ablation studies:**\n\nThe components in OmniBal are inherently interdependent and cannot be decoupled trivially. \n\n**Data-level balancing must be addressed first**, as it lays the foundation for subsequent optimizations at the model and memory levels. \n\nMoreover, memory optimization is intrinsically tied to model structure. As such, we did not perform isolated ablation studies. Instead, we adopted a progressive evaluation strategy, incrementally adding modules to assess their cumulative effectiveness.\n\n**Response to Other Comments Or Suggestions:**\n\nWe sincerely appreciate the reviewer’s valuable feedback and will incorporate the suggested improvements in the revised version of the paper.\n\nFeel free for any other comments.", "claims_and_evidence": "Please see  Strengths and Weaknesses.", "methods_and_evaluation_criteria": "Please see  Strengths And Weaknesses.", "theoretical_claims": "Please see  Strengths and Weaknesses.", "experimental_designs_or_analyses": "Please see  Strengths and Weaknesses.", "relation_to_broader_scientific_literature": "None.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n- This work demonstrates the detrimental impact of imbalanced data (batch size), distribution of model computation load, and re-computation on training efficiency of VLMs explicitly, showing thorough comprehension on the issue.\n- Experiments across various datasets and VLM models verified the effectiveness of the OmniBal as a holistic framework, suggesting that it boosts training speeds by around 1.5 times to 3.5 times.\n\nWeaknesses:\n- In section 3 the paper mentions the differences between the training of VLMs and LLMs. However, it needs more solid explanations and explicit data to tell: (i) Whether present strategies experimented on LLMs can be transferred to VLMs? If they can, the data of their performance is expected. (ii) As mentioned in the same subsection, \"(simple packing) results in computation imbalance problems (on VLM instruct-tuning)\", then how severe the problems are?\n- Effective as the experiments show, yet the solutions provided in the paper seem trivial. In addition, the details of optimizing re-computation in the Balanced Adaptive Re-Computation Strategy needs further explanation in Appendix B.1..\n- In Table 4, the proposed method BMP only shows marginal improvement, compared with present methods.\n- The experiments are not sufficient. (i) The paper demonstrates the effectiveness of OmniBal, but lacks comparisons between OmniBal as a holistic framework and other works on this issue. (ii) More ablation studies of different combinations of components are expected in Tabel 6, data + memory balance and model + memory balance for instance. Although these modules might perform better with prior ones applied, such experiments are still recommended.", "other_comments_or_suggestions": "- There are some expressions and definitions that can be polished to be clearer for readers. In subsection 4.1, the exact term \"Distribution ratio\" should be placed right after the \"Dist ratio\" in bold. In subsection 4.2, when introducing the partitioning strategy, it might be better to use the form of the definition of partition in set theory. Both might make it clearer and avoid confusion.\n- The explanation of Balanced Adaptive Re-Computation Strategy in Appendix B.2. can be moved to the main body to make the paper more organized.", "questions_for_authors": "See Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DN0aCs1Lmk", "reviewer": "Reviewer_vzDC", "summary": "This paper identifies a significant computational imbalance issue in large-scale distributed training for Vision-Language Models (VLMs) due to heterogeneity in vision and language components. To tackle this, the authors propose OmniBal, a comprehensive framework that balances computation across three dimensions—data, model, and memory—by adaptively forming balanced mini-batches, optimizing model partitioning, and strategically adjusting memory re-computation. Extensive experiments demonstrate that OmniBal significantly reduces training time while maintaining accuracy, confirming its efficacy and versatility across various models, datasets, and hardware configurations.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.  Figures and tables are shown at https://anonymous.4open.science/r/O-A/O.pdf.\n\n*Q1: what are the major challenges that prevent applying those approaches to VLMs?*:\n\n**Major Challenges:**\n\n- **Data Level:**\nSimple packing strategies for LLMs lead to a severe imbalance in the vision components when applied to VLMs shown in Figure 1 (https://anonymous.4open.science/r/O-A/O.pdf). \n- **Model Level:**\nExisting pipeline parallelism techniques assume homogeneous architectures for LLM and consistent input lengths. In VLMs, however, the ViT and LLM components differ significantly in both structure and input sequence length, leading to computational imbalance when such methods are naively applied.\n\n- **Memory Level:**\nThe architectural mismatch between ViT and LLM also results in uneven memory consumption, rendering LLM-specific memory optimization strategies ineffective or even infeasible in VLM scenarios.\n\n**Illustrative Example:**\nFigure 1 illustrates a typical data imbalance issue that arises when applying simple LLM-style packing to VLM training.\n\n*Q2:  Section 4 introduces Q'v and Q't to reduce Pad Ratio and Dist Ratio, ... how Q values are chosen to achieve the intended balance.* \n\n**A2:**\n\n**Why not a mathematical formula?**\n\nThe core of ensuring data balance lies in simultaneously fixing both the ViT input and the LLM input, which constitutes a Two-Dimensional Balanced Partition Problem (2D-BPP), an **NP-hard problem**, making it difficult to derive a general mathematical formulation. \n\nOur goal is to approximate equal input lengths for ViT and LLM, but no explicit optimal solution exists. Thus, we introduce Q'v and Q't as dataset-dependent hyperparameters to guide the balancing process.\n\n**How to get Q'v and Q't?**\n\nWe provide a script (https://anonymous.4open.science/r/omnibal_example-E4D7/test_balanced_dynamic_batch.py,  line 227) that uses an offline exhaustive search (≈10 minutes for 1.2M dataset only once, **no full training required**) to automatically determine Q'v and Q't, making the method easy to apply in practice.\n\n*Q3: In Section 5, Q′v is set equal to Qv, while Q′t is defined as Qt − 128. Could the authors provide an intuition or rationale behind these choices?*\n\n**A3:** The choices of Q′v and Q′t are based on the search results discussed in **A2: How to get Q′v and Q′t?**. The goal is to minimize the DistRatio, and Appendix C provides an ablation table showcasing part of our search results supporting this decision.\n\n*Q4: The design of data balancing seems a little bit ad-hoc. It seems like a trial-and-error method that relies heavily on some empirically picked hyper-parameters. Consequently, it's unclear whether this method would generalize effectively to different evaluation datasets.*\n\n**A4:** \n\n**This is not ad-hoc**\n\nThe data balancing strategy was carefully designed based on a thorough analysis of the data distribution and task requirements. \n\n**hyper-parameters**\n\nThe parameters were also designed to solve the Two-Dimensional Balanced Partition Problem (2D-BPP), an **NP-hard problem**, and they are based on a search-based approach using dataset information, making them easily generalizable to other validation datasets. \n\nwe conducted extensive experiments across diverse datasets, as detailed in Section 5.4 (**Generalization Capability**). These results demonstrate that our approach generalizes well beyond the initial evaluation setting.\n\n\n*Q5: In the model partitioning design, it says that \"A candidate set of partition strategies is created by jittering P(1), P(2), . . . , P(N−1) within a radius of r\". Could the authors clarify how these parameters are chosen and whether the selection process is systematic or largely empirical?*\n\n**A5:** \n\n**How these parameters are chosen**\n\n\"P(1), P(2), ..., P(N−1) init\" are obtained by profiling each layer’s forward time. A greedy algorithm computes the anchor partition strategy P⁺ to balance the computation time across all stages Si, which is **systematic**.\n\n\"r\" is **empirically** determined and spans the adjacent layers between the ViT and LLM to ensure a sufficient candidate space.\n\n*Q6: I would like to see the comparison of the proposed method to some other... Adapipe, besides showing the speedup over the baselines.*\n\n**A6:**  \n\n**Comparison with other works:**\n\nDue to the limited number of studies specifically focused on large-scale VLM training, most existing approaches are designed for pure LLMs and cannot be directly applied.\n\nModel balance optimization in AdaPipe (combined with our data balance method ISF) is aligned with the profile-based method in Table 4 (paper) which achieves inferior results compared to our BMP.\n\n*Q7: What is the difference between Figures 2 and 3?*\n\n**A7:**  Same content, shown again for the reviewer to easily reference.\n\nFeel free for any other comments.", "claims_and_evidence": "Please refer to the Question.", "methods_and_evaluation_criteria": "Please refer to the Question.", "theoretical_claims": "N.A.", "experimental_designs_or_analyses": "Please refer to the Question.", "relation_to_broader_scientific_literature": "Please refer to the Question.", "essential_references_not_discussed": "Please refer to the Question.", "other_strengths_and_weaknesses": "Strength:\n\n1. The paper tackles an important challenge of computational imbalance in large-scale distributed training of vision-language models (VLMs).\n2. It introduces a comprehensive framework that systematically addresses this imbalance from three complementary angles: data input balancing, optimized model partitioning, and adaptive memory management.\n3. The motivation/observation of the imbalanced computation in VLMs is clear.\n\nWeakness:\n1. Some claims need a more detailed explanation.\n2. The novelty of this paper needs further justification.\n3. Some design choices seem ad-hoc and need justification.\n\nPlease refer to the Questions section for details.", "other_comments_or_suggestions": "Please refer to the Question.", "questions_for_authors": "Thanks for the submission to ICML. I have the following comments/questions on this submission:\n\n&nbsp;\n1. The data partitioning approach in this paper appears to offer limited novelty. Numerous adaptive data batching methods, such as [1] and [2], have already addressed similar issues, albeit without specific evaluations on VLMs. However, applying these techniques to VLMs does not seem particularly challenging, given that (i) the primary goal—balancing computation both within and across batches—is quite similar, and (ii) the methodology, which involves batching based on profiled or predicted latency, closely resembles the existing works. Could the authors clarify the unique contributions of the proposed data partitioning method in comparison to these established batching approaches? \n\nIn otehr words, what are the major challenges that prevent applying those approaches to VLMs? If it is a simple adaptation, I have concerns on the contribution of this paper. \n\n2. Section 4 introduces Q'v and Q't to reduce Pad Ratio and Dist Ratio, but the paper lacks a clear formula or explanation of how these values affect the ratios. This makes it challenging to understand how Q values are chosen to achieve the intended balance.\n3. In Section 5, Q′v is set equal to Qv, while Q′t is defined as Qt − 128. Could the authors provide an intuition or rationale behind these choices?\n4. The design of data balancing seems a little bit ad-hoc. It seems like a trial-and-error method that relies heavily on some empirically picked hyper-parameters. Consequently, it's unclear whether this method would generalize effectively to different evaluation datasets.\n\n5. In the model partitioning design, it says that \"A candidate set of partition strategies is created by jittering P(1), P(2), . . . , P(N−1) within a radius of r\". Could the authors clarify how these parameters are chosen and whether the selection process is systematic or largely empirical?\n\n6. I would like to see the comparison of the proposed method to some other computation balanacing frameworks such as Adapipe, besides showing the speedup over the baselines.\n\n7. What is the difference between Figures 2 and 3?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zaClF8DuXO", "reviewer": "Reviewer_NXGX", "summary": "This work focuses on addressing imbalanced computational loads in large-scale 3D parallel training of vision-language models by rebalancing across data, model, and memory dimensions. The authors conduct experiments on various models, datasets, and hardware platforms to demonstrate the speed-up ratio for vision-language model training.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.  Figures and tables are shown at https://anonymous.4open.science/r/O-A/O.pdf.\n\n*Q1: The claim that \"vision-language instruct-tuning models have recently made significant progress due to their more comprehensive understanding of the world\" is unclear. The statement implies that the progress is caused by a deeper understanding of the world, yet the authors provide no evidence or explanation for how such understanding is achieved or how it leads to the observed improvements.*\n\n**A1:**  We acknowledge that the original explanation may lack clarity. In the revised version, we will provide a more rigorous and precise description, clarifying the mechanisms behind the improvements in vision-language instruct-tuning models and avoiding ambiguous references to \"understanding of the world.\n\n*Q2: Regarding the heterogeneity between LLM and ViT models, while this disparity is an objective fact, when the two models are fine-tuned as a unified system, it is unclear why one should consider the structure of each module separately. The authors fail to explain in detail what consequences arise from this heterogeneity—for example, how imbalanced model partitioning is affected and how the partitioning is specifically conducted.*\n\n**A2:** \n\n**Reasons for considering the structure of each module separately**\n\n - **Input differences:** ViT and LLM modules receive fundamentally different types of inputs with distinct length distributions.\n - **Architectural differences:** Although both ViT and LLM are based on Transformer architectures, their configurations differ significantly, resulting in unequal computational overhead.\n\n**Consequences arising from this heterogeneity**\n - **Data-Level:** As shown in Figure 1 (https://anonymous.4open.science/r/O-A/O.pdf), when using LLM-style simple packing in a vision-language system, the ViT component suffers from a data imbalance problem.\n - **Model-Level:** Figure 5 (https://anonymous.4open.science/r/O-A/O.pdf) shows the simple example model partition based on LLM's method. Since computational cost depends on both the architecture and input length, the architectural and input discrepancies between ViT and LLM make it difficult to evenly divide computation across pipeline stages. This results in an imbalanced workload and significant pipeline bubbles when applying traditional pipeline parallelism strategies designed for LLMs.\n\n*Q3. The paper focuses on the instruction-tuning stage, but pretraining is equally important for building a strong multimodal LLM (MLLM). The authors do not mention pretraining or explain why it is not addressed in this work.*\n\n**A3:**   We have also conducted experiments on multimodal large model pretraining. Details can be found in Section 5.4(4), Generalization Capability, under \"Different Tasks,\" where we state: \"Besides SFT tasks, pretraining tasks are also tested, as shown in Appendix G, and we observed consistent improvements across all settings.\" These results demonstrate that our approach remains effective even in the pretraining stage.\n\n*Q4: Why does using data balancing affect the maximum sequence length, and how is this reduction achieved?*\n\n**A4:** \n\nFigure 2 (https://anonymous.4open.science/r/O-A/O.pdf) shows details examples of how this reduction is achieved.\n\n**Simple Fixed Batching**\n\nThe \"maximum length\" here refers to the max-seq-len in Table 5, including the batch dimension. The maximum input length is 5K tokens for the ViT and 4K tokens for the LLM. And with a batch size of 4 (as in InternVL-1.5), this yields 20K for ViT and 16K for LLM.  \n\n**ISF Dynamic Batching**\n\n Since ISF adopts a dynamic batching strategy, it can flexibly manage sequence lengths.  If some samples are relatively long, the corresponding batch size will be set smaller to keep the maximum sequence length (VIT 9K and LLM 4K).\n\n*Q5: The authors mention that existing vision-language models suffer from issues such as structural heterogeneity and varied input sizes, which lead to slower training speeds. However, if future MLLMs adopt an architecture similar to the \"fuyu\" model—directly inputting visual tokens into the LLM without a separate visual encoder—would these issues persist? Would the proposed method still be effective under such circumstances?*\n\n**A5:**  \n\n\"fuyu\" model—directly inputting visual tokens into the LLM without a separate visual encoder making the model behave like LLM. The structural heterogeneity and input size imbalance issues will **disappear**. \n\nHowever,  its performance still lags behind more established architectures such as ViT-MLP-LLM. Therefore, our proposed method remains highly relevant to the current mainstream vision-language models.\n\n*Q6: The font in Figures 3 and 4 is hard to read and should be improved for clarity.*\n\n**A6:** We sincerely appreciate the reviewer’s valuable feedback and will update it in the revised version of the paper.\n\nFeel free for any other comments.", "claims_and_evidence": "1. The claim that \"vision-language instruct-tuning models have recently made significant progress due to their more comprehensive understanding of the world\" is unclear. The statement implies that the progress is caused by a deeper understanding of the world, yet the authors provide no evidence or explanation for how such understanding is achieved or how it leads to the observed improvements.\n\n2. Regarding the heterogeneity between LLM and ViT models, while this disparity is an objective fact, when the two models are fine-tuned as a unified system, it is unclear why one should consider the structure of each module separately. The authors fail to explain in detail what consequences arise from this heterogeneity—for example, how imbalanced model partitioning is affected and how the partitioning is specifically conducted.\n\n3. The paper focuses on the instruction-tuning stage, but pretraining is equally important for building a strong multimodal LLM (MLLM). The authors do not mention pretraining or explain why it is not addressed in this work.", "methods_and_evaluation_criteria": "1. The authors propose to enhance training speed from three aspects: balanced dynamic mini-batches, balanced model partition, and balanced adaptive re-computation.\n2. They validate the efficiency improvements of their method across different backends and balancing strategies.", "theoretical_claims": "No significant theoretical claims are made.", "experimental_designs_or_analyses": "1. The experiments first verify, across different backends and model sizes, that the proposed method not only speeds up training but also maintains training effectiveness.\n2. Subsequently, the authors separately evaluate improvements from the data, model, and memory perspectives.", "relation_to_broader_scientific_literature": "1. The authors propose an iterative sampling and filtering strategy to improve data balance.\n2. They utilize the sum of VAR (forward time) and communication time as a metric for partition evaluation, ranking the top K candidates for speed assessment.\n3. Finally, the re-computation strategy is optimized based on actual memory needs.\n4. The combined improvements across these three aspects lead to a significant increase in training speed, which is beneficial for the faster construction of MLLMs.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "The font in Figures 3 and 4 is hard to read and should be improved for clarity.\n\nIf the authors address my concerns, I will raise the score.", "questions_for_authors": "1. Why does using data balancing affect the maximum sequence length, and how is this reduction achieved?\n2. The authors mention that existing vision-language models suffer from issues such as structural heterogeneity and varied input sizes, which lead to slower training speeds. However, if future MLLMs adopt an architecture similar to the \"fuyu\" model—directly inputting visual tokens into the LLM without a separate visual encoder—would these issues persist? Would the proposed method still be effective under such circumstances?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xH0Xvtruos", "reviewer": "Reviewer_rnAP", "summary": "## update after rebuttal: \nScore updated to 3\n\nThis paper suggest that during distributed training of vision-language models, there are computation imbalance due to model architectures, data types and how the mini batches are constructed inside and across devices. This paper proposes a novel method of mini-batch construction, model partitioning and re-computation of memory workload distribution so the computation imbalance is minimized, achieving a 1.8x faster training efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.  Figures and tables are shown at https://anonymous.4open.science/r/O-A/O.pdf.\n\n*Q1:  Not sure I can distinguish between evidence 1 (line 041) and evidence 3 (line 047):*\n\n   \"(1) Varying input sizes of LLM and VIT cause imbalance computational loads across training iterations and devices.\n\n   (3) Input size variation and computational imbalance compel us to use the most aggressive re-computation (checkpointing)  strategy to prevent program crashes, which wastes computational resources.\"\n\n**R1:**  \n- (1) highlights the **idle time** across different devices caused by the variation in sequence lengths between the LLM and ViT, which leads to load imbalance during training.\n- (3) emphasizes the computational overhead introduced by the need for aggressive checkpointing strategies leading to **wasted computational resources**. \n\n\n*Q2： I think the methods and evaluation criteria are more empirical and not mathematical. The authors did a lots of experiments from different angles and they are easy to understand. But I do not see any mathematical formulation of their methods. Therefore it is harder to apply the techniques in other models. For example, the calculation of Qv and Qt is purely experimental. So if I want to find out those values for a new model, without any math formulation, I am left with running a bunch of experimentation.*\n\nConcerns about the mathematical formulation of methods and how to find Qv and Qt.\n\n**A2:**\n \n**Why not a mathematical formula?**\n\nThe core of ensuring data balance lies in simultaneously fixing both the ViT input and the LLM input, which constitutes a Two-Dimensional Balanced Partition Problem (2D-BPP), an **NP-hard problem**, making it difficult to derive a general mathematical formulation. Instead, we adopt a practical and intuitive approach using experiments from multiple angles to tackle the problem effectively.\n\nThe imbalance arises from variable input lengths. Our goal is to approximate equal input lengths for ViT and LLM, but no explicit optimal solution exists. Thus, we introduce Q'v and Q't as dataset-dependent hyperparameters to guide the balancing process.\n\n**How to get Q'v and Q't?**\n\nWe provide a script (https://anonymous.4open.science/r/omnibal_example-E4D7/test_balanced_dynamic_batch.py, line 227) that uses an offline exhaustive search (≈10 minutes for 1.2M dataset only once, **no full training required**) to automatically determine Q'v and Q't, making the method easy to apply in practice.\n\n*Q3: How much overhead does the initial sampling stage adds?*\n\n**A3:**\n\n**Overhead Analysis**\n\nThe sampling overhead is very low for example, with 1.2 million samples, the sampling can be completed within a **few tens of seconds**, which imposes minimal overhead on the overall training process. \n\nThe time **complexity of ISF** is C*O(N+M), where N is the number of samples, M is the number of samples per pack and C is the number of iterations. \n\n*Q4:  How does it scales when you have more than 2 modalities - say language, audio and image? How will this impact the initial sampling overhead? Increase significantly or minimal?*\n\n**A4:**\n\n**Scaling to More Than Two Modalities**\n\nWe begin by categorizing multi-modal tasks using the case of three modalities: language, audio, and image. When a task involves only two modalities (e.g., **ViT + LLM, LLM + Audio, or Audio + ViT**), our method can be directly applied with minimal modifications. These cases resemble standard VLM scenarios and are illustrated in the upper part of Figure 3  (https://anonymous.4open.science/r/O-A/O.pdf).\n\nFor tasks involving all three modalities (**ViT + LLM + Audio**), we also provide adapted examples in the lower part of Figure 3 (https://anonymous.4open.science/r/O-A/O.pdf). In these cases, the underlying principles of our approach remain consistent. The only change is the need to handle more combinations during batch construction.\n\n**How will this impact the initial sampling overhead?**\n\n Increase minimal for our ISF.  The time complexity of ISF is C*O(N+M),  which is essentially an **O(N)-level** method and therefore does not introduce significant overhead.\n\n*Q5: I understand the inter-strage and intra-stage problem as described in figure 2. It would be good if you could show how this figure changes once you apply OnmiBal*\n \nFigure 4 (https://anonymous.4open.science/r/O-A/O.pdf) shows the results of using Omnibal.\n\nFeel free for any other comments.", "claims_and_evidence": "Not sure I can distinguish between evidence 1 (line 041) and evidence 3 (line 047)", "methods_and_evaluation_criteria": "I think the methods and evaluation criteria are more empirical and not mathematical. The authors did a lots of experiments from different angles and they are easy to understand. But I do not see any mathematical formulation of their methods. Therefore it is harder to apply the techniques in other models. For example, the calculation of Qv and Qt is purely experimental. So if I want to find out those values for a new model, without any math formulation, I am left with running a bunch of experimentation", "theoretical_claims": "Table 7 shows ablation of Qv, Qt. But what is the mathematical formulation? Asking this because if I want to apply this technique to a new model, how do I find Qv and Qt for that?", "experimental_designs_or_analyses": "Experiments are well designed as per the table 3 and 4", "relation_to_broader_scientific_literature": "this paper in on the broader topic of efficient distributed training such as LazyBatching (https://arxiv.org/abs/2010.13103) or DeepSpeed Zero etc. But the focus of this paper is on the vision-language model which adds a different flavor to existing works that are generally focused on single modality (specifically LLMs)", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strength:\n1. Well written paper with lots of analysis \n2. The authors are tackling a novel problem of looking into efficient distributed training of vision-language models\n\nWeakness\n1. The paper lacks mathematical formulation so hard to use it in other models beyond the models shown in the experiment. For example if I want to do this for PaliGemma, how do I arrive at Qv and Qt without running a lot of experiments?", "other_comments_or_suggestions": "None", "questions_for_authors": "1. How much overhead does the initial sampling stage adds?\n2. How does it scales when you have more than 2 modalities - say language, audio and image? How will this impact the initial sampling overhead? Increase significantly or minimal?\n3. I understand the inter-strage and intra-stage problem as described in figure 2. It would be good if you could show how this figure changes once you apply OnmiBal", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a large-scale 3D parallel training on the VLMs to solve the imbalanced computation load across different devices. The authors rebalance the computational load from data, model and memory perspectives, which achieves more balanced computation across devices. This paper received mixed reviews initially. The reviewers raised several issues including the reasonability of arguments, unclear paper presentation, unclear technique details, and insufficient experimental validations. During the rebuttal phase, the authors try to address these issues by providing more illustrations and results. These efforts have made reviewers unanimously lean towards positive afterward. Overall, the AC has checked all files, and stands on the reviewers' side. The authors shall incorporate relevant clarifications in the camera-ready version.
Decision: Accept (poster)