背景与动机
近年来，AI 大模型正从单一的文本模态扩展到图像、语音、视频等“全模态”（Omni-Modal）的理解与生成。但要训练一个能“看”、能“听”、能“说”的通用AI模型，工程上面临诸多挑战。传统的大模型分布式训练框架（如 Megatron-LM 等）多为**“以系统为中心”设计，将模型实现与并行策略紧密耦合。这类框架在训练结构规整的纯文本模型时表现良好，但面对多模态模型纷繁复杂的不均衡计算负载时，常出现负载不均、扩展性受限**等问题。而且，当尝试引入新模态或调整模型结构时，往往需要深入底层重写大量分布式代码，工程代价高昂。这使得算法研究与工程实现紧密绑死，研究人员难以快速验证新想法。
为了解决上述痛点，字节跳动旗下 Seed 实验室联合火山引擎团队研发并开源了全模态 PyTorch 原生训练框架——VeOmni。VeOmni 引入了**“以模型为中心”的分布式训练方案**，将复杂的分布式并行逻辑与模型计算彻底解耦。研究者可以像搭积木一样灵活组合并行策略来高效训练全模态模型，大幅降低工程开发工作量。据介绍，以往使用传统框架训练一个全新架构的视觉-语言模型可能需要一周以上的工程实现，而使用 VeOmni 只需约一天就能完成模型构建并开始训练，工程耗时压缩了90%以上。同时，VeOmni 在性能上也表现出色：例如，在128张GPU上训练一个约300亿参数的多模态 MoE 模型时，吞吐量达到每卡每秒 2800+ tokens，并可轻松支持 160K 长度的超长序列窗口。这些显著提升展示了 VeOmni 在全模态大模型训练上的潜力。
架构设计与核心理念
VeOmni 的核心设计理念是模型中心化的分布式训练。与以往系统中心化的方案不同，VeOmni 将模型实现与底层分布式训练代码解耦，在保证通用性的前提下再追求高性能，让系统去适配快速演进的模型，而不是限制模型去迁就固定的系统。框架内部将所有分布式通信逻辑从模型代码中抽离，由框架统一处理，模型各部分则尽可能保持纯粹的“计算模块”。这种设计在通用性和高性能之间寻求平衡：由于全模态模型架构变化迅速，VeOmni 优先确保框架的灵活可扩展，支持研究者快速尝试新结构，然后通过插件化的优化模块来提升性能。
VeOmni 架构（右）将通信并行逻辑从模型代码中剥离，对比传统框架（左）实现了模型-系统解耦。左图：(a) 传统多模态训练框架中，模型计算（Comp. Ops）与通信操作（Comm. Ops）交织耦合；右图：(b) VeOmni 中，各子模型均为纯计算模块，所有分布式并行逻辑由框架层的 Parallel Plan 统一调度。
从上图可以看到，传统方案下分布式通信往往散落在模型实现的各处，而 VeOmni 通过引入Parallel Plan并行计划，将 FSDP、HSDP、EP、SP 等并行“乐高块”在框架层组合应用于不同组件，模型代码不需关心通信细节。这种模型-系统解耦使得 VeOmni 能够同时实现 “三个统一” 的核心目标：统一多模态支持、统一混合并行策略、统一算力底座。简单来说，VeOmni 提供统一的接口支持任意模态模型，整合多种并行策略于一个框架，并兼容不同加速硬件和集群环境，使得在各种算力平台上都能方便地横向扩展大模型训练。
全模态统一接口
为使新增模态变得轻而易举，VeOmni 设计了一套轻量级的全模态模型接口。框架允许将多模态的编码器（Encoder）和解码器（Decoder）任意组合，灵活地连接到基础模型（大型语言模型）的输入端和输出端。具体而言，VeOmni 对模型结构进行了模块划分：底层是文本等基础模态的主干模型（如 LLaMA、Qwen 等大语言模型），在其输入前和输出后可以插入不同模态的编码器和解码器。例如，可以在语言模型前接入图像编码器、音频编码器，在模型后接入图像解码器、视频解码器等，实现任意模态到任意模态的统一建模。
VeOmni 的全模态模型架构示意：中间的基础模型（Foundation）为主干大语言模型（如 Qwen、LLaMA、DeepSeek 等），左侧可以接入多个模态的 Encoder 编码器（如视觉编码器NaViT、音频编码器Whisper等），右侧接出 Decoder 解码模块（如用于重建图像的VQVAE等）。通过在 HuggingFace PretrainedModel 基类上混入 EncoderMixin 和 DecoderMixin，所有编码器和解码器都实现统一的函数接口（如 lm_encode, lm_embed, lm_generate 等），从而与主干模型无缝衔接。
在 VeOmni 中，每种模态的编码器/解码器都只需遵循统一的接口规范，即按照 HuggingFace 的模型格式，实现诸如 lm_encode（提取表示）、lm_generate（生成输出）等方法。框架提供了抽象基类和Mixin混入机制，开发者可以很方便地将现有模型改造成符合接口的编码器或解码器模块。接入新模态时，无需修改核心训练代码，只需实现少量接口函数即可作为“插件”挂接到系统中。这个设计让研究者可以快速扩展任意模态能力，或者灵活替换各个子模块的模型结构，而不影响整体训练流程。例如，用户可以在不改动训练脚本的情况下，将图像编码器从 CLIP 换成 ImageBind，将语言模型从 LLaMA 换成自研模型，只要新模型实现了所需的接口函数，就能立即投入分布式训练。相比传统框架需要为每种新模态手动处理数据张量拼接和并行通信，VeOmni 的统一接口大大降低了多模态模型研发的门槛。
异构并行策略与调度
全模态模型的另一个难点在于异构组件的分布式训练优化。不同子模块（如视觉编码器 vs. 文本模型 vs. MoE专家层）的计算特性差异很大，单一的并行策略往往不能对所有部分都达到最佳效果。针对这一点，VeOmni 强调**“可组合性”的 n 维并行**：将数据并行、张量并行、流水线并行、序列并行、专家并行等多种并行策略以模块化方式组合，按照模型各部分的需要各取所长。
VeOmni 基于 PyTorch 的 DeviceMesh 实现了一个统一的分布式抽象层 Parallel State 来管理所有并行维度。通过 init_parallel_state() 接口，用户可以在配置中一键初始化所需的并行维度大小，例如 data_parallel_size（数据并行总进程数）、data_parallel_shard_size（启用混合并行时每节点FSDP分片数）、expert_parallel_size（专家并行切分数）、ulysses_parallel_size（序列并行段数）等等。框架会自动根据这些参数创建底层的进程分组和通信组，屏蔽繁琐的手工 process group 管理细节。在训练脚本中，开发者通过 get_parallel_state() 获取这个并行状态对象，即可方便地查询各维度通信组或当前rank等信息，进行必要的逻辑处理。
得益于这种统一抽象，VeOmni 可以灵活实现混合并行的自由组合。例如，可以针对视觉编码器部分使用完整 FSDP（Fully Sharded Data Parallel，全参数分片数据并行）来兼顾大模型的显存开销；同时对语言模型主干的注意力层采用 HSDP+SP（Hybrid Shard Data Parallel + Sequence Parallel）以减小跨节点通信、支持长序列；对于稀疏专家 (MoE) 层则进一步应用 FSDP+EP+SP（三维并行：数据+专家+序列）来充分加速。这些策略组合完全由配置驱动，像搭积木一样叠加，不需要改动任何模型代码。下面列举VeOmni支持或实现的一些并行策略亮点：
完全/混合数据并行 (FSDP/HSDP)：VeOmni 内置了 PyTorch 原生的 FSDP 实现，以非侵入式的全参数分片并行来应对结构复杂的多模态模型训练。这与框架的解耦理念高度一致。当集群规模更大时，用户还可以选择HSDP模式：在节点内继续使用FSDP分片，而节点间采用标准数据并行（DDP）。这种两级混合并行通过减少昂贵的跨节点通信显著提升了扩展效率。在 VeOmni 中，只需在配置中将
data_parallel_shard_size参数由1改为比如2，就切换到了 HSDP 模式，使用起来和 FSDP 几乎无差别。序列并行 (SP)：针对超长序列训练场景，VeOmni 集成了 DeepSpeed 提出的 Ulysses 序列并行技术，并进一步实现了 Async-Ulysses 异步序列并行。其核心是在多卡之间对长序列进行切分，并将注意力计算中的部分线性投影与 All-to-All 通信重叠执行，从而并行化长序列处理，消除通信瓶颈。在实现上，VeOmni 提供了一个简单的
flash_attention_forward接口封装序列并行逻辑，模型层无需感知底层的通信操作即可利用SP。通过SP技术，VeOmni 成功支持了十万级长度的序列训练：例如在 7B 参数的多模态模型上验证可支持长达 192K 的序列长度，在更大的72B模型上也达到96K长度，为高分辨率图像、长视频等模态的训练提供了保障。专家并行 (EP)：对于包含稀疏专家层的 Mixture-of-Experts 模型，VeOmni 基于 PyTorch DTensor 实现了简洁的ParallelPlan接口来支持专家并行。开发者只需在模型配置中指出需要切分的参数维度，并在运行时设置
expert_parallel_size等于专家数，框架即可自动将每个专家参数块分配到不同设备进行并行训练。相比手动实现复杂的路由通信，VeOmni 的专家并行接口大幅降低了实现 MoE 模型并行的难度。通信优化 (COMET)：在 MoE 专家模型中，不可避免地存在大量 All-to-All 通信开销（将 Token 路由给不同设备上的专家）。业内一种传统解决方案是构建复杂的流水线并行（如 DualPipe）来掩盖通信延迟，但这种方法在负载不均的多模态场景下非常僵硬，容易产生“Bubble”气泡（设备空转）。VeOmni 则采用了 Seed 团队研发的COMET细粒度计算-通信重叠技术，它不依赖模型结构的规整程度，在多变的全模态 MoE 训练中更为灵活高效。COMET 能最大限度地降低 All-to-All 过程中的等待浪费，提升设备资源利用率。
其他系统优化：除了并行策略，VeOmni 在系统层面还做了多方面优化来提升训练效率和稳定性。例如，框架支持动态批处理（Dynamic Batching）策略，根据当前输入自动调整 batch size 以充分利用算力；集成高效算子（包括 LinkedIn 开源的 Liger Kernel 等）以加速稀疏计算；支持重计算检查点和显存优化来降低显存占用；并提供了高效的分布式检查点系统ByteCheckpoint，用于快速保存和合并大模型的断点数据。所有这些优化组合，使得 VeOmni 能够在保证灵活性的同时，最大程度发挥硬件性能。
插件化模块与扩展设计
VeOmni 从架构上秉承插件化和模块化设计思想，不仅体现在模型接口层面，更贯穿了训练流程的各个部分。首先，如前文所述，各模态的编码器/解码器可视作独立插件，通过注册接口融入框架；而模型注册表机制确保了框架能自动识别已支持的模型类型，对于未注册的模型则回退调用其 Transformers 实现。这意味着如果用户想加入一个新模型，除了按需实现必要的接口函数外，只需在注册表中登记一下模型配置，VeOmni 就会自动加载对应的模型结构和优化器，无需繁琐地修改框架源码。
其次，VeOmni 放弃了繁重的“Trainer类”封装（不同于 PyTorch Lightning 或 HuggingFace Trainer 的高度封装风格），而是采用纯脚本化的训练流程。提供的训练脚本（如 train_torch.py）以清晰的顺序展示了数据加载、模型构建、前向计算、反向优化等过程，用户可以很容易地理解并进行定制。如果默认的 train_torch.py 不能满足特殊需求，用户可以拷贝一份脚本到自定义目录进行修改，就像插拔一个新“任务插件”一样。这种设计降低了定制难度，使研究人员能够按需替换任意环节：比如替换优化器为自定义算法、更改数据加载逻辑为特殊格式、插入自定义的监控钩子等等。VeOmni 还提供了统一的参数管理系统，通过 Python dataclass 定义配置项，并支持命令行覆盖配置功能。这种机制也可以被视为框架的配置插件系统——默认配置可以扩展新的字段，自定义参数会自动融合，不会破坏原有功能。
总的来说，VeOmni 的插件化体现在：模型模块可插拔、并行策略可插拔、训练流程可插拔、配置参数可扩展。这种高度模块化的框架设计，使其既能满足不同项目的定制需求，又能随着新技术的发展不断扩充能力，而不会影响现有组件的稳定性。
实现亮点
VeOmni 在技术实现上有诸多亮点，以下从模型能力和系统优化两个方面进行总结：
任意模态模型融合：借助统一接口和模块化并行方案，VeOmni 真正实现了 “Any-to-Any” 的模型融合能力。开发者可以方便地构建支持多种输入输出模态的大模型。例如，利用 VeOmni，团队成功训练了一个支持文本、图像、语音、视频综合理解与生成的30B参数全模态MoE模型。在包含视觉编码器（如 Qwen2.5-Omni 提供的 NaViT）、音频编码器（Whisper 模型）、图像解码器（MoVQGAN）等组件的复杂架构下，VeOmni 依然保持了高效的分布式训练：该模型在128卡集群上不仅适配了160K长度的超长序列，还取得了每卡每秒2800+ token的吞吐量，证明了框架对大规模多模态融合的卓越支持。
资源管理与系统优化：VeOmni 在大规模训练的资源管理上做了大量细致优化。通过动态批大小调整策略，框架能够根据不同模态的数据特性最大化 GPU 利用率；借助梯度重计算（Activation Checkpointing）和显存异构映射（如将部分缓存搬移至CPU或NVMe），有效降低了长序列和大模型训练的显存占用压力；在通信方面，前述 COMET 技术将通信开销隐匿于计算之中，加上灵活的并行组合，极大减少了集群“空转”浪费。VeOmni 还特别关注了训练稳定性：在内部实验中，使用 VeOmni 的“即插即用”接口构建了多种架构各异的全模态模型，并在包含文本、图像、音频、视频的复杂任务集上训练，所有模型均稳定收敛，展现了框架的鲁棒性。在检查点管理方面，VeOmni 集成了高性能的 ByteCheckpoint 库，支持将分布式权重快速合并回 HuggingFace 格式，方便训练后模型的保存与推理部署