OMNIBAL: TOWARDS FAST INSTRUCT-TUNING FOR VISION-LANGUAGE MODELS VIA OMNIVERSE COMPUTATION BALANCE

Authors: Yongqiang Yao, Jingru Tan, Feizhao Zhang, Jiahao Hu, JinXin, Bo Li, Yazhe Niu, Ruihao Gong, Pengfei Liu, Dahua Lin, Ningyi Xu
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: Fast Vision-Language Training; 3D parallel

Abstract:
Vision-language instruct-tuning models have recently made significant progress due to their more comprehensive understanding of the world. 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.8x speed-up. Our method's efficacy and generalizability are further validated across various models and datasets. Codes will be released at https://github.com/anonymousiclr293/omnibal_example.

Decision: reject
Reviewer scores: 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1730285439469

thread_type:
official_review

<thread_object>
creator: reviewer_fc4x

<content>
summary:
This paper aims at addressing the issue of imbalanced computation loads in large-scale 3D parallel training of vision-language models. It rebalances across data, model, and memory dimensions. Experimental results demonstrate that this method can effectively speed up training on many open-source models.

strengths:
Strength：

1. Extensive experiments demonstrate the approach's effectiveness for rebalancing computation for MLLMs.
2. The motivation is clear and the in-depth analysis and ablations are convincing.
3. This paper is well-organized and clearly-written.

weaknesses:
N/A

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721090880

thread_type:
official_review

<thread_object>
creator: reviewer_r51z

<content>
summary:
This paper focuses on the 3D parallel training of MLLMs. The approach to balancing the computational load concerning data, model, and memory distribution both in-device and cross-device is straightforward yet effective. The proposed method demonstrates significant time reduction while maintaining the model's capabilities.

strengths:
+ The paper centers on the efficient parallel training of MLLMs particularly focusing on balancing the computational load concerning data, model, and memory distribution both in-device and cross-device. 

+ The multi-modal sample balancing strategy stands out for its simplicity, offering a straightforward approach to load management. Despite its simplicity, the method proves to be effective, yielding significant reductions in training GPU hours while preserving the model’s capabilities. TItmakes the approach both accessible for large-scale applications.

+ It is comprehensively validated across different datasets, tasks, resolutions and hardware archs etc

weaknesses:
- Results on other MLLM architectures, especially those commonly used in the community such as the LLaVA series, are expected in a technical report style work.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730742599524

thread_type:
official_review

<thread_object>
creator: reviewer_dcak

<content>
summary:
This paper introduces OmniBal, a training framework that addresses computational imbalances in large-scale Vision-Language Models (VLMs) during 3D parallel distributed training. OmniBal improves efficiency by balancing data, model, and memory across devices. For data, the framework forms balanced mini-batches that reduce padding and minimize computational idle times. For the model, a search-based method is used to optimally partition vision and language components, addressing their inherent structural differences. For memory, adaptive re-computation maximizes memory utilization, enhancing training speed. Experiments show OmniBal reduces training time significantly.

strengths:
1. This paper addresses a critical issue: accelerating distributed training for large-scale vision-language models (VLMs).

2. The authors present a comprehensive optimization framework that targets VLM distributed training from three perspectives: data, model, and memory balance.

weaknesses:
1. The paper lacks sufficient motivation data to illustrate the imbalance problem quantitatively.

2. The contributions of this paper seems incremental.

3. Some figures have fonts that are too small, making them challenging to read and interpret.

Please see my comments for details.

questions:
This paper propose OmniBal, which improves efficiency by balancing data, model, and memory across devices. It achieves significant improvements over baselines. However, I have following questions/concerns:

1. 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 their data partitioning method in comparison to these established batching approaches?

2.  The novelty of the model partitioning and re-computation strategies also seems limited. The AdaPipe framework, as referenced in the paper, performs layer-wise partitioning and adaptive re-computation based on computation and memory costs. Since the information needed for partitioning VLMs can also be easily profiled, could the authors explain the specific advantages of their approach over extending AdaPipe for VLMs?

3. The motivation for addressing data imbalance would be stronger with quantitative results from real datasets illustrating variation in sample sizes or computational demands. The simplified example in Figure 2 does not adequately demonstrate the scope of the issue.

4. The section on memory imbalance in Section 3 lacks clarity. While it states the need for aggressive re-computation, it would benefit from a discussion on potential opportunities to mitigate memory imbalance.

5. The figure in Appendix A.1.2 does not clearly illustrate how computation is balanced, making it difficult to assess the efficacy of the approach.

6. 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.

7. 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?

8. The partitioning approach begins with an anchor partition strategy. What is the computational overhead for generating this anchor partition, and how close is it to the optimal partition? If it is far from optimal, what steps are taken to improve upon it? Conversely, if it is close, is it necessary to generate multiple new candidate partitions?

minor point(s):
1. The fonts in some figures are too small, particularly in Figure 3, where the axis labels and legends are difficult to read.


[1] Yu, Gyeong-In, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. "Orca: A distributed serving system for {Transformer-Based} generative models." In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pp. 521-538. 2022.

[2] Choi, Yujeong, Yunseong Kim, and Minsoo Rhu. "Lazy batching: An sla-aware batching system for cloud machine learning inference." In 2021 IEEE International Symposium on High-Performance Computer Architecture (HPCA), pp. 493-506. IEEE, 2021.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733144099771

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all Reviewers and Area Chairs

comment:
Dear Reviewers and Area Chairs

We sincerely thank the reviewers and the area chair for their insightful feedback and the time they dedicated to evaluating our work. We deeply appreciate their recognition of the strengths of our research. The reviewers collectively agree that:

**The Framework:**

- **Important  Problem**  "This paper addresses a critical issue: accelerating distributed training for large-scale vision-language models (VLMs)." (**Reviewer dcak**)
- **Clear Motivation** "The motivation is clear and the in-depth analysis and ablations are convincing." (**Reviewer Fc4X**)
-  **Comprehensive Solution** "The authors present a comprehensive optimization framework that targets VLM distributed training from three perspectives: data, model, and memory balance." (**Reviewer dcak**)

**Strong Practical Contributions:**
- **Effective** "The paper centers on the efficient parallel training of MLLMs particularly focusing on balancing the computational load concerning data, model, and memory distribution both in-device and cross-device." (**Reviewer r51z**) 
- **Simple but Effective** "The multi-modal sample balancing strategy stands out for its simplicity, offering a straightforward approach to load management ..." (**Reviewer r51z**)

**Comprehensive Experiments:**
-  "It is comprehensively validated across different datasets, tasks, resolutions and hardware archs etc..." (**Reviewer r51z**)
- "Extensive experiments demonstrate the approach's effectiveness for rebalancing computation for MLLMs. "(**Reviewer Fc4X**)
</content>
</thread_object>

</thread>
<thread>
creation date:
1733196260629

thread_type:
official_comment

<thread_object>
creator: area_chair_akna

<content>
title:
Discussion due soon

comment:
Dear all reviewers,

Our reviewer-author discussion will end soon. For each of you, please check all the files and see if anything you'd like to discuss with authors.

Best,
Your AC
</content>
</thread_object>

</thread>


Metareview: