生成式 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.
