OmniBal: Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance

TL;DR

OmniBal addresses computation imbalance in 3D parallel training for vision-language models by balancing data, model, and memory. It introduces Balanced Dynamic Mini-Batch (ISF), Balanced Model Partitioning (BMP), and Balanced Adaptive Re-Computation to create a cohesive omniverse balancing framework. Experimental results across multiple models, datasets, and backends demonstrate significant speedups (up to several-fold) while preserving performance. The approach generalizes across architectures and hardware, enabling faster instruction-tuning for diverse VLMs and settings.

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.

Figures (4)

• Figure 1: Overview of the computation imbalanced problem and our proposed solution in the Standard Vision-Language instruction-tuning framework. We consider the bottleneck issues of data, model, and memory, and propose an omniverse solution addressing these three aspects, each providing the foundation for the next.
• Figure 2: The Problem of Computation Imbalance in VLM Instruction-Tuning 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.
• Figure 3: The pipeline consists of four stages, labeled Stage-1 to Stage-4, each representing a different stage of pipeline parallelism. Within this structure, "ViT" stands for the Vision Transformer layer, while "LLM" refers to the Transformer layer used for large language models (LLM). Regarding computational execution, the darker-colored sections signify forward passes with re-computation. In contrast, the lighter-colored sections denote a standard forward pass without re-computation.
• Figure 4: The convergence of ISF in various scenarios, including (a) different datasets, (b) different patch sizes, and (c) different image resolutions.