Source-linked AI summary
VeOmni: Scaling Any Modality Model Training with Model-Centric Distributed Recipe Zoo
Qianli Ma, Yaowei Zheng, Zhelun Shi, Zhongkai Zhao, Bin Jia, Ziyue Huang, Zhiqi Lin, Youjie Li, Jiacheng Yang, Yanghua Peng, Zhi Zhang, Xin Liu
TL;DR
Omni-modal LLM training is hindered by heterogeneous architectures, limited end-to-end scalability, and frameworks that entangle model definitions with parallel logic. VeOmni separates computation from communication through composable model-centric distributed recipes and lightweight modality customization. It demonstrates competitive scaling across models and GPU configurations, including long-context training under sequence parallelism.
Problem
Existing frameworks provide limited support for scalable end-to-end omni-modal training because model definitions are tightly coupled with parallel logic.
Method
VeOmni uses composable model-centric distributed recipes, decoupled communication and computation, and a lightweight interface for customizing modality-specific modules.
Results
VeOmni achieves competitive efficiency and scalability across 8–128 GPUs and 7B–72B omni-modal models, while supporting up to 192K context for a 7B model with 61.5% MFU.
Takeaways & Limitations
VeOmni provides modular abstractions for efficient distributed training, rapid prototyping, and production-scale deployment of omni-modal LLMs.
Abstract
from arXiv · showhide
Recent advances in large language models (LLMs) have driven impressive progress in omni-modal understanding and generation. However, training omni-modal LLMs remains a significant challenge due to the heterogeneous model architectures required to process diverse modalities, necessitating sophisticated system design for efficient large-scale training. Existing frameworks typically entangle model definition with parallel logic, incurring limited scalability and substantial engineering overhead for end-to-end omni-modal training. We present VeOmni, a modular and efficient training framework to accelerate the development of omni-modal LLMs. VeOmni introduces model-centric distributed recipes that decouples communication from computation, enabling efficient 3D parallelism on omni-modal LLMs. VeOmni also features a flexible configuration interface supporting seamless integration of new modalities with minimal code change. Using VeOmni, a omni-modal mixture-of-experts (MoE) model with 30B parameters can be trained with over 2,800 tokens/sec/GPU throughput and scale to 160K context lengths via 3D parallelism on 128 GPUs, showcasing its superior efficiency and scalability for training large omni-modal LLMs.
1 Introduction
Omni-modal LLM training is difficult because heterogeneous architectures expose gaps in existing end-to-end infrastructure and tightly coupled parallel logic limits scalability. VeOmni addresses this with model-centric, composable distributed recipes and lightweight modality customization, demonstrating competitive scaling across 8–128 GPUs and 7B–72B models.
- Motivation: Existing frameworks rarely support end-to-end any-to-any training, leaving scalable infrastructure for fully omni-modal LLMs underdeveloped.Mature systems primarily target text-to-text training, while only a few support any-to-text tasks.
- Motivation: Tightly coupling model definitions with parallel logic causes load imbalance, poor scalability, and difficult generalization to increasingly heterogeneous omni-modal architectures.Directly applying tensor and pipeline parallelism designed for language models does not adequately accommodate complex omni-modal models.
- VeOmni: VeOmni decouples model computation from communication through model-centric distributed recipes that compose FSDP, SP, and EP into 2D or 3D parallel strategies.A high-level parallel plan API applies distributed strategies to model blocks and supports dense or MoE training recipes.
- VeOmni: VeOmni’s lightweight configuration interface lets developers customize modality-specific encoders and decoders while integrating new modules with minimal engineering effort.Computation modules need not account for distributed concerns, and the interface supports both multimodal understanding and generation.
- Evaluation: VeOmni demonstrates competitive efficiency and scalability across 8–128 GPUs for omni-modal models ranging from 7B to 72B parameters.This result covers the framework’s reported scaling range across model sizes and GPU counts.
2 Related Work
Related work spans mature distributed frameworks for text training and specialized systems for multimodal any-to-text workloads. These systems motivate VeOmni’s model-centric design for natively scaling omni-modal models across modalities.
- Text Training Frameworks: Megatron-LM established optimized transformer blocks with tensor and pipeline parallelism as a foundation for large-scale text training frameworks.Colossal-AI extends this direction with comprehensive 3D parallel training strategies.
- Multimodal Training Frameworks: Specialized multimodal frameworks address any-to-text training through model and data heterogeneity handling, feedback mechanisms, or bubble exploitation.Examples include DistMM, DistTrain, Align Anything, and Optimus, each targeting particular multimodal training challenges.
- VeOmni: VeOmni differs by providing a model-centric framework natively suited to omni-modal training, with distributed strategies that scale arbitrary modalities on large clusters.Its scope extends beyond the specialized any-to-text systems described in related work.
3 VeOmni: Scalable Omni-Modal Training Framework
VeOmni provides a model-centric framework for scalable omni-modal training by decoupling model computation from distributed strategies and supporting modular modality integration. Its recipe system composes parallelism and system optimizations for long sequences, MoE models, and heterogeneous architectures.
- 3.1 Light-Weight Customization for Omni-Modal LLMs: VeOmni uses a plug-and-play architecture that attaches multimodal encoders and decoders to a foundation model through fully decoupled modules.The encoder processes input modalities, the foundation model performs central computation, and decoders support target-specific training and inference pathways.
- 3.2 Model-Centric Distributed Recipe Zoo: VeOmni composes FSDP, SP, and EP across heterogeneous components to support scalable training of long-sequence and MoE-based omni-modal models.The framework targets modality-specific requirements by applying distinct parallelism techniques to different architectural components.
- 3.2 Model-Centric Distributed Recipe Zoo: Its model-centric recipes separate distributed training mechanisms from model definitions, allowing flexible composition without modifying low-level parallelization code.The framework exposes high-level interfaces for applying and composing strategies across model components.
- 3.2.1 Fully Sharded Data Parallel for Large Model Training: HSDP reduces communication overhead by combining sharding within groups with data replication across groups, and switching to it requires only a configuration change.This 2D device-mesh design is intended to improve scalability while preserving the non-intrusive model interface.
- 3.2.5 Other System Optimization Strategies: DeepSpeed-Ulysses and Async-Ulysses address ultra-long sequences by using sequence parallelism and overlapping all-to-all communication with linear projections.The overlap hides part of communication latency behind computation and improves training performance.
- 3.2.5 Other System Optimization Strategies: VeOmni integrates additional system optimizations through modular implementations decoupled from model computation, enabling reuse across model architectures with minimal code changes.This design principle supports the framework’s broader recipe composition and extensibility goals.
4 Experiments
VeOmni is evaluated across diverse omni-modal models, datasets, GPU scales, context lengths, and parallel strategies. The experiments show efficient long-context and MoE training together with stable multimodal convergence.
- Experimental Setup: VeOmni supports training across 8–128 GPUs on dense 7B and 72B models and a 30B MoE omni-modal model.The evaluation covers text, image, video, audio, and image-generation workloads.
- Comparison of Training Recipes under Different Scenarios: 192K context is supported for the 7B model at 61.5% MFU, while the 72B model supports 96K context at 54.82% MFU.Increasing sequence parallelism enables substantially longer training contexts.
- Comparison of Training Recipes under Different Scenarios: 160K context is supported for the 30B MoE model on 128 GPUs using FSDP, sequence parallelism, and expert parallelism.Moderate sequence- and expert-parallel sizes provide a balanced trade-off between context length and throughput.
- Comparison of Training Recipes under Different Scenarios: VeOmni minimizes overhead for short sequences while preserving scalability for long-sequence and MoE training.The results attribute this behavior to effective use of sequence and expert parallelism.
- Convergence Study on Omni-Modal LLMs: All three evaluated omni-modal models converge stably across multimodal understanding and generation tasks.The study covers text, image, video, and audio understanding, plus text and image synthesis, using language-model and decoder losses.
5 Conclusion
VeOmni is a model-centric framework for efficient distributed training of omni-modal models. Its composable recipes and modular APIs support scalable training, while the authors identify pipeline parallelism and modality-aware balancing as future extensions.
- Conclusion: VeOmni combines multiple parallelism methods and system optimizations in a composable recipe-based design for efficient omni-modal training.The framework separates model-centric abstractions from distributed execution and supports large-scale vision-language-audio applications.
- Conclusion: VeOmni provides high throughput, scalability, and developer-friendly abstractions for prototyping and production-scale deployment.Future work targets non-intrusive pipeline parallelism and modality-aware sequence-parallel data balancing.
A Comparison of VeOmni and TorchTitan on Training Large Language Models
VeOmni is benchmarked against TorchTitan on text-only models because existing frameworks do not support direct omni-modal comparison. The results report higher throughput and memory efficiency across tested sizes and sequence lengths, including settings beyond TorchTitan’s limits.
- Evaluation Setup: The benchmark covers Qwen2-7B, Qwen2.5-32B, and Qwen2-72B on 128 GPUs under matched mixed-precision and full activation-checkpointing settings.Each configuration uses micro-batch size 1 and global batch size 128.
- Comparison Results: VeOmni consistently achieves higher throughput and memory efficiency than TorchTitan across model sizes and sequence lengths.The comparison uses controlled text-only experiments because direct omni-modal benchmarking is infeasible.
- Comparison Results: VeOmni supports long-sequence and MoE training beyond TorchTitan’s memory or capability limits.Table 4 evaluates Qwen3-30B-A3B on 128 GPUs, a setting TorchTitan does not support.
B.1 Data Processing and Collating
VeOmni’s data-collation pipeline provides a unified interface for batching text, image, video, and audio inputs. Its packing and concatenation controls support flexible multimodal alignment.
- Data Processing and Collating: The OmniDataCollatorWithPacking class unifies processing and batching for text, image, video, and audio inputs.The pipeline handles modality-specific input and output masks and features.
- Data Processing and Collating: Fine-grained feature packing and concatenation simplify preparation of batched multimodal training data.The design supports token-level and feature-level alignment while easing integration of new modalities.
B.2 Omni-Modal LLMs Initialization
VeOmni initializes omni-modal models through modular encoder–foundation–decoder components and abstracts distributed parallelization behind unified APIs. Its pipeline supports both training and modality-specific generation while coordinating multiple n-D parallelism strategies.
- Model Initialization: VeOmni’s build_omni_model API assembles modality encoders, decoders, and a foundation model from configuration and component paths.The API initializes the resulting model on a designated device while hiding underlying construction logic.
- Model Computation: The OmniModel forward pipeline passes inputs through an encoder, decoder, and foundation model, returning loss and hidden states.The architecture supports a modular encoder–foundation–decoder design across training and inference.
- Inference: During inference, generated text selects modality-specific branches whose decoders produce embeddings for next-token prediction and final modality outputs.The process invokes decoder lm_generate functions after token prediction is complete.
- Model Computation: OmniEncoder converts text and modality features into embeddings and inserts them into the input sequence through modality-specific masks.Image, video, and audio encoders provide the corresponding modality embeddings before the foundation model processes the combined inputs.
- Parallel State Management: VeOmni’s parallel_state API uses DeviceMesh to declaratively organize DP, TP, EP, PP, and SP configurations and expose global groups and ranks.This design decouples low-level distributed logic from model implementation.
- Model Parallelization: The build_parallelize_model API applies FSDP and customized parallel plans while abstracting parameter sharding, precision, activation checkpointing, and optional offloading.The interface is intended to enable large-scale distributed training with minimal model-code changes.
B.6 Long-Context Attention
VeOmni extends FlashAttention to support Ulysses-style sequence parallelism without changing the underlying attention logic. It converts tensor layouts around attention computation and restores the sequence-sharded output layout afterward.
- Long-Context Attention: VeOmni’s flash_attention_forward integrates Ulysses sequence parallelism with HuggingFace FlashAttention while preserving modular compatibility.The extension requires no changes to the model’s attention logic.
- Long-Context Attention: When enabled, sequence-sharded query, key, and value tensors are transformed into head-sharded layouts before standard FlashAttention computation.The query, key, and value transformations use Ulysses all-to-all operations.
- Long-Context Attention: After attention, the output is transformed back from head-sharded to sequence-sharded format for downstream processing.This restores the layout expected by the sequence-parallel pipeline.
B.7 Expert Parallelism Implementation
VeOmni defines expert parallelism declaratively through ParallelPlan, allowing selected MoE modules to be sharded across expert-parallel groups. Wildcard module matching and DTensor placements connect the plan to model parallelization.
- Expert Parallelism Implementation: ParallelPlan specifies which MoE expert projection layers are sharded across expert-parallel groups.The example targets gate_proj, up_proj, and down_proj modules within transformer-layer experts.
- Expert Parallelism Implementation: The expert-parallel plan uses wildcard module-name matching and DTensor-compatible Shard(0) placements before being passed to build_parallelize_model().This declarative interface separates expert sharding choices from lower-level distributed implementation.