摘要：生成式 AI 已经从基于文本的模型发展到多模态模型，最近还扩展到了视频领域，为各行各业开辟了新的潜在用途。视频模型可以为用户创造新的体验，以及大规模模拟场景来训练自主代理。它们正在帮助包括机器人技术、自动驾驶汽车和娱乐业在内的多个行业实现变革。
生成式 AI 已经从基于文本的模型发展到多模态模型，最近还扩展到了视频领域，为各行各业开辟了新的潜在用途。视频模型可以为用户创造新的体验，以及大规模模拟场景来训练自主代理。它们正在帮助包括机器人技术、自动驾驶汽车和娱乐业在内的多个行业实现变革。
由于视频数据的庞大和多样性，视频基础模型的开发面临着独特的挑战。这也凸显了可扩展的数据管理流程和训练框架的必要性，以此来理解时间和空间动态的模型。
我们发布了 NVIDIA NeMo 框架中全新的视频基础模型功能，这是一个端到端的训练框架，用于预训练和微调您的专属视频基础模型。该框架包括高吞吐量的数据管理、高效的多模态数据加载功能、可扩展的模型训练，以及并行化的框架内推理。
高吞吐量视频数据处理链路
NVIDIA NeMo Curator 通过高效处理和准备高质量的数据（包括大型视频数据集），来提高生成式 AI 模型的准确性。
使用 NeMo Curator 的可扩展数据流程，您可以高效地剪辑、标注和过滤 100 PB 或更多的视频数据。为了消除瓶颈并优化性能，NeMo Curator 使用以下组合：
NVDEC：硬件解码器NVENC：硬件编码器Ray：用于扩展 AI 应用的计算框架NeMo Curator 的自动负载平衡技术可以利用具有多种 GPU 类型的异构集群，以充分利用 Ada 架构 GPU 上的 NVENC 以及 Hopper 和 Blackwell 架构优异的性能。
图 1 展示了 NeMo Curator 如何处理 2,000 万小时的视频数据，将处理时间从数年缩短到数天，使用 1,000 个 GPU 相比于 CPU 上未优化的流程，在相同功耗下实现了 89 倍的加速。
图 1. NeMo Curator 将视频数据处理速度提升 89 倍 (图片展示的性能数据来自 NVIDIA 的实验测试，仅供技术讨论和参考)
NeMo Curator 为视频基础模型的训练和微调数据集提供以下相关的功能：
剪辑 (Clipping)分片 (Sharding)剪辑链路从解码和将原始视频分割为短而连续的片段开始，通过分析帧与帧之间的颜色变化来完成分割。在拼接阶段，利用图像嵌入 (embedding) 的相似性来平滑衔接片段，将相邻片段合并在一起。随后，这些片段会被转码为高质量的视频编码格式 (H264)，并通过视频嵌入 (embedding) 和字幕进行标注，这些字幕可以是现有的，也可以是由视觉语言模型 (VLM) 合成生成的，从而增强语义搜索功能。
图 2. 视频整理中的剪辑与分片链路
分片通过为字幕生成文本嵌入 (embedding) 来创建用于训练的最终 WebDataset。NeMo Curator 还利用 Ray 流处理技术构建了一个自动平衡系统，并为链路中的每个阶段部署了最佳数量的工作线程，以避免任何阶段成为瓶颈 (如图 3 所示)。
图 3. 用于匹配整体链路的吞吐量的自动平衡系统
高效的多模态数据加载
视频模型可以在数十亿张图像和数百万个视频上进行训练，这需要一种高效的数据加载策略，从而在训练期间实现高吞吐量。
在 NeMo 框架中，这通过使用 NVIDIA Megatron-Energon 数据加载器实现：
将大规模数据分片：采用 WebDataset 格式，将 TB 级数据集分片为压缩文件，以减少训练期间的 I/O 开销。保存与加载一致性：在训练任务中断时，确保数据集在一次遍历中不重复访问，从而在不同的训练集群配置中保持一致性。序列打包：将长度或分辨率可变的图像和视频打包在一起，直至达到最大序列长度，减少因填充而导致的计算资源浪费，同时简化数据加载逻辑。NeMo 使用 Transformer Engine 中的特殊 THD 注意力算子核 (attention kernel)，通过结合序列打包来加速训练。图 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 在每个注意力头上应用，以保持与张量并行性的兼容性。
图 6. NeMo 视频扩散模型训练链路
视频扩散模型的并行优化
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% 的性能提升。
图 9. 使用上下文并行化的视频生成
图 10. 在不同 GPU 数量下的推理性能
结论
在本文中，我们介绍了 NVIDIA NeMo 框架，帮助您以高效的方式进行视频基础模型的预训练或微调。
NeMo Curator 通过剪辑和分片链路提供高吞吐量的数据整理，Megatron Energon 库则提供高效的多模态数据加载。NeMo 框架通过支持多种针对扩散模型和自回归模型特别优化的模型并行技术，实现了可扩展的视频基础模型训练。此外，它通过将去噪操作分布到多个 GPU 上，并结合 FP8 多头注意力，实现了高效的框架内推理。
您可以通过 NeMo Curator 的早期试用计划处理您的视频数据，对其进行分词，进行预训练（扩散、自回归）、微调（扩散、自回归），并在 NeMo 框架下进行多 GPU 的框架内推理（扩散、自回归）。
您还可以访问 NVIDIA 官网尝试 NVIDIA Cosmos 世界基础模型 (https://www.nvidia.cn/ai/) ，并观看 NVIDIA CEO 黄仁勋在 CES 上的主题演讲，了解更多关于 NVIDIA Cosmos 世界基础模型平台的信息：
• 扩散预训练 GitHub:
• 自回归预训练 GitHub:
• 扩散微调 GitHub:
• 自回归微调 GitHub:
鸣谢
感谢以下贡献者：Parth Mannan, Xiaowei Ren, Zhuoyao Wang, Carl Wang, Jack Chang, Sahil Jain, Shanmugam Ramasamy, Joseph Jennings, Ekaterina Sirazitdinova, Oleg Sudakov, Linnan Wang, Mingyuan Ma, Bobby Chen, Forrest Lin, Hao Wang, Vasanth Rao Naik Sabavat, Sriharsha Niverty, Rong Ou, Pallab Bhattacharya, David Page, Jacob Huffman, Tommy Huang, Nima Tajbakhsh, and Ashwath Aithal.
关于作者
来源：NVIDIA英伟达中国
