福州网页建站维护有哪些,网站建设空间申请,网站轮播图片制作,app定制价格是多少1. xDiT 简介
xDiT 是一个为大规模多 GPU 集群上的 Diffusion Transformers#xff08;DiTs#xff09;设计的可扩展推理引擎。它提供了一套高效的并行方法和 GPU 内核加速技术#xff0c;以满足实时推理需求。
1.1 DiT 和 LLM
DiT#xff08;Diffusion TransformersDiTs设计的可扩展推理引擎。它提供了一套高效的并行方法和 GPU 内核加速技术以满足实时推理需求。
1.1 DiT 和 LLM
DiTDiffusion Transformers是文生图与文生视频的核心网络结构。
DiT 和 LLM 推理任务特点不一样这导致二者软件设计哲学也有差异
1. LLM 有 Prefill 和 Decode 两阶段分别是计算密集和访存密集的而 DiT 的计算和 Prefill 相似是计算密集的。所以 Batching、Paged Attention 等为了增加 Decode 阶段计算密度的设计在 DiT 中不需要其实也没有 KVCache 不断 Append 增长的现象。
2. LLM 模型很大而序列长度有限比如 ShareGPT 的任务生成顶多几百个 token。而 DiT 正好反过来模型不大推理时序列长度很长甚至 Million 级别。这导致张量并行这种 vLLM 中第一优先级的并行方法在 DiT 中性能很低。
3. LLM 模型大多收敛到微调 Llama 架构而 DiT 架构则呈现出较大的差异性。例如SD3、Latte 等网络与 Meta DiT 论文中的 DiT Block 在算子设计上存在显著差异且 DiT Block 之间的连接方式也各不相同常采用类似 U-Net 的 U 型结构进行链接。因此DiT 模型的代码相对 LLM 不够简洁往往涉及很多文件和函数对其进行模型层面的改造也更具挑战性。 1.2 背景与挑战
DiT 模型在生成长视觉序列时的推理延迟显著增加例如Sora 模型的推理延迟可达数分钟单个 GPU 无法满足实际应用的延迟要求。因此DiT 模型的部署必然是多卡并行处理。Attention 机制导致的计算时间与序列长度呈二次增长为了达到同样延迟指标GPU 数应该随着序列长度平方项增长的所有部署 Sore 很大概率需要多机多卡大集群执行单个视频生成DiT 的并行推理扩展性至关重要。尽管大语言模型LLMs中常用的并行化技术如张量并行和序列并行通信量和激活参数正相关在扩散模型中由于大激活小参数特点这些技术效率不高。通信成本往往超过了并行计算的好处。因此DiT 模型的部署仍然需要配备高带宽互连如 NVLinkRDMA的 GPU 集群。
随着 DiTs 中输入上下文长度的增加Attention 机制的计算需求呈二次方增长因此多 GPU 和多机器部署对于满足在线服务的实时需求至关重要。 1.3 Diffusion Model 推理原理和特性 尽管 DiT 和 LLM 的结构相似都是 Transformers 架构看起来工程挑战类似。
扩散模型的训练过程给一个图片经过很多步骤每一步加噪声最后变成一个全是噪音的图片。训练过程就是预测一个 Noise Predictor 的监督学习任务输入是中间状态图片、输入指令和 timestep 来 ground truth 就是该 timestep 的噪声。
扩散模型的推理过程给一个噪声通过 Noise Predictor 来通过多次去噪最后变成一个有意义的图片。每一次去噪就是一个 Diffusion Step。
通过描述可知扩散模型推理和 LLM 的有一个很大的差异点就是扩散模型是重复计算很多相同的 Diffusion Step而且连续的 Diffusion Step 之间输入数据和激活状态之间存在的高度相似性称之为输入时间冗余Input Temporal Redundancy。这种冗余性表明在进行相邻时间 Diffusion Step 去噪时可以利用前一步骤的过期 Activation 来顶替当前步骤生成的 Activation 做一些计算可能结果也不会差太多。
在 DiT 推理中鉴于结构的相似性可以用 LLM 中非常成熟的张量并行TP和序列并行SP也可以使用一些利用扩散模型 Input Temporal Redundancy 特性的并行方法比如 DistriFusion 和 PipeFusion。 1.4 xDiT Overview 正如 vLLM 让 transformers 具备 Batching 的能力xDiT 让 diffusers 有并行的能力。考虑到 DiT 的特点首先提出了一系列创新的并行方式并可以相互混合。更重要的是xDiT 提供了一套优雅的开发接口针对性解决了 DiT 模型更改难度高的问题这套开发接口尽可能可能复用 diffusers 的现有逻辑开发者通过一些 wrapper 实现复杂的混合并行实现高效地扩展方法。将之前的 PipeFusion 项目升级成了通用的 DiT 推理引擎命名为 xDiT。向 vLLM 致敬X表示Scalable希望 xDiT 成为 DiT 推理领域的 vLLM
vLLMhuggingface transformersBatchingxDiThuggingface diffusersParallel
xDiT 是一个为大规模多 GPU 集群上的 Diffusion TransformersDiTs设计的可扩展推理引擎。它提供了一套高效的并行方法和 GPU 内核加速技术以满足实时推理需求。 xDiT 有如下核心能力
第一DiT 主干网络混合并行xDiT 支持任意混合四种基础并行策略方法。基础并行方法有 Pipefusion ParallelTeraPipe 方式流水线并行方法在弱互联比如 pcie/以太网的网络硬件上有优势比如以太网互联的多机多卡情境下效率远高于其他方法 Sequence Parallel混合序列并行使用 ring ulysses sequence parallel 混合方式 USP Data Parallel在输入多个 prompt 或单个 prompt 生成多张图片时在 image 间并行处理 CFG Parallel, a.k.a Split Batch在模型使用 classifier free guidanceCFG时可以开启并行度恒定为 2。
xDiT 中这四种并行方式可以以任何形式排列组合混合策略在通信量和通信拓扑两方面可以取得最佳效果使 xDiT 可达到近似线性的扩展性。
第二Parallel VAE针对扩散模型后处理的解码 VAE 模块在高分辨率图像生成时 OOM 问题xDiT 实现了 Patch Parallel 版本的 VAE。
第三简单灵活的开发接口xDiT 轻松帮助用户支持新的 DiT 模型。它尽可能复用开源社区 diffuser完成对模型推理的并行化改造。这套接口也是让 xDiT 从之前的不同并行方法集合变成一个灵活开发工具。通过 xDiT 接口实测 15 分钟可以让 Pixart-Sigma 模型并行起来 1.5 xDit 并行方法
多GPU并行方法 PipeFusion利用扩散模型的特性通过位移Patch实现Patch级流水线并行Unified Sequence Parallel (USP)将 DeepSpeed-Ulysses 和 Ring-Attention 结合起来的并行方法Hybrid Parallel混合并行可配置多种并行方法以优化通信模式适应底层网络硬件CFG Parallel也称为 Split Batch在使用无分类器指导CFG时激活保持恒定的并行度为2Parallel VAE并行 VAE单GPU加速方法 编译加速xDiT 提供了 GPU 内核加速以提高推理速度torch.compile 和 onediff 这些编译加速与并行化方法结合使用 torch.compilePyTorch 的一个特性利用自动图优化、操作融合以及内核选择等技术加速PyTorch模型的执行效率。在 xDiT 中torch.compile 用于提升 GPU 上的单 GPU 性能onediff用于 PyTorch 的自动微分编译器通过优化编译过程来提高模型的执行速度和效率DiTFastAttnxDiT还为单GPU加速提供了DiTFastAttn它可以通过利用扩散模型不同步骤之间的冗余来降低注意力层的计算成本
xDiT 支持多种 Diffusion Transformers 模型包括 CogVideoX、Flux.1、HunyuanDiT、SD3、Pixart、Latte 等。
针对以上不同的模型xDiT目前实现了不同的方式对它进行并行推理加速。例如Latte是文生视频模型xDiT目前实现了USP方式对它进行并行推理加速PipeFusion还在开发中使用混合序列并行(ulysses_degree2, ring_degree4)可以获得最佳性能。CogVideo 是一个文本到视频的模型xDiT 目前整合了 USP 技术包括 Ulysses 注意力和 Ring 注意力和 CFG 并行来提高推理速度同时 PipeFusion 的工作正在开发中。
在使用不同GPU数目时最佳的并行方案都是不同的这说明了多种并行和混合并行的重要性。 例如HunyuanDiT最佳的并行策略在不同GPU规模时分别是在2个GPU上使用ulysses_degree2在4个GPU上使用cfg_parallel2, ulysses_degree2在8个GPU上使用cfg_parallel2, pipefusion_parallel4。torch.compile带来的加速效果也很可观。stable-diffusion-3最佳的并行策略在不同GPU规模时分别是在2个GPU上使用cfg_parallel2在4个GPU上使用cfg_parallel2, pipefusion_parallel2在8个GPU上使用cfg_parallel2, pipefusion_parallel4。 2. Parallel Methods
2.1 PipeFusion
流水线并行
流水线并行属于层间并行对模型不同的 Transformer 层间进行分割模型被分割并分布不同的设备上每一个设备只保存模型的一部分参数在各层之间进行并行计算。 将模型的不同层放置到不同的计算设备可以降低单个计算设备的显存消耗。 如下图所示例如把模型分成四个模型层分别放置到四个不同的计算设备即第 1 层放置到GPU 0第 2 层放置到GPU 1第 3 层放置到GPU 2第 4 层放置到GPU 3。
前向计算过程中输入数据首先在GPU 0 上通过第 1 层的计算得到中间结果并将中间结果传输到GPU 1然后在GPU 1 上计算得到第 2 层的输出并将模型第 2 层的输出结果传输到GPU 2在GPU 2 上计算得到第 3 层的输出将模型第 3 层的输出结果传输到GPU 3GPU 3 经由最后一层的计算最后得到前向计算结果。反向传播过程类似。
以上是训练过程推理过程中没有反向传播只有前向传播最后一层的输出结果即为整个模型的推理输出。 流水线并行训练的一个明显缺点是训练设备容易出现空闲状态因为后一个阶段需要等待前一个阶段执行完毕导致计算资源的浪费。
PipeFusion 论文地址https://arxiv.org/pdf/2405.14430 PipeFusion 将输入图像分割成 M 个不重叠的 Patch并将 DiT 网络均匀分成 N 个阶段并按顺序分配到 N 个计算设备上。 例如下图表示在 N 4 和 M 4 的情况将输入图像分割成了4个不重叠的PatchP0、P1、P2、P3并且将 DiT 网络均匀分成 4 个阶段step0、1、2、3反向扩散过程推理的流水线工作流程横轴分为两个时间步时间步 T深灰色和时间步 T1浅灰色回顾SD的原理可知反向扩散从T1步开始纵轴 device 0 到 device 3 表示并行计算的 GPU。
流水线处理的概念是在一个设备完成当前任务并将结果传递给下一个设备时它可以同时开始处理下一个Patch。例如GPU0先处理patch P0完成后将其传递给GPU1自己开始处理patch P1此时GPU1继续处理从GPU0接收的Patch0以此类推。 不同PatchP0、P1、P2、P3的计算在不同设备上同时进行在每个时间步内不同设备处理不同的任务整个计算过程是并行化的。 时间冗余与流水线优化在并行计算的流水线架构中存在一个流水线初始化阶段即当开始计算时设备之间需要传递数据以保持同步。这个过程会引入一些“空隙”或“延迟”因为每个设备都需要等待上一个设备完成部分计算并传递给它。
为了克服这一延迟采用了输入时间冗余的策略即每个设备使用前一个时间步的旧输出作为计算的上下文避免了当前时间步的完全等待。这使得流水线能够尽早开始而不必等待所有输出完全计算完成。
流水线的有效计算比例 M输入图像被划分为的不重叠patches数量 NDiT 网络被划分为的阶段stages数量 S扩散时间步的数量扩散模型通常需要多步迭代来逐渐去噪 M⋅S表示总的计算量M⋅S 表示所有Patch在所有时间步中需要计算次数 N-1流水线前面几步需要填补流水线中的空隙
假设 M N 4扩散时间步 S 50带入公式 说明流水线的有效计算比例为98.5%当扩散时间步S增加时这个空隙占总计算时间的比例变小整体效率提升。 2.2 USP: Unified Sequence Parallelism
序列并行 注意力机制的基本工作原理Q 与 K 的转置进行点积计算生成一个大小为 S×S 的注意力矩阵注意力矩阵经过softmax操作后与 V 相乘生成最终输出这整个过程的内存复杂度是与序列长度呈二次方关系
近年来随着生成性AI模型上下文长度的不断增长例如Claude在大型语言模型LLMs中将序列长度扩展到了100K个token而OpenAI的GPT-4则将上下文长度扩展到了128K个token为了能够训练超长上下文通常需要使用一些复杂的并行策略。 序列并行Sequence ParallelismSP是一种将输入张量的序列维度分到多个计算设备上的技术将输入序列分割成多个子序列并将每个子序列输入到其相应的设备即 GPU中。在序列并行中模型通常不会被直接切分而是每个设备都运行同一模型的一个完整副本但输入数据是不同的序列片段。 上图展示了序列并行Sequence Parallelism的工作原理。在这种并行方式中微批次Micro Batches被沿着序列维度均匀地分布到不同设备上每个设备负责处理不同的序列片段这样能够减少单个设备的内存开销同时加速长序列的处理。每层的输出在设备之间同步传递。 上图展示了如何将序列数据划分到不同的设备上并在每个设备上并行地执行注意力机制
输入序列被拆分成不同的部分Qi, Ki, Vi等Q、K、V分别代表查询Query、键Key、值Value序列中的每个元素的Q、K、V向量被分配到不同的设备上分别在不同设备上进行计算图中有多个设备Device i-1, i, i1, i2等每个设备负责一部分序列的计算注意力计算每个设备需要访问其他序列元素的K和V来计算注意力因此每个设备都会与相邻设备交换KV信息每个设备计算其对应序列元素的注意力输出
通过将序列数据分割并行处理多个设备可以同时处理不同部分的数据提高模型的性能和效率。 综上序列并行的关键在于它能够有效地分配和管理大规模输入序列的计算任务让模型能够处理更长的序列不会受到单个设备内存限制的束缚。序列并行已成为训练和推理更长序列的一种有效方法。DeepSpeed-Ulysses和Ring-Attention标志着序列并行技术的成熟这是目前两种最重要的序列并行方法。
DeepSpeed-Ulysses 论文地址https://arxiv.org/pdf/2309.14509 DeepSpeed-Ulysses 是微软提出的简单、易用且高效用于支持具有极长序列长度的高效可扩展LLM训练的方法
多头注意力机制计算过程 输入向量 X输入序列每个元素的维度为 [N, d]其中 N 是序列长度d 是隐藏层的维度 线性变换输入 X 通过三个不同的线性变换使用权重矩阵$ W_Q, W_K, W_V$生成查询Q、键K和值V这三个矩阵的维度都是 [N, d]。 分头处理生成的 Q, K, V 各自被分成多个“头”这里的头数 $h_c$ 4每个头处理输入数据的一个子集这可以帮助模型从多个子空间中并行学习信息每个头的维度是 [N, $h_s$]其中 hs 是每个头的大小等于 $d/h_c$ 点积注意力 计算点积对于每个头计算查询 Q 和键 K’K 的转置的点积 缩放点积通常这个点积还会除以一个缩放因子 $h_s$ 的平方根 应用 softmax 函数对缩放后的点积结果应用 softmax 函数生成一个注意力权重矩阵 S 生成输出 加权和使用注意力权重矩阵 S 对值 V 进行加权求和得到每个头的输出 $P_h$ 连接将所有头的输出连接起来形成维度为 [N, d] 的矩阵 P最终线性变换P 通过另一个权重矩阵$W_0$ 进行线性变换得到最终的多头注意力输出 DeepSpeed- Ulysses 计算过程 输入向量 X输入数据 X 的维度为 [N, d]其中 N 是输入序列长度d 是隐藏层的维度输入数据被分割成 P 个部分每部分的维度为 [N/P, d]分别送到 P 个处理器GPU这里假设 P hc 4即 GPU 数量等于注意力头数 线性变换每个GPU上的数据通过三个不同的[d, d]维度的线性变换矩阵$W_Q, W_K, W_V$生成局部的查询Q、键K和值V all-to-all 通信 发送每个GPU需要将其计算得到的局部 Q、$K^T$和 V 发送到其他所有GPU每个GPU都有完整的Q、$K^T$和 V 接收每个GPU将接收其他GPU发送过来的Q、$K^T$和 V 分片组合成全局 Q、$K^T$和 V 计算注意力 计算点积每个GPU并行计算不同的注意力头的 Q 和 $K^T$的点积 缩放点积除以一个缩放因子 hs 的平方根 应用 softmax 函数对缩放后的点积结果应用 softmax 函数生成局部的注意力权重矩阵 S 生成输出 加权和使用局部的注意力权重矩阵 S 对局部的值 V 进行加权求和得到每个头的输出 $P_h$ all-to-all 通信再次使用 all-to-all 通信将所有处理器上的$P_h$聚集起来形成完整的输出 P 最终线性变换聚集后的输出 P 通过权重矩阵 $W_0$ 进行线性变换得到最终的多头注意力输出
对于每一个GPU送入处理的序列长度只是 N/P这样大大减少了单个GPU的运算量但同时因为 All-To-All 通信方式的存在做 Attention 时还是考虑了整体的长度。DeepSpeed-Ulysses 方法的特点是当序列长度和计算设备成比例增加时通信量保持不变。
局限性每个GPU并行计算不同的注意力头Ulysses的SP并行度受限于注意力头数量并行度不能超过头的数量
Ring-Attention GitHubhttps://github.com/gpu-mode/ring-attention Ring Attention顾名思义环形注意力是另一种序列并行方法。
数据块的计算顺序可以任意在环形结构中各个数据块如Q、K、V的计算顺序可以灵活安排不必严格按照序列顺序将QKV序列分割到N个主机上序列中的查询Q、键K和值V被分割并分配到多个处理单元图中例为四个GPU构成环状结构交换KV数据每个GPU不仅计算分配到的数据块各个GPU形成一个环形网络其中每个GPU会向相邻的GPU发送和接收K和V数据。这样每个GPU都能够获取到整个序列的K和V信息来计算注意力所需的K和V块不在本GPU则使用点对点P2P通信从其他GPU获取单次循环完成所有KV部分的交换通过环形网络的一次循环每个GPU都能看到并处理所有的K和V数据长序列处理的零开销同时进行计算和通信在处理较长序列时没有额外开销这意味着随着序列长度的增加系统的效率不会下降 Ring Attention Algorithm 详细算法步骤 输入和初始化 输入序列 x 主机数 Nh指的是参与计算的设备数量GPU这些设备将形成一个环状结构来交换数据 分割输入序列输入序列被分割成 Nh 个块每个块由一个GPU处理 计算查询、键、值Q, K, V每个GPU计算其对应块的查询Q、键K和值V 循环处理每个Transformer层 对于Transformer中的每一层执行以下步骤 发送和接收K和V每个GPU将其K和V块发送给下一个GPU并从上一个GPU接收K和V块这个过程在所有GPU中并发执行 计算注意力每个GPU使用本地的Q、K、V块来计算注意力 计算前馈网络使用注意力输出每个GPU并发地计算feedforword前馈网络 结束经过所有Transformer层的处理后算法结束 USP 论文地址https://arxiv.org/pdf/2405.07719 GitHubhttps://github.com/feifeibear/long-context-attention?tabreadme-ov-file USP(Unified Sequence Parallel) 结合了DeepSpeed-Ulysses和Ring-Attention两种序列并行方法提出了“2D混合序列并行”的概念这种混合并行方式不仅继承了两种方法各自的优点还克服了它们各自的局限性。通过在序列维度上的混合并行策略允许在不同的计算设备上分别运行DeepSpeed-Ulysses和Ring-Attention从而优化了通信模式和计算效率。对比主流的DeepSpeed-Ulysses和Ring-Attention序列并行方式USP在DiT场景下生图的性能提升到24%在LLAMA2-7B场景的性能提升2倍以上。
USP方法核心思想是将计算过程分为两个正交的过程组SP-Ring过程组和SP-Ulysses过程组形成二维网格结构它的算法流程包括
1) 使用DeepSpeed-Ulysses方法的all-to-all交换query、key、value
2) 各组内进行独立的Ring-Attention计算
3) all-to-all通信使各组具有完整的Head Output结果
算法流程如下 对image的切分逻辑如下 即在外层维度切分ring blocks在ring block的内部切分ulysses block每个device持有对应的ulysses block在进行计算时同一个ulysses group所属的devices先进行all2all相当于沿hidden dim方向切分image hidden state保证每个device上有attention中一部分heads所需要的属于目前ulysses group的所有kv随后属于同一ring group的devices进行ring attention通过online softmax的特性完成对整个attention结果的计算值得注意的是ulysses的通信中需要对query,key,value三者均进行all2all而ring attention中仅有kv需要进行p2p通信。并且在ring attention中kv通信和attn技术是可重叠的。因此在低带宽连接的设备中采用较高的ring degree与较低的ulysses degree通常能取得更好的加速效果 将上述算法展示成图的方式如下图所示 输入数据分布输入数据按序列维度拆分到4个GPU上每个GPU上都有部分的序列并且每个GPU上处理的注意力头head也是不一样的。比如GPU0和GPU2负责head0GPU1和GPU3负责head1相当于GPU0和GPU1之间进行Ulysess并行GPU0和GPU2之间进行Ring Attention并行 AlltoAll q k v在进行多头自注意力计算之前需要对query (q)、key (k)、value (v) 进行AlltoAll操作。在这个过程中GPU0和GPU1之间会交换部分序列数据GPU2和GPU3之间也会交换。这样GPU0和GPU1都拥有相同的部分序列但它们分别处理不同的headGPU2和GPU3同理 Ring Attention计算这一阶段的目的是在多个GPU之间计算Attention分数。Ring Attention是通过在GPU0和GPU2之间、GPU1和GPU3之间的通信完成的从图中可以看出GPU0和GPU2具有所有数据GPU0有一半数据GPU2有另一半数据相当于它们共享的序列数据是完整的但是只有head0数据所以它们会一起计算head0的Attention结果同样GPU1和GPU3有完整token数据但是只有head1的数据故GPU1和GPU3一起计算head1的attention结果 AlltoAll输出在完成Ring Attention计算之后再次执行AlltoAll操作将不同GPU上的head结果进行交换。GPU0和GPU1交换它们的计算结果GPU0上就有了完整序列的head0和head1的结果类似地GPU2和GPU3也完成结果交换每张卡上都有完整的head和序列的计算结果。例如序列0-7拆分成是0-3和4-7那么GPU0有0-3且head为0的计算结果GPU1有0-3且head为1的计算结果GPU2有4-7且head为0的计算结果GPU3有4-7且head为1的计算结果此时GPU0和GPU1的结果进行AlltoAllGPU2和GPU3的结果进行AlltoAll所以卡均获得完整head的数据且序列完全拆分 2.3 Parallel VAE 项目代码https://github.com/xdit-project/DistVAE 在高分辨率图片生成中最后用于从潜在空间到像素空间转换的VAE开销很大是整个生成过程的显存尖峰Memory Spike在主干网络生成结束后添加了并行VAE的实现以支持高分辨率图片解码提出DistVAE避免OOM并行高效地处理高分辨率图像 DistVAE 结合了两种关键策略 序列并行将潜在空间中的特征映射划分为多个Patch并在不同设备上执行序列并行VAE解码将中间激活所需的峰值内存减少到1/N其中N是所使用的设备数量 Chunked Input Processing分块输入处理与MIT patch conv类似将输入特征图分割成块并将其顺序输入卷积算子这种方法最大限度地减少了临时内存消耗
通过协同这两种方法极大地扩展了VAE解码的能力。成功处理了高达 10240px 的图像分辨率与默认 VAE 实现相比提升了 11 倍。
