点击下方卡片，关注「集智书童」公众号
导读
视频基础模型（VFMs）近期被用于模拟现实世界，以训练物理AI系统和开发创意视觉体验。然而，训练大规模、高质量的VFMs以生成高质量视频存在重大挑战。作者提出了一种可扩展的开源VFM训练流程，使用NVIDIA NeMo，提供加速的视频数据集整理、多模态数据加载以及并行化的视频扩散模型训练和推理。作者还提供了一份全面的表现分析，突出了高效VFM训练和推理的最佳实践。
- 引言
生成式AI的演进已超越基于文本的模型，进入了新的多模态模型阶段。这些模型现在可以处理图像生成、图像标题和视觉问答等任务，标志着向能够跨越各种模态与人类交互的AI系统的迈进。最近，研究重点已从文本和图像转向视频，解锁了各行业的新潜在应用。
视频基础模型（VFMs）将革新多个行业，包括机器人、自动驾驶汽车和娱乐业。由于视频数据的规模和多样性，构建VFMs面临着独特的挑战。训练VFMs需要可扩展的训练 Pipeline ，能够处理大量数据并训练出能够理解时空动态以模拟世界的模型。
本文提出了一种使用NVIDIA NeMo进行可扩展VFM训练的框架。作者的端到端训练框架允许用户创建或微调自己的VFMs。作者提供NeMo Curator的高吞吐量数据整理，通过Megatron Energon实现高效的跨模态数据加载功能，利用Megatron Core进行可扩展的扩散模型训练，并在NeMo中实现并行化的视频生成 Pipeline 。
- 高吞吐量视频整理与NeMo Curator
有效训练视频帧级模型（VFMs）需要互联网规模的高质量视频数据。NeMo Curator 允许用户高效地对
的视频进行剪辑、标注和筛选。为此，NeMo Curator 提供了两个模块化 Pipeline ，可以处理原始视频并输出高质量的 VFM 训练/微调数据集。每个 Pipeline 由一系列阶段组成，每个阶段在数据整理过程中执行一个步骤。这些 Pipeline 的模块化特性使用户能够使用自定义模型来定制数据整理过程中的不同部分。
2.1. 剪辑 Pipeline
剪辑 Pipeline 以原始、未经编辑的视频为输入，输出短小、连续（无跳切或其他过渡）的剪辑，并附带相关元数据。它采用一种激进的剪辑方法，通过分析帧之间的颜色变化来分割剪辑，这些变化通过计算相邻剪辑图像嵌入的相似性来平滑，以潜在地合并它们。然后，这些剪辑被转码为高质量的视频编码（H264），并使用视频嵌入进行标注，以方便语义搜索功能。当没有提供视频字幕时，作者可以从VLM创建合成字幕，这已被证明比预先编写的字幕提高了下游模型的质量。在此 Pipeline 之后，建议用户手动检查生成的字幕，以验证其质量并根据需要做出任何修改。
2.2. 分片 Pipeline
分片技术为字幕生成文本嵌入，并创建用于训练的最终WebDataset。WebDataset将大型数据集划分为更小、更易于管理的分片。这种划分促进了并行数据访问和处理，允许多个GPU同时处理不同的分片。这在训练具有PB级多模态数据的VFM并在数千个GPU上进行训练时尤为重要。分片 Pipeline 将数据分割成不同的POSIX tar存档文件。在训练过程中，这些tar文件可以通过纯顺序读取操作进行访问，这可以显著提升来自云存储解决方案的I/O性能。
2.3 Pipeline 优化
NeMo Curator中的流水线设计用于扩展至数百个节点，同时达到最大吞吐量。为了在扩展规模下有效运行，作者在流水线的所有阶段都使用了GPU加速，并利用Ray自动调整每个阶段的工人数量以平衡吞吐量。除了在GPU上运行模型推理外，作者发现使用NVIDIA GPU上的硬件视频解码器（NVDEC）和硬件视频编码器（NVENC）在解码和转码阶段带来了3倍的速度提升。
由于每个阶段都有不同的吞吐量，很容易受到吞吐量最低的阶段限制。例如，在剪辑 Pipeline 中，作者使用不同的模型来生成视频嵌入和视频字幕。视频字幕是通过具有数十亿参数的视觉-语言模型生成的。另一方面，视频嵌入是通过规模较小的模型生成的，大约有几千百万参数。在这种情况下，视频字幕模型比视频嵌入模型的吞吐量低得多，会成为 Pipeline 中的 Bottleneck 阶段。在不浪费GPU周期的情况下对这些模型进行推理，需要一个系统来自动平衡工作负载。为了解决这个问题，作者构建了一个自动平衡系统，如图3所示，该系统将为 Pipeline 中的每个阶段部署最优数量的工作者。这会导致吞吐量较低的阶段拥有更多的工人，以跟上吞吐量较高的阶段。如图3所示，通过利用作者的自动平衡系统以及Ray在GPU工作者之间进行 Stream 分布式执行，作者可以在整个 Pipeline 中实现显著的加速。
自动平衡异构模型NeMo Curator将剪辑、分片和自动平衡系统整合到一个平台上，使用户能够扩展其视频内容制作工作。这些系统对于训练高质量的视频基础模型至关重要，因为模型的质量直接取决于训练数据。NeMo Curator简化了尝试各种数据内容制作策略的实验过程，从而加速了视频基础模型（VFM）的研究与开发进程。
- Megatron Energon的高效多模态数据加载
VFMs通常使用
张图像和
段视频进行训练。在训练过程中，需要一个高效的加载数据策略以实现高吞吐量。作者的VFM训练框架利用Megatron Energon的多模态高效加载数据能力来实现这一点。在本节中，作者讨论了NeMo中实现可靠和可扩展VFM训练的核心Megatron Energon特性。
3.1 数据集准备与融合
准备和存储用于VFM训练的大规模数据带来了许多挑战。假设使用一个可以将时间和空间维度压缩8倍的分词器，并且固定图像/视频的分辨率和长度，那么将需要
的存储空间。在大多数计算集群上存储如此大量的数据通常是不切实际的，因此作者使用像AWS S3这样的云存储解决方案来存储作者的训练数据集。Energon使用WebDataset格式从云存储中提取数据集碎片。Energon还允许将多个数据源混合在一起，使用户能够进行实验并找到最适合其VFM的最佳数据混合。
数据集融合的核心应用之一是混合图像-视频训练与序列打包。作者的数据集包含图像和视频，序列长度各异。通常，作者可以使用SBHD注意力格式，并将相同长度的序列分组到同一个训练批次中。然而，这要求作者将视频模型训练分为不同的阶段，从而在作者的工作流程中引入复杂的加载数据逻辑。序列打包是一种技术，它沿着序列维度连接多个训练样本。这种方法消除了过度填充的需要。由于每个微批次处理的 Token 更多，GPU计算和GPU内存利用率显著提高。这对于VFM训练特别有帮助，因为它允许在同一微批次中将不同长度和分辨率的图像和视频打包，如图4所示。尽管序列打包将微批次大小限制为1，但它有助于最小化计算浪费，并消除了复杂逻辑的需求，从而显著简化了加载数据的过程。
3.2 优化网络使用
在AWS S3等云服务上存储大规模视频数据时，由于每个模型并行等级通常需要下载相同的数据分片以保持其数据加载器一致，这往往会在分布式训练期间对网络带宽造成压力。随着扩展到数千个GPU，这一问题变得更加严重，从而降低了整体训练吞吐量。作者的解决方案是为每个等级分配一个独特的数据分片，然后使用全收集操作将这些分片分配到各个等级。这种策略保持了每个等级的数据加载器的一致性，同时减少了冗余下载，从而在低带宽环境下提高了训练吞吐量——即使增加了通信开销。
- 视频基础模型训练的扩展
NeMo提供了一个用户友好的界面，用于扩展视频扩散模型训练。视频基础模型可以通过自回归下一个 Token 预测或完整序列扩散目标进行训练。NeMo在大语言模型（LLMs）上建立的一套工具可以用于自回归模型的重用。在本节中，作者讨论了作者新的可扩展扩散训练能力，包括对最先进的扩散 Transformer （如NVIDIA Cosmos扩散世界基础模型）的实现。作者的扩散训练框架提供了各种模型并行选项，对于多个模型配置实现了超过40%的模型FLOPs利用率（MFU）。
4.1. 扩散公式
扩散模型通过迭代采样过程学习从随机噪声中生成高质量视频。具体来说，扩散模型被训练以逆向一个固定的时间相关随机过程，该过程通过添加噪声来破坏数据。作者训练了一个神经网络
，它估计添加到创建损坏视频潜在变量
中的噪声
。作者的扩散模型可以基于多个信号进行条件化，例如噪声调度时间步长
和文本嵌入
。作者使用以下去噪分数匹配损失公式来训练模型：
其中
是一个函数，用于对贡献于训练目标的各个去噪任务进行加权。此外，作者使用了与EDM中提出的相同去噪网络预条件、损失加权以及噪声分布。
作者可以通过迭代地对给定条件信号
的随机噪声进行去噪过程
步，从作者的扩散模型中采样视频。在这里，作者可以利用各种随机采样器来生成视频。NVIDIA-NeMo提供了二阶EDM Heun采样器以及用于加速视频生成的更高阶RES采样器。
4.2 训练流程
训练视频扩散 Transformer 需要多个不同组件。图5展示了从视频扩散 Transformer 训练和生成视频样本的端到端过程。
第一步是视频分词。给定一个输入图像或视频，作者使用因果时间3D分词器生成时空 Token 。通过3D patchify操作将3D Token 展平成一系列3D时空块。此外，NVIDIA-NeMo还提供了视频分词器的微调能力。关于分词器训练的详细信息请见第4.3节。
接下来，作者训练扩散 Transformer （DiT）以去噪被高斯噪声污染的3D时空块。扩散 Transformer 将输出完整视觉 Token 序列的去噪版本。在训练过程中，作者通过扩散噪声调度时间步长
和描述输入视频样本的 Prompt 文本嵌入来条件化扩散 Transformer 。时间步长
通过自适应层归一化（AdaLN）机制应用，其中作者根据
条件化层归一化的伽马和贝塔参数。使用AdaLN的一个问题是，模型总参数中很大一部分是由自适应层归一化层的MLP参数组成的。作者提供了使用AdaLN-LoRA的选项，通过将操作分解为两个低秩矩阵乘法来减少AdaLN层使用的总参数。作者观察到AdaLN-LoRA也提高了训练过程中的MFU。文本条件化通过每个 Transformer 块中的交叉注意力层应用，其中作者计算视频 Query 张量与文本嵌入键值张量之间的注意力。
最后，作者使用扩散 Transformer 预测的噪声，通过并行化的EDM扩散 Pipeline 计算扩散损失。作者的自定义EDM扩散 Pipeline 与模型并行性兼容，例如张量并行、 Pipeline 并行和上下文并行。
4.3. 训练自定义视频分词器
直接在原始视频或高分辨率图像上进行训练通常计算量过大。一种广泛采用的方法是使用视频 Token 器将原始数据编码到潜在空间中，这可以通过将数据集大小减少超过100倍来显著加速学习过程。如图4所示，训练过程从3D因果视频 Token 器开始，用于压缩视频输入，从而实现高效的VFM训练。认识到视频 Token 器在其中的关键作用，NVIDIA-NeMo提供了一个全面的工具包，使用户能够自定义或微调视频 Token 器模型，以更好地满足其特定需求。
4.3.1. 个性化分词器架构
NVIDIA-NeMo利用了来自Hugging Face Diffusers的开源模型，以实现无缝的视频分词器定制。用户可以通过简单地调整配置JSON文件中的几行来修改视频分词器架构。例如，默认的视频分词器提供了8倍的空间压缩。用户可以通过添加和微调额外的编码器/解码器块来增加空间压缩率，这可以通过进一步减少总数据集大小来显著降低下游VFM训练成本。这种方法为用户提供了一种灵活性，以平衡生成质量与可用计算资源。
4.3.2. 视频 Token 化器的微调
作者支持对预训练的 Token 器或具有从预训练 Token 器部分权重初始化的定制 Token 器进行微调，以适应专有用户数据。微调预训练的 Token 器使用户能够快速将现有模型适应其特定的视频数据。微调定制 Token 器允许用户修改视频 Token 器架构，根据其计算资源和下游任务需求调整数据压缩 Level 。作者还为用户提供使用多种损失函数的选项，包括均方误差（MSE）、库尔巴克-莱布勒（K-L）散度损失、学习感知图像块相似度（LPIPS）损失和生成对抗网络（GAN）损失。
4.4 并行化扩散Transformer
扩散模型在数据、模型大小和计算方面表现出令人印象深刻的扩展行为。DiT架构在生成和感知任务中显示出理想的扩展特性。现代视频基础模型已采用类似的基于transformer的架构。然而，DiT架构在训练和推理过程中都存在一些挑战。在长序列视频上训练具有数十亿参数的DiT模型，具有高激活内存使用，需要使用各种模型并行策略。在本节中，作者将讨论NeMo通过Megatron Core为视频扩散训练提供的各种训练并行选项。
张量并行（TP）。增加DiT模型的整体规模也增加了训练过程中存储激活的内存需求。TP将每层的参数张量分割，并将它们分布到多个GPU上。因此，每个GPU只持有和处理总参数的一小部分。这种策略不仅减轻了内存需求——因为模型参数和激活在每个GPU上占用的空间更少——而且还有助于容纳那些否则会超出单个GPU容量的极大规模模型。然而，由于每个GPU的 Kernel 工作负载较小，CPU开销有所增加。TP的额外CPU开销以及通信成本可能会导致整体MFU下降。
完全分片数据并行（FsDP）。FsDP将模型的参数、梯度以及优化器状态分片存储在多个GPU上。在传统的数据并行设置中，每个GPU维护模型参数的完整副本。FsDP将模型参数分割成多个分片，并将每个分片存储在不同的GPU上。因此，没有任何单个GPU持有全部参数集。在正向传播过程中，每个GPU只需使用其拥有的局部分片，从而减少内存使用。FsDP利用数据并行GPU之间的集体操作。对于参数计算，在正向和反向传播过程中都应用了全聚合操作。对于参数梯度，仅在反向传播过程中应用了reduce-scatter操作。FsDP将所有通信与模型计算重叠，提供了一种无缝的模型并行体验，无需像使用3D并行性时那样寻找最佳配置。然而，FsDP的性能高度依赖于训练集群的GPU间通信速度，尤其是在大型模型训练中，通信开销相当大。这个问题可以通过将FsDP与TP结合使用来缓解。通过将FsDP分片组限制在数据并行组大小，TP可以帮助减少每个GPU的模型状态和激活大小，从而有效减少FsDP的通信和激活内存开销。
上下文并行（CP）。为了在长序列视频上训练作者的DiT模型，由于高激活内存成本，作者必须在序列维度上并行化transformer计算。为了实现长上下文训练，作者利用CP通过在序列维度上分片输入张量来并行化网络激活，这有效地将所有层的激活进行了分片，与序列并行（SP）不同，后者仅分片LayerNorm和Dropout层的激活。
在使用CP时，除了注意力机制之外的所有模型组件均未进行修改，因为它们没有跨 Token 操作。对于注意力机制，作者必须使用序列中所有 Token 的键（K）和值（V）来计算每个 Query Token （Q）。CP将Q和KV分割成多个块，因此作者必须通过利用通信集合来确保计算的正确性。具体来说，作者将GPU主机排列成环形拓扑结构，其中每个GPU处理一个Q块，并迭代地与每个KV块进行块级注意力计算。
在正向传播过程中，全聚合操作组装完整的KV序列，以便每个Q块可以处理所有 Token 。在反向传播过程中，reduce-scatter操作将梯度分布回各自的碎片。为了减少激活内存成本，作者在正向传播过程中只在每个GPU上存储序列块的一个KV块，然后在反向传播过程中再次聚合KV。通过TransformerEngine将全聚合和reduce-scatter操作作为对等（P2P）通信实现，允许在块级注意力计算的同时进行KV通信，从而进一步提高训练效率。
流水线并行（PP）。通过P，作者可以将Transformer模型的层分配到多个GPU设备上，每个设备分配相同数量的Transformer块。作者将数据批次拆分为更小的微批次，然后在微批次之间流水线执行模型。PP的一个主要挑战是流水线气泡，即GPU空闲的时间块。近年来，已经提出了各种流水线调度方案。在NeMo中，作者提供了默认的流水线，使用GPipe算法和交错流水线以减少流水线气泡的大小。
对于DiT模型，作者特别探讨了如何高效地在 Pipeline 阶段之间传递条件信息。对于Transformer模型，通常仅在各个PP阶段之间传递隐藏状态。然而，对于基于DiT的视频基础模型，作者可能还有几个额外的条件信号，包括文本嵌入、图像嵌入、相机参数嵌入等。这些嵌入通常在每个Transformer块中的交叉注意力块中使用。为了确保正确的训练，作者必须能够在每个 Pipeline 阶段使用这些条件信号进行交叉注意力计算。
作者考虑了两种缓解该问题的方法，如图6所示。第一种方法是将在 Pipeline 阶段之间传递额外的条件信号以及隐藏状态。这需要在第一个DiT块之前计算一次嵌入，然后在每个PP阶段之间进行通信。第二种方法是在每个PP阶段分别计算每个嵌入，然后使用计算出的嵌入对模型进行条件化。在第一种方法中，每次作者通过增加通信缓冲区大小以添加新的嵌入来引入新的条件信号时，都必须修改PP调度，这会增加通信成本。在第二种方法中，作者不需要修改PP调度，但作者必须花费更多的计算资源，因为作者需要在每个PP阶段重新计算条件信号。作者观察到，第二种方法相对于第一种方法，其额外的计算成本与MFU之间的权衡更为有利。此外，第二种方法也更容易实现，并且随着新的条件信号的添加而容易扩展。
4.5. DiT 并行化性能研究
本节提供了对NVIDIA-NeMo的详细评估和性能基准测试。作者首先探索算法-系统协同设计空间，以找到性能最佳的模式配置。随后是详细性能基准测试结果，与其他开源替代方案如Fast-DiT（Fang等，2024）的比较以及扩展实验。本节最后强调了灵活训练框架对于扩散模型的重要性，并对作者的学习成果进行了简要总结。
4.5.1 算法-系统协同设计
为了在模型大小和上下文长度扩展的同时实现高计算效率，作者探索了算法与系统协同设计空间以及并行化策略。
AdaLN-LoRA架构。作者对DiT架构所做的关键改动之一是改变AdaLN投影。AdaLN投影消耗了模型的大部分参数[隐藏维度 x
隐藏维度]。
例如，使用基础DiT模型架构，一个70亿参数的模型将会有一个AdaLN投影，其本身就会消耗27亿参数。
此外，由于AdaLN是在时间步嵌入（而不是完整上下文长度）上进行的，并且是一个单一的GEMM，它也不允许采用直接模型并行策略来平衡计算与通信。为了改进这一点，作者实现了AdaLN-LoRA。这不仅使作者能够减少AdaLN中关注的参数密度，还允许作者在给定的固定模型大小下使用更大的隐藏维度。这导致更高效的GPU计算利用率和更好的计算-通信比率，如图7所示，这有助于作者将计算性能提高高达1.2倍。
模块化上下文并行。在探索并行化策略时，作者做出的另一个选择是在使用CP（跨层并行）时禁用层中CrossAttention部分的上下文并行。作者做出这个决定是因为KV上下文长度不足以要求并行或在使用并行时的重叠通信。
QK归一化。作者面临的另一个挑战是在点积注意力计算之前对
和
进行归一化。在张量并行性下，
和
通常被分散到注意力头的数量上，而在整个隐藏维度上计算归一化会需要过多的通信。作者选择限制每个头的归一化，以提高计算吞吐量，并且发现这对模型性能没有不利影响。
4.5.2. 性能基准测试
作者研究了具有不同模型大小和上下文长度的训练工作负载的计算吞吐量性能，并采用多种并行化策略和计算规模，结果表明作者的训练框架可以达到高达48.2%的MFU。
性能评估。表1描述了所使用的负载配置。作者报告了所有工作负载在
节点上实现的最佳性能，除了28B - 阶段3，该性能是在
节点上报告的。
作者以Fast-DiT作为基准，与NVIDIA-NeMo进行比较。为了进行公平的比较，作者在Fast-DiT中启用了AdaLN-LoRA和FSDP支持，使用Huggingface Accelerate。图8显示了使用NVIDIA-NeMo和Fast-DiT实现的每个工作负载的最佳计算性能。作者的训练框架能够在所有工作负载中持续实现高吞吐量，并在7B模型上比Fast-DiT高出最多
。当将模型规模扩展到28B时，NVIDIA-NeMo仍能保持强大的性能，而Fast-DiT由于内存容量不足，无法运行28B模型。
扩展效率。NVIDIA-NeMo还实现了高强扩展效率。作者在固定工作负载下，不断增加8xH100节点的数量，持续测量计算吞吐量，发现NVIDIA-NeMo实现了如图9所示95%或更高的强扩展效率。
观察与学习。如图8所示，每个工作负载都需要独特的并行策略组合以实现高效率，这突出了灵活训练框架的必要性。为了进一步强调全4D并行（TP、CP、PP和FSDP）的重要性，作者在扩展上下文长度的情况下，研究了DiT-7B与AdaLN-LoRA在各种并行配置下实现的最佳吞吐量。图10显示了不同策略相对于最佳4D并行策略的性能表现。
在每种上下文长度中，最佳并行策略都是独特的，没有任何单一的并行策略能够单独持续地实现良好的性能。在较短的上下文长度和较小的模型尺寸下，FSDP通常足以实现可观的训练吞吐量。然而，随着上下文长度和激活内存足迹的增长，FSDP单独已不再是最佳配置，需要TP、CP和FSDP的组合。在保持高计算效率的同时，PP也成为在分布式大型模型时的必要策略。以下是作者的一些学习总结。
并行策略的上下文缩放性能。在扩展扩散 Transformer 训练时，考虑不同并行策略之间的权衡至关重要。TP、CP和PP在应用于DiT架构时各有优缺点：
-
- 上下文并行性：CP主要减少了训练过程中每个设备上激活的内存占用。在长视频序列的场景中，作者可以有效地隐藏CP的通信开销。然而，当序列较短时，这种开销变得很大程度上暴露出来。
-
- 张量并行+序列并行：TP和SP结合可以有效减少参数和激活的内存使用。然而，由于自适应层归一化（AdaLN）在DiT架构中的广泛应用，TP通信无法很好地与计算重叠。此外，当模型规模不是很大（即
参数）时，参数过度拆分会导致矩阵乘法效率降低。 3. 3. Pipeline 并行性：虽然PP无法帮助减少激活内存的使用，并在处理开始和结束时引入GPU空闲时间气泡，但它可以有效降低参数内存的使用。因此，在处理较大模型尺寸时，尤其是在GPU间通信带宽较低的环境中，它是有益的。
总结来说，在权衡了并行技术的优缺点后，可以提出以下建议：
-
- 当模型大小和上下文长度相对较小时，FSDP（分布式深度学习框架）可以满足需求。
-
- 当模型规模相对较小且视频序列较大时，应优先考虑CP。
-
- 随着上下文长度的增加，应优先考虑CP，而当模型规模增长时，对于任何模型分片，应优先考虑TP/FSDP。
-
- 如果模型规模较大，应在节点内部优先考虑TP，但程度要保证较小的GEMMs仍然高效。超出这个范围，则应使用PP。
-
- 在模型大小和上下文长度特别大的情况下，可能需要TP、PP和CP的组合。一般来说，TP/SP通信应保持在节点内部，而CP和PP组可以进行跨节点通信，同时仍然足够地隐藏通信开销。
4.6 空间-时间DiT（ST-DiT）
除了全注意力外，ST-DiT（时空DiT）引入了空间和时序注意力，分别应用于帧内和帧间 Token 。这些注意力在Transformer层中堆叠，由于它们的批量大小和序列长度存在显著差异，其计算复杂度差异很大。
全注意力机制将所有帧的所有 Token 扁平化为一个序列，这导致小批量大小（例如，1或2）和长序列长度（例如，256K、512K等）。空间注意力分别应用于每个帧，这相当于相对较大的批量大小（例如，数百个帧）但短序列长度（例如，每个帧中几千个 Token ）。时间注意力类似于空间注意力，因为它们的批量大小和序列长度是相互转置的。输入形状的变化导致在并行化策略上存在不同的权衡：较大的批量大小更倾向于数据并行（DP），而较长的序列长度则需要上下文并行（CP）。
由于DP将批维度拆分，因此它不会触发任何用于注意力的跨GPU通信。因此，每个GPU可以独立地使用其分配的批拆分进行工作。CP将长序列分割成块，并在GPU之间分配它们，因此每个GPU只需要计算和保存序列块的激活。然而，CP插入了一个全收集操作来收集完整的KV序列，因为每个 Token 的Q需要与同一序列中所有 Token 的KV进行计算。TransformerEngine中的CP实现将全收集操作改为环形拓扑结构中的P2P通信，如果序列足够长，这可以在注意力计算下隐藏通信开销。
作者提出将CP应用于全注意力，但将DP应用于空间和时间注意力，如图1l所示。由于输入/输出形状不匹配，这不能直接实现。假设每个Transformer层的输入形状为[b,
，d]，其中
是批大小，
和
是每个帧的高度和宽度，
是帧数，
是隐藏大小。每个注意力的序列长度由
、
和
确定。CP的全注意力输出形状为[b,
d]，但后续的空间注意力需要输入形状为
，
, d]，而最后的时序注意力期望输入形状为
。为了获得每个注意力期望的输入形状，作者在注意力之间插入全全通信以进行张量形状转换。
现有的ST-DiT实现要么没有同时考虑所有三个注意力机制，即DSP，要么在通信开销方面存在显著问题，例如DeepSpeed Ulysses中的每个注意力机制都涉及所有节点之间的通信。与之前的实现相比，作者的 Proposal 不仅考虑了所有注意力模块，而且能够完全隐藏全注意力机制的通信开销，并将空间和时间注意力机制的通信开销降至最低（仅需两个所有节点到所有节点的通信）。表2比较了不同并行策略在时空注意力上的性能。可以看出，对于时空注意力使用全节点到节点的通信比CP策略性能更好，并且对于作者的12B ST-DiT配置，在NVIDIA H100 GPU上能够实现高达40%的MFU。
- 高效视频生成推理
视频扩散模型必须在推理时对长序列的3D时空 Token 进行去噪。为了生成高分辨率视频样本，这些模型必须对成千上万的 Token 进行去噪，而这些 Token 可能仅相当于几秒钟的实际视频。采样过程是迭代的，这可能导致每个采样视频的生成时间较长。
作者提出了一种利用多个GPU上的上下文并行化来提高视频生成效率的方法。在接下来的章节中，作者将详细介绍作者的推理方法和在不同配置下的性能结果。
5.1 并行推理算法
作者采用了一种简单而有效的方法，通过上下文并行化来并行化视频扩散模型推理。该方法通过在GPU和去噪序列块之间并行分配去噪工作负载，显著加速了推理过程。
图12展示了生成单个视频的过程。作者首先将输入噪声的潜在表示在序列维度上进行划分，并将每个块分配给特定的GPU。一旦输入噪声的潜在表示被分配，每个GPU上的DiT独立对其分配的子序列执行去噪过程。这个过程重复进行
次，其中
是去噪步骤的总数。去噪操作是扩散模型的核心组件，涉及迭代优化以产生连贯且高质量的输出。通过在多个GPU上并行执行这些操作，作者有效地减少了与GPU数量成比例的总计算时间。最后，将去噪后的潜在张量收集并连接起来，在潜在空间中重建视频的完整序列。然后，使用Cosmos Video Tokenizer对组合的潜在序列进行解码。
5.2 性能研究
作者在多种配置下对视频扩散模型的推理性能进行了基准测试。具体来说，作者分析了NVIDIA Cosmos-1.0-Diffusion-7B-Text2World模型的性能。在推理过程中，作者使用了无分类器引导（CFG）（Ho和Salimans，2022年），这需要同时生成条件和无条件输出。因此，在所有实验中，作者使用了全局和微批大小为2。作者观察到，在推理时使用上下文并行和张量并行，其扩展效率可达80%-90%，最高可扩展到32个H100 GPU。通过为上下文并行启用计算/通信重叠，作者可以进一步提高加速比，实现接近线性的扩展。作者还注意到，在H100 GPU上，TransformerEngine的FP8多头注意力（MHA）比BF16在1个GPU上提高了约28%，在32个GPU上提高了约48%。作者的推理流程能够实现快速的视频生成，即使是在多节点系统中也能使用大型模型，这对于向客户提供服务以及用于物理AI应用的合成数据生成流程至关重要。还有其他几种策略也可以提高推理性能，包括模型量化、CFG并行和模型蒸馏。作者将这些探索留待未来的工作。
- 结论
本文介绍了一种使用NVIDIA NeMo训练VFMs的全面且可扩展的框架。NVIDIA-NeMo通过NeMo Curator集成了强大的视频编辑能力，通过Megatron Energon实现了多模态数据加载，以及通过Megatron Core实现了可扩展的视频扩散模型训练。作者展示了算法-系统协同设计的重要性，强调了在训练大型扩散 Transformer 时进行High-Level4D并行化的必要性。作者应用并展示了算法-系统协同设计在作者新的混合并行方法中为定制扩散 Transformer 带来的优势。
作者的广泛基准测试表明，与现有方法相比，在训练和推理时间上均表现出卓越的性能和可扩展性，实现了接近线性的扩展，并且比之前的 Baseline 有显著更高的MFU。通过系统地探索各种并行策略，作者为不同模型大小和序列长度提供了最优并行配置的指导方针。作者的灵活框架为用户提供探索和扩展视频基础模型所需的工具，从而加速新型研究和应用，奠定了坚实的基础。
参考
[1]. Training Video Foundation Models with NVIDIA NeMo
扫码加入👉「集智书童」交流群
（备注：方向+学校/公司+昵称）