LongVILA: Scaling Long-Context Visual Language Models for Long Videos

🌈 Abstract

The article introduces LongVILA, a comprehensive solution for long-context visual-language models (VLMs). It addresses the challenges of training and deploying long-context VLMs by co-designing the algorithm and system. The key points are:

• LongVILA extends the context length of VLMs from 8 frames to 1024 frames, improving performance on long video captioning and instruction-following tasks.
• The training pipeline includes five stages: multi-modal alignment, large-scale pre-training, short supervised fine-tuning, context extension for language models, and long supervised fine-tuning.
• The Multi-Modal Sequence Parallelism (MM-SP) system efficiently scales the context length up to 2 million tokens without gradient checkpointing, achieving 2.1 to 5.7x speedup compared to ring sequence parallelism.
• LongVILA-8B demonstrates consistent accuracy improvements on long videos as the number of frames increases, and the inference system supports longer sequences than existing methods.

🙋 Q&A

[01] LongVILA Training Pipeline

1. What are the five stages of the LongVILA training pipeline? The five stages of the LongVILA training pipeline are:

• Stage 1: multi-modal alignment
• Stage 2: large-scale pre-training
• Stage 3: short supervised fine-tuning
• Stage 4: context extension for language models
• Stage 5: long supervised fine-tuning

2. What is the purpose of the context extension stage (Stage 4)? The purpose of the context extension stage (Stage 4) is to enhance the context length of the language model prior to engaging in supervised fine-tuning with long video datasets. This is done through a continuation of pre-training on text-only datasets to increase the context length to 262,144 tokens.

3. How does the long supervised fine-tuning stage (Stage 5) address the challenges of training on long videos? The long supervised fine-tuning stage (Stage 5) utilizes the newly constructed long video dataset, which contains 15,292 videos with diverse content and question-answer annotations. To efficiently train on these long videos, the authors developed the Multi-Modal Sequence Parallelism (MM-SP) system, which can scale the context length up to 2 million tokens.

[02] Multi-Modal Sequence Parallelism (MM-SP)

1. What are the key limitations of existing sequence parallelism systems when applied to multi-modal LLMs? The key limitations are:

• Modality heterogeneity: Existing systems treat visual and text modalities the same, leading to workload imbalance.
• Networking heterogeneity: Ring-style sequence parallelism ignores the significant difference in intra-node and inter-node network bandwidth.
• Limited maximal sequence length: DeepSpeed-Ulysses is limited by the number of attention heads.

2. How does the MM-SP system address these limitations?

• Modality heterogeneity: MM-SP uses a two-stage sharding strategy that first distributes images evenly, then aggregates global inputs for token-level sharding.
• Networking heterogeneity: MM-SP adopts a 2D-attention mechanism that combines A2A communication for head-dimension parallelism and P2P for sequence-dimension parallelism.
• Scalability: MM-SP can scale the context length up to 2 million tokens, outperforming existing systems.

3. How does the MM-SP inference system differ from the training system? The MM-SP inference system is designed to address the challenge of KV cache memory usage, which becomes a bottleneck for very long sequences. It deploys a distributed inference mode that leverages sequence parallelism to process longer sequences efficiently, achieving a speedup of 4.8x compared to the HuggingFace Pipeline parallelism approach.