VeOmni: Scaling Any Modality Model Training with Model-Centric Distributed Recipe Zoo

发表时间: 2025-08 · arXiv:2508.02317

原文: https://arxiv.org/abs/2508.02317

文章标题: VeOmni: 通过以模型为中心的分布式配方库扩展任何模态模型的训练
作者/机构: Qianli Ma, Yaowei Zheng, Zhelun Shi, Zhongkai Zhao, Bin Jia, Ziyue Huang, Zhiqi Lin, Youjie Li, Jiacheng Yang, Yanghua Peng, Zhi Zhang, Xin Liu (ByteDance Seed)

A1 主要贡献

核心问题: 随着大型语言模型（LLMs）向全模态理解和生成发展，其模型架构变得日益异构和复杂。现有的训练框架通常将模型定义与并行逻辑紧密耦合，导致在端到端全模态训练中可扩展性有限，并带来巨大的工程开销。特别是，现有框架很少支持“任意到任意”（any-to-any）的全模态学习场景，这揭示了构建下一代全模态LLMs需要可扩展基础设施的必要性。

研究目标: 本文旨在解决全模态LLMs训练中的效率和可扩展性挑战，提出一个名为VeOmni的模块化、高效的训练框架，以加速全模态LLMs的开发。该框架的目标是解耦模型定义与并行逻辑，支持灵活的并行策略组合，并以最小的工程成本集成新的模态。

创新点:
1. 提出以模型为中心的分布式配方 (Model-Centric Distributed Recipes)：通过可组合的n-D并行（如FSDP, SP, EP）来高效扩展全模态训练，同时将工程开销降至最低。这种设计将通信与计算分离，允许灵活地为密集或MoE模型量身定制训练配方。
2. 引入轻量级配置接口: 该接口支持灵活定制全模态LLMs，使用户能够轻松修改特定模态的编码器和解码器，从而支持多模态的理解和生成，且只需极少的代码改动。
3. 验证了VeOmni的卓越效率和可扩展性: 在8到128个GPU上，针对7B到72B参数的模型，在全模态场景下展示了VeOmni的性能。例如，一个30B参数的全模态MoE模型，在128个GPU上通过3D并行，训练吞吐量超过2800 tokens/sec/GPU，并能扩展到160K的上下文长度。

A2 方法细节

3. VeOmni：可扩展的全模态训练框架

VeOmni设计了一个以模型为中心的框架，该框架原生适用于训练全模态模型。此框架包含多样化的分布式训练策略，使得任何模态都能轻松扩展至大规模集群。

3.1 面向全模态LLMs的轻量级定制

即插即用的模块化架构。为了支持任意模态的训练，我们提出了一个即插即用的架构，该架构允许将任何多模态编码器和解码器的组合灵活地附加到基础模型的输入和输出端，如图2所示。这种模块化且完全解耦的设计将系统级操作与模型特定的计算解耦，使得流式管线能够独立于模型管线运行。因此，VeOmni能够轻松扩展和缩放多样化的全模态LLM。

统一接口协议。我们的核心设计围绕为每个组件量身定制的轻量级协议展开。具体来说，编码器和解码器模块通过继承HuggingFace的PreTrainedModel类并结合一个任务特定的mixin来实现一个统一接口。这种设计确保了与标准训练工作流的兼容性，并使得所有模块都能被一致地注册、初始化和执行。

训练阶段的输入输出处理。在训练阶段，对于输入模态，每个编码器实现一个lm_encode函数，该函数将原始模态数据编码为词元嵌入（token embeddings），然后将其插入到基础模型的输入嵌入中。同样地，解码器模块也实现lm_encode函数，以提供代表目标输出词元的训练时输入。基础模型的最终输出使用lm_head函数解码为目标模态，从而完成端到端的映射。

推理阶段的生成流程。在推理阶段，对于每个预测步骤，输出的隐藏状态会通过相应模态解码器的lm_embed函数，该函数生成用作基础模型下一词元预测输入的嵌入。一旦所有词元都被预测完毕，解码器的lm_generate函数将被调用，以生成最终的特定模态输出。训练和推理过程的详细信息在附录B.3中列出。

3.2 以模型为中心的分布式配方库

模块化的分布式策略套件。为了支持全模态LLM的大规模分布式训练，VeOmni提供了一套模块化的分布式训练策略，包括完全分片数据并行（FSDP）、序列并行（SP）和专家并行（EP），旨在应对诸如超长序列训练和高效扩展大型专家混合（MoE）模型等挑战。VeOmni使用户能够根据其模型架构和特定模态的需求，自由组合分布式训练配方，而无需管理底层实现细节。

设计原则。该框架的设计遵循以下原则：
1. 非侵入式的分布式训练API。VeOmni将模型架构与底层的分布式训练机制解耦，使用户能够灵活地组合全模态模型，而无需修改底层的并行化代码。
2. 支持长序列训练。VeOmni能够适应涉及跨不同模态的极长序列的训练工作负载，这是全模态训练中的一个常见挑战。
3. MoE模型的高效扩展。通过集成的专家并行，VeOmni促进了大规模MoE模型的高效训练，实现了可扩展性和计算效率。

后续章节介绍。在接下来的章节中，我们将介绍VeOmni中的每一种分布式技术，并展示它们如何被集成以支持可扩展的全模态训练。

3.2.1 用于大型模型训练的完全分片数据并行

FSDP的优势。完全分片数据并行（FSDP）是PyTorch中零冗余优化器（ZeRO）【46，Zero: Memory optimizations toward training trillion parameter models，2020，SC20】的一种高效实现。其主要优势在于显著减少了训练期间每个GPU所需的内存。通过将模型的参数、梯度和优化器状态分片到所有可用设备上，FSDP允许训练那些否则会超出单个GPU内存容量的超大型模型。

FSDP的非侵入式设计。FSDP的一个关键优势是其非侵入式设计，它将模型的架构与并行化策略解耦。这意味着开发人员可以专注于设计他们的模型，而无需修改底层代码来适应分布式训练设置。这种非侵入式的特性使得FSDP非常适合训练全模态LLM。由于这些模型以其复杂和异构的架构为特征，FSDP能够在不需对模型代码进行任何更改的情况下并行化训练，使其成为一个理想的解决方案。VeOmni集成了FSDP1【80，Pytorch fsdp: experiences on scaling fully sharded data parallel，2023，arXiv】和FSDP2【36，TorchTitan: One-stop pytorch native solution for production ready llm pretraining，2025，ICLR】作为其分布式训练栈的基础组件，并提供统一的API以便于配置。更多细节请参见附录B.5。

HSDP的集成。为了进一步提高训练效率，VeOmni集成了混合分片数据并行（HSDP），这是FSDP的一个扩展，旨在最小化通信开销。HSDP利用一个2D设备网格，在“分片组”内结合FSDP，并在“副本组”间使用分布式数据并行（DDP）。这种混合方法极大地减少了节点间通信，从而实现了更大的可扩展性。切换到更高效的HSDP就像在配置中更改一个参数一样简单，无需对底层代码进行任何修改。

3.2.2 用于长序列训练的序列并行

长序列训练的挑战。随着最先进的全模态LLM的发展，训练所需的序列长度已显著增长，尤其是在高分辨率图像或视频理解与生成等领域【18，Seedance 1.0: Exploring the boundaries of video generation models，2025，arXiv；66，Qwen2-vl: Enhancing vision-language model’s perception of the world at any resolution，2024，arXiv；78，Video instruction tuning with synthetic data，2024，arXiv】。这种向更长序列发展的趋势在计算成本和内存消耗方面都构成了重大挑战。因此，高效地应对长序列训练的需求对于扩展模型容量和释放全模态架构的全部潜力至关重要。

DeepSpeed-Ulysses的采用。为了解决这个问题，VeOmni采用了DeepSpeed-Ulysses【29，Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models，2023，arXiv】，这是一种专为在超长序列上进行Transformer训练而设计的高效序列并行技术。DeepSpeed-Ulysses沿序列维度切分激活，并在注意力计算期间策略性地编排all-to-all通信，确保在序列长度和设备数量按比例扩展时通信量保持不变。这种设计为长序列训练提供了高效率和可扩展性。

非侵入式接口与FlashAttention集成。VeOmni的一个核心优势是其为部署DeepSpeed-Ulysses提供了非侵入式接口。开发人员可以无缝地利用DeepSpeed-Ulysses，而无需修改模型构建代码或直接与底层实现细节交互。特别是，VeOmni通过集成DeepSpeed-Ulysses增强了FlashAttention【9，FlashAttention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS】的实现，同时保持与默认注意力的完全兼容，从而使得在序列并行工作流中能够直接采用先进的注意力加速方法。更多细节请参见附录B.6。

Async-Ulysses的引入。此外，为了优化训练吞吐量，我们引入了Async-Ulysses，这是一个旨在重叠通信和计算的增强实现。该版本调度all-to-all通信操作与线性投影计算并发执行。通过将大部分通信延迟有效地隐藏在计算背后，这种方法在训练性能上取得了显著的提升。

3.2.3 用于MoE模型扩展的专家并行

MoE架构的重要性。专家混合（MoE）架构已成为高效扩展大型模型的主流解决方案【10，Deepseek-v3 technical report，2025，arXiv】，这主要归功于其在推理和训练期间稀疏激活参数子集的能力，从而在计算效率和模型容量方面实现了显著提升。在全模态基础模型的背景下，采用MoE架构作为骨干不仅降低了训练成本，还通过根据输入内容动态分配专家容量，增强了模型在不同模态上的性能。

VeOmni对专家并行的支持。为了支持基于MoE的全模态LLM的可扩展训练，VeOmni引入了一个用户友好且灵活的专家并行接口，允许用户轻松地在设备间应用专家分片而无需手动配置。该接口与广泛使用的MoE变体兼容，并支持即插即用集成，从而显著降低了实现分布式专家并行的门槛。更多细节请参见附录B.7。

通信计算重叠优化。大规模MoE训练的一个主要瓶颈在于将词元路由到其指定专家所需的all-to-all通信，这会引入巨大的延迟。为了缓解这个问题，VeOmni融合了受近期进展【4，Flux: Fast software-based communication overlap on gpus through kernel fusion，2024，arXiv；76，Comet: Fine-grained computation-communication overlapping for mixture-of-experts，2025，arXiv】启发的细粒度通信-计算重叠技术。这些工作通过将集体操作与本地专家计算并发调度来隐藏通信延迟，从而无需像DualPipe【10，Deepseek-v3 technical report，2025，arXiv】那样复杂的流水线级解决方案。与通常僵化且对特定模态不平衡敏感的以流水线为中心的设计不同，VeOmni的算子级方法是轻量级、模型无关的，并且特别适合多模态设置。这在大规模MoE训练期间带来了更高的利用率和更快的迭代速度。

3.2.4 可组合的n-D并行训练

并行策略的组合性。在VeOmni中，核心并行策略（即完全分片数据并行FSDP、序列并行SP和专家并行EP）被设计为完全可组合的，并且可以灵活地应用于全模态LLM的不同组件。这种可组合性对于全模态训练尤其有利，因为不同的模态或架构组件可能会从不同的并行技术中受益。例如，视觉编码器可以根据其规模采用FSDP或标准数据并行，而语言主干可以利用EP用于MoE层和SP来支持长序列处理的组合。这种细粒度的灵活性使得在多样化的模型架构上能够进行高效和可扩展的训练。

n-D并行设计概览。图3概述了VeOmni的n-D并行设计。图的左侧说明了全模态模型的通信和计算流，其中不同的特定模态编码器处理各自的输入并将中间特征分发给下游模块。右侧可视化了n-D并行策略在模型不同部分的应用，展示了各种并行技术如何在一个单一的训练配置中共存。

统一的分布式训练抽象。为了支持这种灵活的并行组合，VeOmni引入了一个基于全局设备网格【36，TorchTitan: One-stop pytorch native solution for production ready llm pretraining，2025，ICLR】的统一分布式训练抽象。与需要手动管理多个进程组的方法【41，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC；52，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】不同，VeOmni显著简化了n-D并行的配置和协调（附录B.4）。这不仅降低了进程组管理的复杂性，还提高了可扩展性，使其更容易适应和扩展到未来的并行策略或新的模型组件。

3.2.5 其他系统优化策略

模块化优化集成。除了上述的并行策略，VeOmni还集成了一系列广泛的系统级优化，遵循了VeOmni关键的设计原则，即这些优化实现与模型的计算逻辑解耦。这种模块化架构允许无缝集成，使得这些系统级增强可以轻松应用于各种模型架构，而只需对模型代码进行最少甚至无需修改。这些优化包括但不限于以下策略：

• 动态批处理 (Dynamic Batching)。将批次中的所有样本填充到固定的序列长度通常会导致大量的计算效率低下，尤其是在样本间序列长度差异显著时。为了缓解此问题并提高训练效率，VeOmni采用了一种动态批处理机制，该机制在缓冲区中累积样本，并策略性地将它们打包以接近目标序列长度。利用FlashAttention【8，Flashattention-2: Faster attention with better parallelism and work partitioning，2023，arXiv】，这种打包策略能够在最小化填充开销的情况下高效利用批次预算，同时保持跨样本注意力计算的正确性。

• 高效核函数 (Efficient Kernels)。为最大化训练吞吐量，VeOmni集成了一套高度优化的算子核，包括来自liger-kernel【23，Liger kernel: Efficient triton kernels for llm training，2025，arXiv】的RMSNorm、LayerNorm、RoPE、SwiGLU和in-place cross-entropy，以及FlashAttention【8，Flashattention-2: Faster attention with better parallelism and work partitioning，2023，arXiv；9，FlashAttention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS；50，Flashattention-3: Fast and accurate attention with asynchrony and low-precision，2024，NeurIPS】和MoE特定操作【4，Flux: Fast software-based communication overlap on gpus through kernel fusion，2024，arXiv；76，Comet: Fine-grained computation-communication overlapping for mixture-of-experts，2025，arXiv】。这些核函数经过精心设计，以实现高性能和广泛的兼容性，从而能够在多样化的基于Transformer的架构和模型变体上高效执行。

• 内存优化 (Memory Optimization)。VeOmni集成了层级重计算【6，Training deep nets with sublinear memory cost，2016，arXiv】、激活卸载和优化器状态卸载，以大幅减少训练期间的内存消耗。这些节省内存的技术使得每个GPU可以使用更大的微批次大小，从而改善通信成本的摊销，并促进与完全分片数据并行（FSDP）引入的通信开销更好地重叠，最终提高整体训练效率。

• 高效的分布式检查点 (Efficient Distributed Checkpointing)。VeOmni利用ByteCheckpoint【64，Bytecheckpoint: A unified checkpointing system for large foundation model development，2024，arXiv】来支持在不同的分布式配置下高效地保存和恢复检查点，且开销极小。除了促进弹性训练和提高集群利用率外，我们的框架进一步扩展了ByteCheckpoint以支持全模态模型。此扩展确保了异构模型组件的一致和可靠的保存与加载。

• 元设备初始化 (Meta Device Initialization)。VeOmni支持在元设备上进行大型模型初始化，而无需在初始化期间分配物理内存。在元设备上实例化模型后，我们通过将参数转换为DTensor格式【36，TorchTitan: One-stop pytorch native solution for production ready llm pretraining，2025，ICLR】来执行参数分片和并行加载，从而显著加速了大规模模型的初始化和加载过程。

A4 实验环境

• 硬件配置:
  • GPU: 在包含8到128个GPU的大规模生产级GPU集群上进行评估。
• 模型与数据集:
  • 模型:
    • 密集模型: Qwen2-VL 7B 和 72B 【66，Qwen2-vl: Enhancing vision-language model’s perception of the world at any resolution，2024，arXiv】。
    • MoE模型: 基于Qwen3-30B-A3B【58，Qwen3 technical report，2025，arXiv】构建的全模态专家混合（MoE）LLM。
    • 其他对比模型: LLaMA#Omni (使用LLaMA【61，Llama: Open and efficient foundation language models，2023，arXiv】作为基础模型) 和 Janus 【68，Janus: Decoupling visual encoding for unified multimodal understanding and generation，2025，CVPR】。
  • 数据集:
    • 文本理解: FineWeb100T 【43，The fineweb datasets: Decanting the web for the finest text data at scale，2024，NeurIPS】。
    • 图像理解: ShareGPT4V 【5，Sharegpt4v: Improving large multi-modal models with better captions，2024，ECCV】。
    • 视频理解: LLaVA-Video 【78，Video instruction tuning with synthetic data，2024，arXiv】。
    • 音频理解: Voice Assistant 【70，Mini-omni: Language models can hear, talk while thinking in streaming，2024，arXiv】。
    • 图像生成: ImageNet 【13，Imagenet: A large-scale hierarchical image database，2009，CVPR】。
• 实验设置:
  • 训练方式: 所有实验中，特定模态的编码器和解码器被冻结，而包括基础模型和多模态投影器在内的其余组件则进行完全微调。
  • 工作负载: 输入上下文长度从8K逐步扩展到256K个词元。
  • 评估指标:
    • 训练吞吐量 (tokens per second per GPU)。
    • 模型FLOPs利用率 (MFU) 【7，Palm: Scaling language modeling with pathways，2023，JMLR】。
    • 损失收敛行为。

A4 实验结果

训练配方在不同场景下的比较

• 实验内容: 在不同模型规模（7B, 72B, 30B MoE）和GPU数量（8, 128）下，比较了不同训练配方（2D并行: FSDP+SP; 3D并行: FSDP+SP+EP）的效率。通过逐步增加上下文长度来评估系统的可扩展性。
• 实验结果:
  • Qwen2-VL 7B (8 GPUs): 如图4所示，增加序列并行度可以处理更长的上下文。使用4路序列并行，模型可以支持高达192K的上下文长度，MFU达到61.5%。
  • Qwen2-VL 72B (128 GPUs): 如图5所示，对于更大的模型，序列并行同样有效。使用8路序列并行，模型可以处理96K的上下文长度，MFU为54.82%。
  • 30B MoE模型 (128 GPUs): 如图6(a)所示，在3D并行设置下，中等程度的SP和EP组合实现了良好的平衡，支持高达160K的上下文长度，同时保持了有竞争力的吞吐量。
• 分析结论: 实验结果强调了VeOmni在处理长序列训练和扩展MoE模型方面的效率。VeOmni在短序列场景下开销很小，同时能有效利用序列和专家并行来维持长序列和MoE设置下的可扩展性和可行性。

全模态LLMs的收敛性研究

• 实验内容: 在多模态理解（文本、图像、视频、音频）和生成（文本、图像）任务上，评估了三个结构不同的全模态LLM（Janus, LLaMA#Omni, Qwen3-MoE#Omni）的收敛行为。
• 实验结果: 如图6(b)所示，所有模型在理解和生成任务上都表现出稳定的收敛性。"LM loss"（文本词元的交叉熵损失）和"decoder loss"（图像词元的交叉熵损失）均平稳下降。
• 分析结论: 结果表明VeOmni能够为大型全模态LLM提供高效且鲁棒的训练支持。

A5 结论

本文介绍了VeOmni，一个以模型为中心的分布式训练框架，旨在高效地扩展任何模态模型的训练。通过将多种并行方法和系统优化集成到一个可组合的、基于配方的设计中，VeOmni实现了在全模态LLM上无缝且高效的分布式训练策略。我们进一步展示了系统级优化、模块化API以及在大型视觉-语言-音频模型上的实际训练应用。实验分析表明，VeOmni不仅实现了高吞吐量和可扩展性，还为快速原型设计和生产规模部署提供了对开发者友好的抽象。

未来工作:
1. 计划将VeOmni扩展以支持非侵入式的流水线并行，从而进一步解耦模型定义和并行执行。
2. 旨在通过具有模态感知的数据平衡策略来增强序列并行，以更好地支持多模态训练场景。

A6 附录

A. VeOmni与TorchTitan在训练大型语言模型上的比较

引言。在主论文中，我们展示了VeOmni支持全模态模型的高效大规模训练，这是其他现有框架尚不支持的能力。因此，在全模态训练场景中进行直接比较是不可行的。为了提供一个公平的基准，我们对大规模纯文本模型进行了受控实验，以比较VeOmni与最先进的分布式训练框架TorchTitan【36，TorchTitan: One-stop pytorch native solution for production ready llm pretraining，2025，ICLR】。

基准测试结果。这些基准测试结果表明，VeOmni在一系列模型大小和序列长度上始终实现了更高的吞吐量和内存效率。值得注意的是，VeOmni支持长序列和MoE模型训练，这些训练超出了TorchTitan的内存或能力限制。

表1 VeOmni与TorchTitan在使用Qwen2-7B [66] 于128个GPU上的性能比较。两种配置均采用混合精度训练和完全激活检查点，微批次大小为1，全局批次大小为128。

表2 VeOmni与TorchTitan在使用Qwen2.5-32B [45] 于128个GPU上的性能比较。两种配置均采用混合精度训练和完全激活检查点，微批次大小为1，全局批次大小为128。

表3 VeOmni与TorchTitan在使用Qwen2-72B [72] 于128个GPU上的性能比较。两种配置均采用混合精度训练和完全激活检查点，微批次大小为1，全局批次大小为128。

表4 VeOmni在使用Qwen3-30B-A3B [58] 于128个GPU上的性能。我们采用混合精度训练和完全激活检查点，微批次大小为1，全局批次大小为128。（TorchTitan不支持此设置。）

B. VeOmni的API设计

B.1 数据处理与整理

统一接口。列表1展示了VeOmni中的数据整理流水线，它为处理和批处理全模态输入（包括文本、图像、视频和音频）提供了一个统一的接口。OmniDataCollatorWithPacking类能够对特征打包和拼接进行细粒度控制，使得为不同模态准备批处理训练数据变得简单直接。这种模块化设计简化了新模态的集成，并支持灵活的词元级和特征级对齐策略。最终的train_dataloader通过build_dataloader()构建，它会自动与并行训练配置对齐。

# 列表1 数据处理与整理
from veomni.data import (
    OmniDataCollatorWithPacking, build_dataloader,
)

data_collate_fn.append(
    OmniDataCollatorWithPacking(
        packing_features=[
            "input_ids",
            "attention_mask",
            "labels",
            "position_ids",
            "image_input_mask",
            "image_output_mask",
            "video_input_mask",
            "video_output_mask",
            "audio_input_mask",
            "audio_output_mask",
        ],
        concat_features=[
            "image_input_features",
            "image_output_features",
            "video_input_features",
            "video_output_features",
            "audio_input_features",
            "audio_output_features",
        ],
    )
)

train_dataloader = build_dataloader(
    dataset=train_dataset,
    micro_batch_size=micro_batch_size,
    global_batch_size=global_batch_size,
    collate_fn=data_collate_fn,
)

B.2 全模态LLM初始化

统一入口点。列表2展示了VeOmni的build_omni_model API的用法，它作为构建全模态模型的统一入口点。给定配置文件和模型组件路径（包括特定模态的编码器、解码器和基础模型），此函数会自动在指定设备上组装和初始化模型。通过抽象底层的初始化逻辑，VeOmni确保用户可以轻松地以模块化和可扩展的方式实例化复杂的全模态模型，同时保持与分布式并行化工作流的兼容性。

from veomni.models import build_omni_model

model = build_omni_model(
    config_path=config_path,
    weights_path=model_path,
    encoders=encoders,
    decoders=decoders,
    foundation=foundation,
    init_device=init_device,
)

B.3 全模态模型计算

即插即用架构实现。列表3-4展示了VeOmni中名为OmniModel的即插即用架构的部分前向过程实现，该架构采用编码器-基础模型-解码器结构。代码片段演示了在训练和推理模式下的前向逻辑如何将每个模块（即编码器、基础模型和解码器）集成到一个无缝且模块化的流水线中。在推理期间，控制流会根据生成的文本特殊词元动态切换到特定模态的生成分支，从而实现灵活的全模态输出。在推理期间，OmniModel支持并行化生成和无分类器引导【68，Janus: Decoupling visual encoding for unified multimodal understanding and generation，2025，CVPR】。这种架构突显了编码器、基础模型和解码器模块之间的清晰解耦，允许每个组件独立定义和执行其自己的前向逻辑。

# 列表3 OmniModel训练过程
class OmniEncoder:
    def forward(self, **inputs):
        inputs_embeds = self.text_encoder(inputs['input_ids'])
        image_input_embeds = self.image_encoder.lm_encode(inputs.get("image_input_features"))
        video_input_embeds = self.video_encoder.lm_encode(inputs.get("video_input_features"))
        audio_input_embeds = self.audio_encoder.lm_encode(inputs.get("audio_input_features"))
        inputs_embeds = self.masked_scatter(
            inputs_embeds,
            image_input_embeds, inputs.get("image_input_mask"),
            video_input_embeds, inputs.get("video_input_mask"),
            audio_input_embeds, inputs.get("audio_input_mask"),
        )
        return inputs_embeds

class OmniDecoder:
    def forward(self, inputs_embeds, **inputs):
        image_output_embeds = self.image_decoder.lm_encode(inputs.get("image_output_features"))
        video_output_embeds = self.video_encoder.lm_encode(inputs.get("video_output_features"))
        audio_output_embeds = self.audio_encoder.lm_encode(inputs.get("audio_output_features"))
        inputs_embeds = self.masked_scatter(
            inputs_embeds,
            image_output_embeds, inputs.get("image_output_mask"),
            video_output_embeds, inputs.get("video_output_mask"),
            audio_output_embeds, inputs.get("audio_output_mask"),
        )
        return inputs_embeds

    def lm_head(self, hidden_states, **inputs):
        loss = self.image_decoder.lm_head(outputs.hidden_states, **inputs)
        loss += self.video_decoder.lm_head(outputs.hidden_states, **inputs)
        loss += self.audio_decoder.lm_head(outputs.hidden_states, **inputs)
        return loss

class OmniModel:
    encoder: OmniEncoder
    decoder: OmniDecoder

    def forward(self, **inputs):
        inputs_embeds = self.encoder(**inputs)
        inputs_embeds = self.decoder(inputs_embeds, **inputs)
        outputs = self.foundation(inputs_embeds, **inputs)
        hidden_states = outputs.hidden_states
        loss = None
        if self.training:
            loss = outputs.loss
            loss += self.decoder.lm_head(hidden_states, **inputs)
        return OmniOutput(loss=loss, hidden_states=hidden_states)

# 列表4 OmniModel推理过程
class OmniModel:
    def setup_image_generation(self):
        self.generation_type = "image"
        self.setup_parallelize_generation()
        self.setup_classifier_free_guidance()
        self.setup_position_id_map()

    def setup_video_generation(self):
        ...

    def setup_audio_generation(self):
        ...

    def setup_text_generation(self):
        ...

    def prepare_inputs_for_generation(self, **inputs):
        model_inputs = self.foundation.prepare_inputs_for_generation(**inputs)
        if cache_position[0] == 0:
            inputs_embeds = self.encoder(**inputs)
            model_inputs["inputs_embeds"] = inputs_embeds

        if self.generation_type == "text":
            if input_ids[0][-1] == self.image_start_token:
                self.setup_image_generation()
            if input_ids[0][-1] == self.video_start_token:
                self.setup_video_generation()
            if input_ids[0][-1] == self.audio_start_token:
                self.setup_audio_generation()

        if self.generation_type == "image":
            hidden_states = model_inputs["hidden_states"]
            input_embeds, next_tokens = self.decoder.image_decoder.lm_embed(
                hidden_states, **self.image_generation_configs
            )
            model_inputs["inputs_embeds"] = input_embeds
            self.image_tokens.append(next_tokens)
            if len(self.image_tokens) == self.image_token_num:
                self.setup_text_generation()
                self.images.append(self.decoder.image_decoder.lm_generate(self.image_tokens))

        if self.generation_type == "video":
            ...
        if self.generation_type == "audio":
            ...

        return model_inputs

B.4 并行状态管理

统一声明式初始化。列表5展示了VeOmni的parallel_state API的核心接口，它提供了所有n-D并行配置的统一和声明式初始化。与传统方法【52，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】基于进程组管理并行状态不同，VeOmni利用DeviceMesh【36，TorchTitan: One-stop pytorch native solution for production ready llm pretraining，2025，ICLR】抽象来表示和组织各种并行拓扑，包括数据并行（DP）、张量并行（TP）、专家并行（EP）、流水线并行（PP）和序列并行（SP）。一旦通过init_parallel_state初始化，用户就可以通过全局parallel_state对象方便地管理进程组、rank和通信拓扑。这种设计简化了n-D并行策略的编排，并将底层分布式逻辑与模型实现解耦。

# 列表5 并行状态管理
from veomni.distributed.parallel_state import (
    get_parallel_state, init_parallel_state,
)

init_parallel_state(
    dp_size=data_parallel_size,
    dp_replicate_size=data_parallel_replicate_size,
    dp_shard_size=data_parallel_shard_size,
    tp_size=tensor_parallel_size,
    ep_size=expert_parallel_size,
    pp_size=pipeline_parallel_size,
    ulysses_size=ulysses_parallel_size,
    dp_mode=data_parallel_mode,
)

# 获取全局 parallel_state
parallel_state = get_parallel_state()

# 获取 DP Mesh
dp_mesh = parallel_state.dp_mesh
# 获取 DP Group
dp_group = parallel_state.dp_group
# 获取 DP Shard Rank
dp_shard_rank = parallel_state.dp_shard_rank
# 获取 Ulysses Group
ulysses_group = parallel_state.ulysses_group

B.5 全模态模型并行化

抽象底层细节。列表6展示了build_parallelize_model API的用法，该API将完全分片数据并行（FSDP）和自定义的并行计划应用于目标模型。此接口抽象了底层的分布式训练细节，例如参数分片、精度配置和激活检查点，从而能够以最少的代码更改实现高效的大规模分布式训练，并在不牺牲可用性的情况下实现可扩展性。

# 列表6 全模态模型并行化
from veomni.distributed.torch_parallelism import build_parallelize_model

model = build_parallelize_model(
    model,
    weights_path=model_path,
    enable_full_shard=enable_full_shard,
    enable_mixed_precision=enable_mixed_precision,
    enable_gradient_checkpointing=enable_gradient_checkpointing,
    init_device=init_device,
    enable_fsdp_offload=enable_fsdp_offload,
)

B.6 长上下文注意力

Ulysses风格序列并行支持。列表7展示了VeOmni中flash_attention_forward的实现，它扩展了HuggingFace的原生FlashAttention接口以支持Ulysses风格的序列并行【29，Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models，2023，arXiv】。当启用序列并行时，查询、键和值张量首先通过all-to-all集体操作从序列分片布局转换为头分片布局。在通过标准的_flash_attention_forward函数执行注意力计算后，输出再从头分片格式转换回序列分片格式。此函数保持了与HuggingFace模型的模块化兼容性，无需更改注意力逻辑，并实现了序列并行与FlashAttention【9，FlashAttention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS】的无缝集成。

# 列表7 长上下文注意力
def flash_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    **kwargs,
) -> Tuple[torch.Tensor, None]:
    # Ulysses all-to-all
    ulysses_enabled = get_parallel_state().ulysses_enabled
    if ulysses_enabled:
        ulysses_group = get_parallel_state().ulysses_group
        query = gather_seq_scatter_heads(query, seq_dim=1, head_dim=2, group=ulysses_group)
        key = gather_seq_scatter_heads(key, seq_dim=1, head_dim=2, group=ulysses_group)
        value = gather_seq_scatter_heads(value, seq_dim=1, head_dim=2, group=ulysses_group)

    attn_output: torch.Tensor = _flash_attention_forward(
        query, key, value, attention_mask, **kwargs,
    )

    # Ulysses all-to-all
    if ulysses_enabled:
        ulysses_group = get_parallel_state().ulysses_group
        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2, group=ulysses_group)

    return attn_output, None

B.7 专家并行实现

声明式设计。列表8演示了使用VeOmni的ParallelPlan接口定义专家并行（EP）计划。这种声明式设计允许用户指定哪些模块（例如，MoE块中的专家投影层）应在专家并行组中进行分片。该语法遵循类似于PyTorch模块名称的通配符匹配约定，并利用了与DTensor兼容的放置原语，如Shard(0)。当传递给build_parallelize_model()时，该计划会自动应用于沿专家维度对目标模型进行分区。这种灵活的机制不仅简化了专家并行，还实现了与张量并行（TP）和FSDP等其他并行策略的清晰集成。

# 列表8 专家并行实现
def get_parallel_plan():
    ep_plan = {
        "model.layers.*.mlp.experts.gate_proj": Shard(0),
        "model.layers.*.mlp.experts.up_proj": Shard(0),
        "model.layers.*.mlp.experts.down_proj": Shard(0),
    }
    parallel_plan = ParallelPlan(
        ep_plan=ep_plan,
    )
    return parallel_plan

方法细节中的引用汇总

• 【46】Zero: Memory optimizations toward training trillion parameter models (2020, SC20):

  • 引用位置: 3.2.1节，第一段。
  • 引用内容: 描述FSDP是ZeRO [46] 在PyTorch中的高效实现，其核心优势是通过分片模型状态来大幅减少单卡内存占用。

• 【80】Pytorch fsdp: experiences on scaling fully sharded data parallel (2023, arXiv):

  • 引用位置: 3.2.1节，第二段。
  • 引用内容: 指出VeOmni集成了FSDP1 [80] 作为其分布式训练栈的基础组件之一。

• 【36】TorchTitan: One-stop pytorch native solution for production ready llm pretraining (2025, ICLR):

  • 引用位置: 3.2.1节，第二段；3.2.4节，第三段；3.2.5节，第五段；附录A。
  • 引用内容: 提及VeOmni集成了FSDP2 [36]；采用其全局设备网格（DeviceMesh）抽象来简化n-D并行管理；使用其DTensor格式进行元设备初始化；并在附录中与TorchTitan进行性能比较。

• 【18】Seedance 1.0: Exploring the boundaries of video generation models (2025, arXiv), 【66】Qwen2-vl: Enhancing vision-language model’s perception of the world at any resolution (2024, arXiv), 【78】Video instruction tuning with synthetic data (2024, arXiv):

  • 引用位置: 3.2.2节，第一段。
  • 引用内容: 共同作为例子，说明在高分辨率图像或视频理解与生成等领域，训练所需的序列长度显著增长，带来了计算和内存挑战。

• 【29】Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models (2023, arXiv):

  • 引用位置: 3.2.2节，第二段；附录B.6。
  • 引用内容: 描述VeOmni采用DeepSpeed-Ulysses [29] 作为高效的序列并行技术，通过沿序列维度切分激活和编排all-to-all通信来应对长序列训练。附录中展示了对其的扩展实现。

• 【9】FlashAttention: Fast and memory-efficient exact attention with io-awareness (2022, NeurIPS):

  • 引用位置: 3.2.2节，第三段；3.2.5节，第一段；附录B.6。
  • 引用内容: 提到VeOmni通过集成DeepSpeed-Ulysses增强了FlashAttention [9]的实现；在动态批处理中利用FlashAttention来高效利用批次预算；附录展示了将序列并行与FlashAttention无缝集成的实现。

• 【10】Deepseek-v3 technical report (2025, arXiv):

  • 引用位置: 3.2.3节，第一段和第三段。
  • 引用内容: 引用MoE架构作为高效扩展大模型的主流方案[10]；并提到VeOmni的算子级通信计算重叠方法避免了像DualPipe [10] 那样复杂的流水线级解决方案。

• 【4】Flux: Fast software-based communication overlap on gpus through kernel fusion (2024, arXiv), 【76】Comet: Fine-grained computation-communication overlapping for mixture-of-experts (2025, arXiv):

  • 引用位置: 3.2.3节，第三段；3.2.5节，第二段。
  • 引用内容: 指出VeOmni的细粒度通信-计算重叠技术受到了这些近期进展[4, 76]的启发；并且集成了这些工作中提出的MoE特定优化操作核。

• 【41】Efficient large-scale language model training on gpu clusters using megatron-lm (2021, SC), 【52】Megatron-lm: Training multi-billion parameter language models using model parallelism (2019, arXiv):

  • 引用位置: 3.2.4节，第三段；附录B.4。
  • 引用内容: 与VeOmni基于全局设备网格的抽象进行对比，指出这些方法[41, 52]需要手动管理多个进程组，增加了复杂性。

• 【8】Flashattention-2: Faster attention with better parallelism and work partitioning (2023, arXiv):

  • 引用位置: 3.2.5节，第一段和第二段。
  • 引用内容: 提到VeOmni利用FlashAttention[8]进行动态批处理；并集成了FlashAttention作为高效核函数之一。

• 【23】Liger kernel: Efficient triton kernels for llm training (2025, arXiv):

  • 引用位置: 3.2.5节，第二段。
  • 引用内容: 指出VeOmni集成了一系列来自liger-kernel [23] 的高效算子核。

• 【50】Flashattention-3: Fast and accurate attention with asynchrony and low-precision (2024, NeurIPS):

  • 引用位置: 3.2.5节，第二段。
  • 引用内容: VeOmni集成的高效核函数中包含了FlashAttention [50]。

• 【6】Training deep nets with sublinear memory cost (2016, arXiv):

  • 引用位置: 3.2.5节，第三段。
  • 引用内容: 提及VeOmni采用层级重计算[6]作为内存优化技术之一。

• 【64】Bytecheckpoint: A unified checkpointing system for large foundation model development (2024, arXiv):

  • 引用位置: 3.2.5节，第四段。
  • 引用内容: 描述VeOmni利用ByteCheckpoint [64] 实现高效的分布式检查点功能，并将其扩展以支持全模态模型。

• 【68】Janus: Decoupling visual encoding for unified multimodal understanding and generation (2025, CVPR):

  • 引用位置: 附录B.3。
  • 引用内容: 提到OmniModel在推理时支持无分类器引导[68]。