通过 NVIDIA NeMo 框架的最新功能全链路加速视频基础模型的训练和微调

生成式 AI 已经从基于文本的模型发展到多模态模型，最近还扩展到了视频领域，为各行各业开辟了新的潜在用途。视频模型可以为用户创造新的体验，以及大规模模拟场景来训练自主代理。它们正在帮助包括机器人技术、自动驾驶汽车和娱乐业在内的多个行业实现变革。

由于视频数据的庞大和多样性，视频基础模型的开发面临着独特的挑战。这也凸显了可扩展的数据管理流程和训练框架的必要性，以此来理解时间和空间动态的模型。

我们发布了 NVIDIA NeMo 框架中全新的视频基础模型功能，这是一个端到端的训练框架，用于预训练和微调您的专属视频基础模型。该框架包括高吞吐量的数据管理、高效的多模态数据加载功能、可扩展的模型训练，以及并行化的框架内推理。

图 1. NeMo Curator 将视频数据处理速度提升 89 倍 (图片展示的性能数据来自 NVIDIA 的实验测试，仅供技术讨论和参考)

NeMo Curator 为视频基础模型的训练和微调数据集提供以下相关的功能：

• 剪辑 (Clipping)

• 分片 (Sharding)

剪辑链路从解码和将原始视频分割为短而连续的片段开始，通过分析帧与帧之间的颜色变化来完成分割。在拼接阶段，利用图像嵌入 (embedding) 的相似性来平滑衔接片段，将相邻片段合并在一起。随后，这些片段会被转码为高质量的视频编码格式 (H264)，并通过视频嵌入 (embedding) 和字幕进行标注，这些字幕可以是现有的，也可以是由视觉语言模型 (VLM) 合成生成的，从而增强语义搜索功能。

图 2. 视频整理中的剪辑与分片链路

分片通过为字幕生成文本嵌入 (embedding) 来创建用于训练的最终 WebDataset。NeMo Curator 还利用 Ray 流处理技术构建了一个自动平衡系统，并为链路中的每个阶段部署了最佳数量的工作线程，以避免任何阶段成为瓶颈 (如图 3 所示)。

图 3. 用于匹配整体链路的吞吐量的自动平衡系统

图 4. 通过序列打包加速图像-视频混合训练

• 减少网络带宽压力： 每个模型并行等级 (rank) 只下载数据集的一个不同子集，而不是整个数据集，然后所有等级通过全聚合 (all-gather) 在各卡之间共享数据，从而确保每张卡获得一致的数据加载。

图 5. 减少网络带宽压力以提高训练吞吐量

视频基础模型可以是自回归模型或扩散模型。

NeMo 框架里面现有的大语言模型功能可以复用于上述自回归模型，同时 NeMo 新增了对扩散 Transformer 的支持，如 DiT、MovieGen，以及最新的 NVIDIA Cosmos 世界基础模型（用于物理 AI）。

NeMo 技术栈经过高度优化，在最新基准测试中提供了超过 40% 的模型浮点运算利用率 (Model FLOPs Utilization, MFU)（表 1）。

表 1. NVIDIA NeMo 框架在扩散 Transformer (DiT) 上的 GPU 利用率和吞吐量基准测试

标注: CP=context parallelism; TP=tensor parallelism; SP=sequence parallelism; PP=pipeline parallelism

视频扩散训练链路通常由以下几个主要步骤组成：

• 使用因果时序 3D 分词器对输入的图像和视频进行分词，生成 3D 时空分词。

  
  

• 使用由扩散噪声调度 @ 时间步 t 和文本输入作为条件的 Transformer 解码器。

  
  

  o 时间步条件调节通过 Adaptive LayerNormalization (AdaLN) 机制进行应用，并提供使用 AdaLN-LoRA 的选项，后者在训练过程中进一步提高模型浮点运算利用率 (Model FLOPs Utilization，MFU)。

  
  

  o 文本条件调节通过每个 Transformer 块中的交叉注意力层呈现。

  
  

  o NeMo 框架使您能够基于标准的 DiT 架构或 MovieGen Llama 架构来初始化您的 Transformer 解码器，后者使用分组查询注意力机制 (GQA, Grouped-Query Attention)。

  
  

• 使用扩散 Transformer 的噪声预测，通过并行化的 EDM 扩散链路计算扩散损失。

NeMo 还在每个注意力块之前，对 query 和 key 向量使用额外的均方根层归一化 (Root Mean Square Layer Normalization, RMSNorm) 以稳定扩散训练。RMSNorm 在每个注意力头上应用，以保持与张量并行性的兼容性。

NeMo 和 NVIDIA Megatron-Core 支持多种模型并行技术：

• 张量并行 (Tensor Parallel, TP)

• 序列并行 (Sequence Parallel, SP)

• 流水线并行 (Pipeline Parallel, PP)

• 上下文并行 (Context Parallel, CP)

然而，这些技术在应用于视频扩散 Transformer 时面临独特的挑战。以下是 NeMo 如何解决这些挑战，实现可扩展和高性能训练的方式：

• 用于条件调节的高效流水线并行

• 支持时空 DiT (ST-DiT) 架构

• 定制的随机种子机制

传统方法是通过流水线各阶段之间传递条件信息，这会产生额外的通信开销，并且需要对流水线调度进行较大修改。NeMo 通过在每个流水线阶段即时计算条件嵌入 (embedding) 来解决这个问题。与通信成本相比，条件调节的高效流水线并行计算成本要低得多，从而提高了训练吞吐量。

图 7. 在条件调节流水线并行中，以通信换取计算资源的优化策略

时空 DiT (ST-DiT) 架构在每个 Transformer 块中引入了额外的空间和时间自注意力层，以替代在长视频序列上进行全自注意力训练。因为这些层在输入序列比较短，所以进行计算时计算量较小，以至于暴露出在上下文并行中的通信开销。NeMo 通过使用局部注意力计算和 A2A 通信来处理空间/时间注意力层里的短序列输入，同时保持 P2P 环形拓扑结构用于全自注意力里的长序列输入。该混合方法有效地减少了时间/空间注意力的带宽需求，同时仍然从上下文并行中获益，避免了全自注意力层的计算开销（表 2）。

图 8. 时空 DiT Transformer 架构

表 2. NeMo 针对每种层的通信策略

标识: b=batch size; h*w=spatial size; t=temporal size; cp=context parallel size; d=hidden size, with input size being (b, t*h*w, d).

定制的随机种子机制的目标是确保以下组件中的随机种子得以正确初始化：

• 时间步

• 高斯噪声

• 实际模型权重

表 3 显示了 NeMo 框架里面对于随机种子的初始化策略。

表 3. 针对并行化扩散 Transformer 架构的定制随机种子设置

标识: Diff=Different random seed from other parallel ranks; Same=Same random seed as other parallel ranks.

推理的时候，NeMo 框架通过上下文并行，将去噪操作分布到多个 GPU 上，从而加速推理。在并行去噪后，潜在张量被合并以重建视频序列，然后使用 Cosmos 视频分词器进行解码。

基准测试显示，在最多 32 个 NVIDIA Hopper 架构 GPU 上，NeMo 的扩展效率达到 80-90%。在 1 个和 32 个 GPU 上，FP8 多头注意力相比 BF16 提供了分别 28% 和 48% 的性能提升。

在本文中，我们介绍了 NVIDIA NeMo 框架，帮助您以高效的方式进行视频基础模型的预训练或微调。

NeMo Curator 通过剪辑和分片链路提供高吞吐量的数据整理，Megatron Energon 库则提供高效的多模态数据加载。NeMo 框架通过支持多种针对扩散模型和自回归模型特别优化的模型并行技术，实现了可扩展的视频基础模型训练。此外，它通过将去噪操作分布到多个 GPU 上，并结合 FP8 多头注意力，实现了高效的框架内推理。

您可以通过 NeMo Curator 的早期试用计划处理您的视频数据，对其进行分词，进行预训练（扩散、自回归）、微调（扩散、自回归），并在 NeMo 框架下进行多 GPU 的框架内推理（扩散、自回归）。