B站文生视频模型工程实践

q4vqL2pB

一、前言

近年来，AI 内容生成（AIGC）领域的快速发展令人雀跃，OpenAI 在 2023 年初推出大型语言模型（LLM）GPT-4 受到了学术界和工业界的极大关注。OpenAI 随后在 2024 年初推出文生视频（T2V）模型Sora，能够根据文本指令制作出具有现实风格和富有想象力的场景视频，更是展示了令人惊喜的“世界模拟器”能力。
B站作为UGC内容丰富的视频网站，在视频生成模型领域有着天然数据优势和广泛应用场景。在此之前我们已经有了一段时间的LLM模型训练经验，文生视频模型结构、语料以及训练过程有一定的差异性，本文重点介绍B站TTV团队在文生视频模型上积极探索后的经验及感悟。

二、TTV model




在OpenAI提供的公开信息中，Sora模型实际上是一个Diffusion Model+Transformer架构。本文基于B站在生成式TTV在自研道路上的探索、结合行业进展和工程实践，先后尝试了几种TTV（Text to video，后面将简称为TTV）模型。文中重点介绍由colossal-ai发布的类Sora模型Open-Sora，以及由智谱AI发布的CogVideoX模型。


2.1 OpenSora


OpenSora的核心是Stdit（Spatial Temporal Diffusion Transformer）结构。
DiT（Diffusion Transformer）模型是一种结合了去噪扩散概率模型（DDPM）和Transformer架构的扩散模型，通过模拟数据的逐步去噪过程来生成新的文本，在此基础上发展出了STDiT，STDiT模型就是一种使用Cross-Attention的DiT变体。
STDiT模型 的特点包括：

• 融合空间 - 时间注意力机制：结构巧妙地串联起二维空间注意力模块和一维时间注意力模块，能够在捕捉视频帧内的空间特征的同时，精准模拟视频帧与帧间的时序关联。
  
• 交叉注意力集中模块：紧随在时间注意力模块后，确保文本语义与生成视频的深度对齐。
  
• 计算资源需求降低：相较于全注意力机制显著减少了计算资源需求。
  
• 利用预训练权重：能够更好地利用预训练好的图像 DiT 权重迁移学习至视频场景。
  

图2-1 opensora模型架构

如图2-1所示，OpenSora的Transformer Block的部分包括了Spatial Attention（空间维度上的注意力计算）、Temporal Attention（时间维度上的注意力计算），以及通过Cross Attention将两个维度信息进行交叉计算，此外整个架构还包括一个预训练好的视频编码器（VAE），一个文本编码器（T5）。整个模型的训练阶段如图2-2所示，首先采用预训练好的 Variational Autoencoder (VAE) 的编码器将视频、图片数据进行压缩，将对应的文本描述通过Text Encoder进行embedding化，然后在压缩之后的隐层空间中与文本嵌入(text embedding)一起训练 STDiT扩散模型。

图2-2 opensora训练流程示意图

Open-Sora 模型的训练采用多阶段训练方法，包括大规模图像预训练、大规模视频预训练和高质量视频数据微调等。在数据收集和预处理方面，涉及到对视频美学、动态性、运镜、字幕等信息的处理，构建高质量的训练数据集。
其功能特点有：支持视频数据预处理、加速训练、推理等全套流程；提供视频切割和字幕工具，支持剪辑和 T5 文本调节；实现可变长宽比、可变分辨率和可变时长等功能；能够生成各种风格的视频内容，支持多种生成任务等。


2.2 CogVideoX


CogVideoX【8】使用了由智谱提出的3D Causal VAE以及专家Transformer模块，训练流程与STDIT基本一致，都是通过VAE与Text Encoder将视频数据与文本描述进行embedding化后，输入Transformer Block进行训练。与STDiT不同的是，视频embedding与文本embedding会直接拼接起来做 Full Attention 计算，以便更好地对齐视觉和语义信息。但是这两种模态的特征空间差异显著，它们的嵌入甚至可能具有不同的数值范围。为了在同一序列中更好地处理它们，Transformer模块采用了专家自适应层归一化（Expert Adaptive Layernorm）来独立处理每种模态，如下图所示。处理后的模态信息会继续拼接到一起进行3D Full Attention进行计算，在Attention计算前使用了 3D-RoPE将相对位置信息添加到视觉模态上。

图2-3 cogvideo模型结构

三、工程实践





3.1数据存储与加载


TTV模型训练相较于LLM模型，训练数据集有较大差异性。处理后的训练数据包括视频切片、图片和对应的文字描述。相比于LLM语料库，TTV单个数据文件较小、但整体数据量大，给训练造成了一些不便。
此前的LLM模型训练数据均以5G一个分片整块存储在boss上，每个训练epoch开始前会拉取一个5G的分片到训练机器上，IO效率较高。但视频切片的数据平均大小在1M左右，如果在一个epoch的开始拉取所需的训练数据，boss拉取小文件的效率无法满足我们的训练需求，会造成GPU资源闲置。
结合上述特点，我们和B站基础架构存储团队讨论过多个方案：
方案一：HDFS+文件打包

相较于boss，HDFS能够提供更大的数据访问的带宽上限。HDFS 的负载通常较为平稳，但也存在抖动场景，且DN和NN受任务影响明显，当前HDFS数据99%为2副本存储，当有2台DN受到任务影响时，数据读写性能会有明显下降。此外因为小文件的特点，拉取效率较难满足训练要求，基于这个特点，我们考虑将多个小文件打包在一起，形成一个大chunk file，加速HDFS拉取效率，此外在训练的脚本中，加入一个dataset reader，解析相应的chunk file，重新还原成小视频、图片文件。显然这个办法可以解决文件拉取效率问题，并且还节约了一定的存储空间和文件句柄，对存储媒介友好。但是存在对代码的侵入，以及chunk内容调整不灵活的问题，我们短期没有采用。
方案二：HDFS->backend & Alluxio->frontend

Alluxio 是一个开源的分布式内存文件系统，能大幅提升数据访问性能，加速大数据分析和处理任务。在机器学习训练中经常使用alluxio，快速提供数据给模型。并且他能灵活适配，可与多种存储系统（如 HDFS、S3、Azure Blob Storage 等）和计算框架（如 Hadoop、Spark、Flink 等）集成，适应不同的大数据架构。经基架同学一同努力，最终我们采用了Alluxion Fused+HDFS backend方案，加载数据的方式近似于访问本地文件系统，训练框架无感。alluxio定期同步HDFS上的数据到本地，能较好的满足我们海量小文件的读写以及快速上线的诉求。


3.2 数据预处理优化


TTV数据的预处理过程更为繁琐，除了对视频的文字描述进行tokenizer外，还需要对视频数据进行抽帧、归一化和vae编码，尤其是vae编码这一步在单epoch训练中会占到总时长的18%，并且极为消耗显存。为了优化这部分的速度，我们做了两部优化。

•   数据并行
  

参考训练中的数据并行策略（Data Parallel），在训练任务启动的初始阶段，初始化一个包含全部显卡的全局通信组，将一个epoch训练中最耗算力的编码部分拆分到所有显卡上，待每张卡处理完成后再将所有的结果gather到0号卡上进行聚合和保存，后续的训练阶段只需读取处理好的embedding数据即可。

•  VAE离线化
  

考虑到一份训练数据可能会在后续下训练任务中被反复使用，并且数据处理非常消耗算力和显存，因此在第一步优化的基础上，我们会把聚合好的训练结果持久化，后续直接从HDFS\Alluxio中读取，整体显存和训练时间都有较大优化。
上述过程如图3-1所示：

图3-1 数据读取、预处理示意


3.3 模型并行优化


虽然TTV模型通常参数量较小，但视频训练数据更为复杂，序列长度较长。因此TTV模型训练过程中的激活内存成本极高，导致训练速度明显更慢，且限制了训练数据的清晰度与抽帧率的提高。
在Transformer Block中，激活显存占用的大头为Attention的计算，在不使用Flash-attention的情况下，attention激活占用与序列长度的平方成正比，而使用了时空Cross-Attention设计的STDIT模型，训练过程中更是包含了3次attention的计算，使得激活占用更是大大增加。针对Transformer训练过程中，序列长度过长导致的激活占用高的问题，业界已有一些序列并行方案提被出，序列并行是一种专门为跨多个设备分发长序列和激活而设计的技术，以下四种是主要的序列并行技术， Ring-Attention、Megatron-SP、DeepSpeed-Ulysses和Dyncamic SP，首先将介绍下各方案的设计。


业界方案

• Ring Attention
  

Ring Attention 的核心思想是将输入序列分割成多个块，并将这些块分布在多个计算设备上进行并行处理，通过使用 Online Softmax 机制，在不保留完整序列长度的情况下计算注意力分数。每个设备首先对自己的数据块进行局部的自注意力计算，然后将关键信息（key-value pair）传递给下一个设备，通过一个环形的数据传输策略，实现在不增加单个设备内存负担的情况下处理超长序列。然而，环形注意力对 P2P 通信的依赖在高延迟环境中可能效率较低。

• Megatron-SP
  

Megatron框架中的的序列并行，是在张量并行(Tensor Parallelism)的基础上，将Transformer Block中的LayerNorm以及Dropout层的输入按Sequence Length维度进行了切分，使得各个设备上面只需要做一部分的Dropout和LayerNorm。虽然减少了每张卡上的激活占用，但在通信过程中引入了额外的all-gather和reduce-scater操作。因在设计上依托切分注意力头来实现并行，使用上也会受到注意力头数量的限制。

图3-2 megatron-sp

• DeepSpeed-Ulysses
  

DeepSpeed-Ulysses 【7】引入了一种创新的方法，通过利用全对全（all-to-all）集体通信来训练长序列。在处理长序列时，它将查询（query）、键（key）和值（value）矩阵在注意力头之间进行切分，但保留原始注意力计算结构。这个过程通过两组全对全通信来实现，这两组通信在序列分割和注意力头分割之间交替进行。这样的设计使得在处理长序列时，能够在保持计算结构的同时，有效地在多个 GPU 之间分配数据，减少通信开销。

图3-3 deepspeed-ulysses

• Dyncamic SP
  

与以上三种针对单序列维度内的并行性方案不同，DSP（Dynamic Sequence Parallel）【6】方案是针对多维序列的并行问题所设计。在多维Transformer模块中，每个序列维度的计算其实是独立进行的，因此可以分别在不同纬度切分和计算，仅在维度计算交换的时间点使用高效的全对全操作（all-to-all）来为中间序列切换并行维度并重新进行动态切分。这种方法使 DSP 能够独立于模块内的计算逻辑外，消除了模块内许多不必要的通信。

图3-4 dsp



相关实践

•  基于OpenSora SP实现
  

OpenSora模型使用了时空交叉注意力机制，会分别计算视频空间维度的Attention、时间维度的Attention，并最后和文本embedding进行交叉Attention计算，其设计更适合使用DSP方案。具体实现上，我们会在空间注意力计算前，先在空间维度进行切分，SP并行组内的各rank分别计算不同的seq片段。随后在时间维度注意力计算前，进行一次all-to-all通信进行同步，并交换切分纬度到时间维度上。因为与文本embedding的交叉注意力计算，只与空间信息有关，因此可在交叉Attention计算后再进行一次all-to-all通信来同步计算结果。

• 基于CogvideoX的SP实现
  

考虑到CogVideoX模型虽然使用了视频信息与文字信息，但两种embedding是在单一维度进行拼接，并进行全局Attention计算的，因此本身属于单维Transformer。对于单维Transformer，我们选择实现了DS-UIysses方案，在实现上：
a. 先在Transformer Block之前，沿sequence维度进行切分，但只对视频的hidden state的seq维进行切分。文字embedding的部分不切分，并与切分后的视频embedding拼在一起。
b. 相对位置编码是需要对全序列长度进行计算的，因此在每个3D-Rope计算前，进行一次all-to-all通信，回收sequence维度的信息。
c. 因为每个seq段都拼接了文字embedding，因此首先需要通过remove_extra_encoder方法，移除每段seq冗余的text embedding。在计算完Attention后，进行一次all-to-all通信，在attention head维度回收信息，并回到sequence切分的状态。在MLP Block计算前，通过 add_text_encoder补上每段seq都有的文字embedding部分。
d. 在所有的Transformer Block后，进行gather sequence操作，合并SP组各组上的计算结果。
通过使用SP并行策略，CogVideoX由单机16卡只能训练45帧1080p的数据，提升至可训练221帧1080p的训练数据。

图3-5 基于CogVideoX的SP实现示意图

四、文生视频模型在NPU架构上的工程实践




目前我们的训练算力构成为GPU + NPU，但GPU与NPU在芯片设计上差异较大，软件栈和生态也存在较大区别，因此需要做相应的适配和优化才能充分利用NPU资源进行训练。


4.1 基础适配

• 模型适配
  

目的是让模型能够在npu环境下开箱启用（训练），保证pipeline可运行。其主要工作为检查依赖cuda硬件及精度在npu下受限的算子，查询入参匹配的npu-wrapped算子，进行替换与轻量的代码适配，例如替换T5模块中的LayerNorm为NPURmsNorm、在nn.Conv3d算子中显式指定使用torch.float16精度等。

• 框架移植
  

我们的训练框架中部分模块开采用megatron-core作为加速手段，需要移植到对应NPU的版本，主要参考华为Megatron-NPU仓库的范例和实现，进行移植。

• 精度验证
  

可训练后，需要进行精度的验证。精度验证需注意保证GPU版本与NPU版本的第三方依赖保持一致，并固定代码中随机的部分，例如随机数设置、数据的抽样、Vision encoder中的加躁信息等。可以借助华为的精度验证工具，dump主要算子的输入输出，并逐步进行API级别、模块级别与整网级别的精度验证。
不过值得注意的是，因为底层实现有所区别，配置不同，以及fusion的使用与否等等，loss是无法完全对齐的。下面列出几个常见的，会引起loss差异的参数以供参考。

• transformer_impl: local -> transformer_engine
  
• attention_softmax_in_fp32: False -> True
  
• apply_rope_fusion: False -> True
  
• rotary_fusion: False -> True
  
• swiglu_fusion: False -> True
  

案例


基于内部基座版本：global batch size 160，Run 1曲线代表NPU，Run 2曲线代表GPU。
2940-steps:
Loss diff max = 0.46540455520153046, diff mean = 0.006955942064507858

图片

1500-steps:
Loss diff max = 0.2761184275150299, diff mean = 0.008173752022286256

图片

对比趋势和diff，基本一致，可以确认模型参数基本可用，更细粒度则需要做到分层精度对比以及生成视频benchmark对比。
基于CogVideo-5B基座版本：8卡，25帧，480P


图片


4.2 优化

• profiler性能分析调优
  

a.后向耗时问题定位
下图profiler中可以发现，在后向过程中执行了一次前向，占用接近20%，但其实我们并未显式地配置重计算策略。



后续分析代码发现，训练中会默认使用torch.utils.checkpoint对全层数的transformer进行重计算，但基于如下几点原因可以对其优化：
A.  显存仍未用尽，有tradeoff的余地
B.  需要重计算的层数应该可控
C.  重计算作用在一个融合算子上，其中涉及MLP，GN，QKV计算，attention计算等，但其实只有attention计算使用该策略的价值最高
经过一系列的代码调整后，最终训练速度提升约12%。
b.融合算子替换
根据profiling分析，红框中gelu算⼦实际执⾏时是以多个⼩算⼦拼接的形式下发和执⾏的。

图片

可以使⽤融合算⼦F.gelu 进⾏替换，优化下发和执⾏。跟据profiling中的API调⽤栈可以定位到该算⼦主要出现在T5模块和megatron中。
c.使用连续张量减少切分操作

图片
图片

如上图profiling所示，耗时占比最大的3个op中，第二第三都是计算密集型的，可以从策略或算法的角度优化。第一位StridedSlice是访存密集型，改变使用StridedSlice 的上层算子的输入tensor的存储方式（使用torch.contiguous）即可连续分配，优化原理说明如下：
torch.contiguous操作将非连续张量转换为连续张量，后续的切片和访问操作大幅简化，甚至避免StridedSlice的调用。
主要体现在三种场景：
其一，连续张量的切片不需要依赖复杂的stride信息。硬件可以更高效地预取数据，提高计算速度。
其二，非连续张量的 StridedSlice 需要动态计算目标地址，频繁调用可能带来性能瓶颈。连续张量的切片是直接基于线性偏移量完成的，减少了计算需求。
其三，某些操作（如 view）要求张量是连续的。如果张量已经是连续的，相关操作无需隐式调用 StridedSlice 或创建新的张量拷贝，直接提升性能。
在我们的训练中得到加速的主要是第三种场景。上右profiling是优化后的结果。

• FlashAttention(NPU)
  

在torch2.0之前的时代，业界都会使用朴素的standard_attention进行注意力机制的计算，但当其attention_mask为精度f32时会引起巨大的显存占用及 NaN-bug。后续torch推出了一个改良算子scaled_dot_product_attention(sdpa)，进行算子融合，优化内存使用，适合长序列，内置支持 Dropout 和 Causal Mask。现今，主流的大模型结构中都会采用FlashAttention，是一种sdpa的实现，相较于原版scaled_dot_product_attention，其对QKT进行了分块处理，避 免存储完整的注意力矩阵，可处理更长的序列。
图片

在迁移至NPU架构后，我们逐步将原有的注意力机制代码，改进为基于 NPU底层实现的FlashAttention，其中涉及的抽象、封装，引入必要的设计模式等。以下仅对flashAttention作简单介绍。
FlashAttention本质是算子和数据的融合：

• 将多个算子合并为一个，简化计算过程，减少计算量，提高计算效率。
  
• 将多个中间结果合并为一个，减少内存占用，提高内存利用率；减少不同算子之间的传输，提高数据处理效率。
  
• 简化代码实现，减少代码量，提高代码可读性和可维护性。
  

FlashAttention实现原理【1】：三基石

• Tiling切片：利用高速SRAM代替内存，但SRAM内存小，无法一次性完成所有数据的注意力计算，需要进行分块计算，对应上文中的QKT分块。
  
• 重计算：放弃中间结果写回，需要使用时重新计算，用计算换访存。
  
• Kernel Fusion：将多个操作融合为一个操作，基于Tiling利用一个kernel完成整个计算，对应上文中的算子融合。
  

以下是前向、反向过程的公式描述，具体细节不在此讨论。
前向Forward: FlashAttentionScore[2]

   

   

   

   

后向Backward: FlashAttentionScoreGrad[3]

   

   

   

• 虚拟内存特性：expandable_segments【4】
  

一般情况下，由PyTorch自己管理虚拟地址与物理地址映射，降低内存碎片。其原理大致如下：对于大于 2MB 的分配，分配器会调用 aclrtMalloc，以获取与用户请求大小完全相同的内存分配。后续计算中，如果这些分配中的某些部分空闲，它们可以被重新用于其他请求。这种方式在程序多次请求相同大小或者是该大小整数倍的内存时效果很好。许多深度学习模型的行为符合这一特点。
然而，有一种常见的例外情况是，批次大小在每次迭代中会略有变化，例如在批量推理中。当程序最初以批次大小N运行时，会为该大小进行合适的内存分配。如果后续运行的批次大小变为N−1，现有的内存分配仍然足够使用。然而，如果批次大小变为N+1，则需要进行新的内存分配，这些后续分配并非所有张量的大小都相同。一些张量的大小可能是(N+1)×A，而另一些可能是(N+1)×A×B，其中A和B是模型中与批次无关的维度。由于分配器会在现有分配足够大时重用它们，因此某些(N+1)×A 的分配可能会勉强适应已经存在的N×B×A 段，尽管不完全匹配。当模型运行时，这些段会被部分填充，导致在段末尾留下不可用的空闲内存切片。最终，分配器可能需要调用 aclrtMalloc 为新的(N+1)×A×B 段分配内存。如果没有足够的内存，就会抛出异常，结束程序。对于有 50 层以上的模型，这种模式可能会重复 50 多次，从而产生许多小的内存碎片。
通过分析训练时采集的profiling文件，发现我们研发中常见的显存瓶颈符合上述场景描述：
a.动态shape场景，比如VAE中的升降采样，shape随step增加而增大，从而导致显存块不能复用，碎片上升
b.transformer一般由多层组成，其带来的大量激活数据与激活数据size不统一，会给显存复用带来困难，不仅会影响性能，严重时导致OOM
采用NPU自研的“内存池扩展段”功能，有助于优化上述场景。
“内存池扩展段”在最初创建一个内存段后，可按需扩展其大小。与每次分配都创建一个新的内存段不同，它尝试为每个流（stream）创建一个内存段，并根据需求动态增长。当程序运行到批次大小N+1的情况时，这些分配会很好地排列到单个大型内存段中，直到段被填满。然后，分配器会请求更多内存，并将其追加到该段的末尾。这种方式不会产生太多无法使用的内存碎片，因此更有可能成功找到所需的内存。

五、后续工作方向





1.引入流水线并行（Pipeline Parallelism）


如上文所述，现在已采用的优化方案多集中在tensor计算的某些维度，比如sp/cp针对的是tensor的序列维度，flash_attention针对的tensor的特征维度。在业界主流优化方案中，以layer或模块作为优化对象的Pipeline Parallelism，PP也是值得探索的方向。
优点：

• 将模型分成多个阶段，每个阶段只在一个NPU上运行，每个NPU只需要存储它负责的部分模型，而不是整个模型。这大大减少了显存消耗。类似FSDP与ZERO机制。
  
• 模型并行MP以及张量并行TP会导致频繁的跨设备通信，而PP通过流水线操作仅传递每个阶段最终结果，减少了通信负担。
  
• 扩展性强，比如transformer这类的堆叠结构
  

挑战：

• 流水线填充与同步延迟：前向传播和反向传播之间存在同步延迟（Pipeline Bubble）。增加batch size是常见的优化策略，但可能导致内存压力增大。
  
• 负载均衡：如果模型切分不均匀，会导致某些NPU过载，而其他NPU闲置。
  
• 实现复杂性：需要在代码层面将模型和数据流切分为多个阶段，并设计高效的通信方案。
  

2.突破torch.nn.GroupNorm的限制


在处理VAE encoding时，虽然我们在序列维度进行了切分，但在处理groupnorm时，常规做法需要统计全序列的数值才可计算。但全序列groupnorm，计算时的激活tensor过大，单份显存~26.9G = 128 * 9 * 720 * 1088 * 16 * 2 / 1024 / 1024 / 1024，如果后续增加视频帧数或是视频分辨率，则将对NPU的64G显存上限是一个不小的挑战。
一种方案是对group维度进行切分，有2个关注点：

• 序列维度合并后才能进行group维度的切分，在此过程中是否会发生用于存储中间结果的tensor无法申请到显存的情况
  
• 对预训练的groupnorm权重进行提取，并新建对应的batchnorm算子，计算后，还需还原序列维度的切分状态
  

另一种方案，全序列groupnorm计算需传输高维张量。但从原理上，只需要获得序列切分后的统计值(低维)即可，单卡获取全量统计值后，可以处理自身的序列切分数据，其中有3个难点：

• 本质上是使用理论公式对groupnorm进行函数级别的分拆实现，涉及eps，有偏无偏参数，统计值精度等细节对齐，以保证与torch.nn.GroupNorm的误差可接受
  
• 因为是一个FusedOp-reverse的过程，会造成速度的损失，需考量可接受度
  
• 与第一种方案类似，需对预训练的groupnorm权重进行提取；另外在某些场景下，还需要实现backward代码，实现难度较高
  

3.分层ZERO3


Deepspeed 的 zero-3：是一种用于深度学习优化的技术。在分布式训练框架下，zero-3 将训练状态（包括权重、梯度和优化器状态）分布到不同的显卡上，以优化显存利用。与 zero-2 相比，它还对模型参数进行了分区，显存减少幅度与使用的 GPU 数量成正比。在TTV这种显存为瓶颈的场景，可以支持更大参数量的模型，提高batch的大小，优化训练效率。

参考文献

• https://arxiv.org/pdf/2205.14135
  
• https://gitee.com/ascend/cann-ops-adv/blob/master/docs/FlashAttentionScore.md
  
• https://gitee.com/ascend/cann-ops-adv/blob/master/docs/FlashAttentionScoreGrad.md#https://gitee.com/link?target=https%3A%2F%2Fcreativecommons.org%2Flicenses%2Fby%2F4.0%2Flegalcode
  
• https://gitee.com/ascend/pytorch/blob/master/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp#L240
  
• https://arxiv.org/abs/2205.05198
  
• https://arxiv.org/abs/2309.14509
  
• https://arxiv.org/abs/2403.10266
  
• https://www.alphaxiv.org/abs/2408.06072v1