Accelerate Custom Video Foundation Model Pipelines with New NVIDIA NeMo Framework Capabilities
Generative AI
has evolved from text-based models to multimodal models, with a recent expansion into video, opening up new potential uses across various industries. Video models can create new experiences for users or simulate scenarios for training autonomous agents at scale. They are helping revolutionize various industries including robotics, autonomous vehicles, and entertainment.
The development of video
foundation models
presents unique challenges due to the vast and varied nature of video data. This also underscores the necessity of scalable pipelines for curating data and effectively training models that can comprehend temporal and spatial dynamics.
We are announcing brand new video foundation model capabilities in the
NVIDIA NeMo framework
, an end-to-end training framework that enables you to pretrain and fine-tune your own video foundation models. The framework includes a high-throughput data curation, efficient multimodal data loading functionality, scalable model training, and a parallelized in-framework inference.
Video 1. NVIDIA Cosmos is a world foundation model platform that accelerates the development of physical AI applications like humanoid robots and autonomous vehicles.
High-throughput video curation through optimized pipelines
NeMo Curator
improves generative AI model accuracy by efficiently processing and preparing high-quality data, including large video datasets.
Using NeMo Curator’s scalable data pipelines, you can efficiently clip, annotate, and filter 100 PB or more of videos. To remove bottlenecks and optimize performance, NeMo Curator uses the following combination:
NVDEC
: Hardware decoder
NVENC
: Hardware encoder
Ray
: Compute framework for scaling AI applications
The NeMo Curator autobalancing techniques can leverage heterogeneous clusters with multiple GPU types to take advantage of NVENC on L40S GPUs and the performance of H100 and GB200 GPUs.
Figure 1 shows how NeMo Curator can process 20M hours of video data, reducing the processing time from years to days, achieving 89x speed up using 1K GPUs compared to unoptimized pipelines on CPUs for ISO power usage.
Figure 1. NeMo Curator delivers 89x faster video data processing
NeMo Curator provides the following relevant pipelines for video foundation model training and fine-tuning datasets:
Clipping
Sharding
The clipping pipeline starts with decoding and splitting raw videos into short, continuous clips by analyzing frame-to-frame color changes. The stitching stage smooths the clips out by using image embedding similarities to potentially merge adjacent clips together. These clips are then transcoded to the high-quality video encoding (H264), and they are annotated with video embeddings and captions, either existing or synthetically generated by a VLM, to facilitate semantic search capabilities.
Figure 2. Video curation clipping and sharding pipelines
Sharding generates text embeddings for captions to create the final WebDataset used for training. NeMo Curator also uses Ray streaming to build an auto-balancing system and deploy an optimal number of workers for each stage in the pipeline to avoid being bottlenecked by any stage (Figure 3).
Figure 3. Auto-balancing system to match the throughput of the overall pipeline
Efficient multimodal dataloading
Video models can be trained on billions of images and millions of videos, necessitating an efficient data loading strategy to achieve high throughput during training time.
This is accomplished in the NeMo framework through the use of
Megatron-Energon data loader
:
Shard large-scale data:
Uses the WebDataset format to shard a TB-size dataset into compressed files to help reduce I/O overhead during training.
Deterministic save and load:
Enables the dataset to be visited in one pass without repetition when the training job is disrupted, ensuring consistency across different training cluster setups.
Sequence packing:
Packs variable length or resolution images and videos together up to the max sequence length, minimizing compute wastage due to padding while simplifying data loading logic. NeMo uses the special THD attention kernel from the
Transformer engine
to support accelerated training with sequence packing.
Figure 4. Mixed image-video training with sequence packing
Reduce network bandwidth strain:
Each model parallel rank downloads a different subset of data instead of the whole dataset, and then all-gathers the data across ranks to get an identical dataloader.
Figure 5. Reducing network bandwidth strain to improve training throughput
Scaling video foundation model training
Video foundation models can be either autoregressive or diffusion models.
The well-established suite of NeMo tools on
large language models
(LLMs) can be reused for autoregressive models, while support for diffusion transformers such as DiT, MovieGen, and the latest
NVIDIA Cosmos
world foundation models for physical AI have been newly added.
The NeMo tech stack is highly optimized and provides more than 40%
Model FLOPs utilization
(MFU) in the latest benchmark (Table 1).
Model size
Context length
Training config
GPU used (TFLOPS/s)
Throughput (token/s/GPU)
DiT 7B
8k
baseline, no optimization
OOM
DiT 7B
8k
CP=2
457
8,969
DiT 7B
74k
TP=4 SP CP=4
414
2,933
DiT 28B
8k
TP=2 SP PP=2
435
2,392
DiT 28B
74k
TP=8 SP CP=4 PP=4
411
994
Table 1. GPU utilization and throughput benchmark for NVIDIA NeMo framework on diffusion transformers (DiT)
Legend:
CP
=context parallelism;
TP
=tensor parallelism;
SP
=sequence parallelism;
PP
=pipeline parallelism
Overview of the video diffusion pipeline
A video diffusion training pipeline is generally composed of the following major steps:
Tokenize the input image and video
with a causal temporal 3D tokenizer to generate 3D spatio-temporal tokens.
Use a transformer decoder
conditioned by the diffusion noise schedule timestep t and text input.
Timestep conditioning is applied through an Adaptive LayerNormalization (AdaLN) mechanism, with an option to use AdaLN-LoRA, which further improves Model FLOPs Utilization (MFU) during training.
Text conditioning is applied through a cross attention layer in each transformer block.
The NeMo framework enables you to initialize your transformer decoder based on the canonical DiT architecture or the MovieGen Llama architecture, which uses Grouped-Query Attention (GQA).
Compute the diffusion loss
with the parallelized EDM diffusion pipeline using the noise prediction from the diffusion transformer.
NeMo also applies additional Root Mean Square Layer Normalization (RMSNorm) on the queries and keys before attention blocks to stabilize diffusion training. RMSNorm is applied per attention head to remain compatible with tensor parallelism.
Figure 6. NeMo video diffusion training pipeline
Parallelism optimizations for video diffusion models
NeMo and Megatron-Core enable various model parallelism techniques:
Tensor parallel (TP)
Sequence parallel (SP)
Pipeline parallel (PP)
Context parallel (CP)
However, these techniques face unique challenges when applied to video diffusion transformers. Here’s how NeMo solves these challenges to achieve scalable and performant training:
Efficient pipeline parallelism for conditioning
Support for Spatio-Temporal DiT (ST-DiT) architecture
Customized random seeding mechanism
The traditional approach is to communicate conditioning information across pipeline stages, incurring additional communication cost and requiring nontrivial modifications to the pipeline schedule. NeMo solves this problem by computing the conditional embeddings at each pipeline stage. The computation cost with efficient pipeline parallelism for conditioning is much less than the communication cost and improves training throughput.
Figure 7. Trading communication for compute in conditioning pipeline parallelism
The Spatio-Temporal DiT (ST-DiT) architecture introduces additional spatial and temporal self-attention layers to each transformer block, as an alternative to training with full self attention on long video sequences. This approach exposes communication overhead during context parallelism due to smaller compute over short input sequence for these layers. NeMo addresses this by using local attention computation with A2A communication for spatial/temporal attention, while maintaining P2P ring topology for full self-attention. The hybrid approach effectively reduces bandwidth needs for temporal/spatial attention while still benefiting from context parallelism over full self-attention layer (Table 2).
Figure 8. Spatial-temporal DiT transformer block
Layer
Input Seq
Communication primitive
Communication bandwidth
Temporal self-attention
Short seq
Local compute & A2A
(bhw/cp, t, d)
Spatial self-attention
Short seq
Local compute & A2A
(bt/cp, hw, d)
Full attention
Long seq
CP with P2P
(b, h*w*t/cp, d)
Table 2. NeMo communication strategies for each kind of layer
Legend:
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).
The customized random seeding mechanism goal is to make sure that random seeds are correctly initialized across the following components:
Time step
Gaussian noise
The actual model weights
Table 3 shows NeMo’s initialization strategy.
RNG seed
Data parallel
Context parallel
Pipeline parallel
Tensor parallel
Time step (t)
Diff
Same
Same
Same
Gaussian noise
Diff
Diff
Same
Same
Weight initialization
Same
Same
Diff
Diff
Table 3. Customized random seeding for parallelized diffusion transformers
Legend:
Diff=
Different random seed from other parallel ranks;
Same
=Same random seed as other parallel ranks.
Efficient in-framework inference
The NeMo framework accelerates inference by distributing denoising operations across multiple GPUs through context parallelism. After parallel denoising, the latent tensors are combined to reconstruct the video sequence before decoding with the Cosmos video tokenizer.
Benchmarks show 80–90% scaling efficiency on up to 32 H100 GPUs, with FP8 Multi-Head Attention providing 28% and 48% performance improvements over BF16 on 1 and 32 GPUs respectively.
Figure 9. Parallelized video generation with context parallelism
Figure 10. Inference performance at different GPU count​s
Conclusion
In this post, we covered all the features of NVIDIA NeMo framework that will help you pretrain or fine-tune video foundation models in an effective and efficient manner.
NeMo Curator offers high-throughput data curation through clipping and sharding pipelines, and the Megatron Energon library offers efficient multimodal data loading. NeMo Frameworks enables scalable video foundation model training by supporting various model parallelism techniques specially optimized on diffusion and autoregressive models. In addition, it provides efficient in-framework inference by distributing denoising operations across multiple GPUs and incorporating FP8 Multi-Head Attention.
You can curate your video data with
NeMo Curator early access program
,
tokenize
them, pre-train(
diffusion
,
autoregressive
), fine-tune (
diffusion
,
autoregressive
), and perform multi-GPU in-framework inference (
diffusion
,
autoregressive
)with NeMo Framework  today.
You can also try the NVIDIA Cosmos world foundation models at
build.nvidia.com
and watch the
CES keynote
from NVIDIA CEO Jensen Huang to learn more about the NVIDIA Cosmos world foundation model platform
.
Acknowledgements
Thanks to the following contributors: 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 NeMo 框架的最新功能全链路加速视频基础模型的训练和微调
生成式 AI
已经从基于文本的模型发展到多模态模型，最近还扩展到了视频领域，为各行各业开辟了新的潜在用途。视频模型可以为用户创造新的体验，以及大规模模拟场景来训练自主代理。它们正在帮助包括机器人技术、自动驾驶汽车和娱乐业在内的多个行业实现变革。
由于视频数据的庞大和多样性，视频
基础模型
的开发面临着独特的挑战。这也凸显了可扩展的数据管理流程和训练框架的必要性，以此来理解时间和空间动态的模型。
我们发布了
NVIDIA NeMo 框架
中全新的视频基础模型功能，这是一个端到端的训练框架，用于预训练和微调您的专属视频基础模型。该框架包括高吞吐量的数据管理、高效的多模态数据加载功能、可扩展的模型训练，以及并行化的框架内推理。
高吞吐量视频数据处理链路
NVIDIA NeMo Curator 通过高效处理和准备高质量的数据（包括大型视频数据集），来提高生成式 AI 模型的准确性。
使用 NeMo Curator 的可扩展数据流程，您可以高效地剪辑、标注和过滤 100 PB 或更多的视频数据。为了消除瓶颈并优化性能，NeMo Curator 使用以下组合：
NVDEC：硬件解码器
NVENC：硬件编码器
Ray：用于扩展 AI 应用的计算框架
根据实验测试数据（仅做技术讨论和参考）显示，NeMo Curator 的自动负载平衡技术可以利用具有多种 GPU 类型的异构集群，以充分利用 L40S GPU 上的 NVENC 以及 H100 和 GB200 GPU 的性能。
图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 框架里面现有的
大语言模型
（LLMs）功能可以复用于上述自回归模型，同时 NeMo 新增了对扩散 Transformer 的支持，如 DiT、MovieGen，以及最新的
NVIDIA Cosmos
世界基础模型（用于物理 AI）。
NeMo 技术栈经过高度优化，在最新基准测试中提供了超过 40% 的
模型浮点运算
利用率
（Model FLOPs Utilization，MFU）
（表 1）。
Model size
Context length
Training config
GPU used (TFLOPS/s)
Throughput (token/s/GPU)
DiT 7B
8k
baseline, no optimization
OOM
DiT 7B
8k
CP=2
457
8,969
DiT 7B
74k
TP=4 SP CP=4
414
2,933
DiT 28B
8k
TP=2 SP PP=2
435
2,392
DiT 28B
74k
TP=8 SP CP=4 PP=4
411
994
表 1. NVIDIA NeMo 框架在扩散 Transformer（DiT）上的 GPU 利用率和吞吐量基准测试
标注:
CP
=context parallelism;
TP
=tensor parallelism;
SP
=sequence parallelism;
PP
=pipeline parallelism
视频扩散模型链路概述
视频扩散训练链路通常由以下几个主要步骤组成：
使用因果时序 3D 分词器对输入的图像和视频进行分词，生成 3D 时空分词。
使用由扩散噪声调度@时间步 t 和文本输入作为条件的 Transformer 解码器。
时间步条件调节通过 Adaptive LayerNormalization（AdaLN）机制进行应用，并提供使用 AdaLN-LoRA 的选项，后者在训练过程中进一步提高模型浮点运算利用率（Model FLOPs Utilization，MFU）。
文本条件调节通过每个 Transformer 块中的交叉注意力层呈现。
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 架构
Layer
Input Seq
Communication primitive
Communication bandwidth
Temporal self-attention
Short seq
Local compute & A2A
(bhw/cp, t, d)
Spatial self-attention
Short seq
Local compute & A2A
(bt/cp, hw, d)
Full attention
Long seq
CP with P2P
(b, h*w*t/cp, d)
表 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 框架里面对于随机种子的初始化策略。
RNG seed
Data parallel
Context parallel
Pipeline parallel
Tensor parallel
Time step (t)
Diff
Same
Same
Same
Gaussian noise
Diff
Diff
Same
Same
Weight initialization
Same
Same
Diff
Diff
表 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 世界基础模型，并观看 NVIDIA CEO 黄仁勋在
CES 上的主题演讲
，了解更多关于 NVIDIA Cosmos 世界基础模型平台的信息。
鸣谢
感谢以下贡献者: 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.