AI 平台/部署

使用 NVIDIA TensorRT 优化基于 Transformer 的扩散模型以生成视频

先进的图像扩散模型需要数十秒才能处理单张图像。这使得视频扩散更具挑战性,需要大量计算资源和高昂成本。通过在搭载 NVIDIA TensorRT NVIDIA Hopper GPU 上利用最新的 FP8 量化功能,可以显著降低推理成本,并以更少的 GPU 为更多用户提供服务。虽然量化扩散器的部署可能比较复杂,但 TensorRT 背后的完整生态系统可以帮助克服这些挑战。

借助此方法,Adobe 将延迟降低了 60%,TCO 降低了近 40%,从而加快了推理速度并提高了响应速度。使用在由 Hopper GPU 加速的 Amazon Web Services (AWS) EC2 P5/P5en 上运行的 TensorRT 进行的优化部署,提高了可扩展性,以更少的 GPU 为更多用户提供服务。

本文将探讨为增强 Adobe Firefly 视频生成模型的性能而实施的策略和优化,重点关注降低延迟、降低成本和加速市场部署。

借助 AWS 上的 Adobe Firefly 和 NVIDIA TensorRT 革新创意 AI

借助 Firefly,用户可以在瞬间根据文本提示生成详细的图像,从而简化创意流程。例如,单个视频 diffusion 模型需要比单个图像 diffusion 模型更多的计算。

对快速、高效和可扩展的 AI 推理和训练解决方案的需求推动了 Firefly 和 Adobe 与 NVIDIA 持续合作的快速发展。TensorRT 为 Adobe 提供硬件加速和模型优化工具,以快速、大规模地部署其创新生成式模型,确保它们始终处于创意 AI 技术的前沿。

Project timeline from July to October 2024 showing milestones in video generation model development, including model training, quantization with TRT BF16 and FP8, final checkpoint, and private beta launch at Adobe Max.
图 1。从研究到私人测试,Adobe Firefly 视频生成模型在 2024 年的上市时间不到四个月

Adobe Firefly 发布

Adobe Firefly 发布 (2024 年 10 月) 非常精彩。这是 Adobe 有史以来最成功的测试版之一。这些数字不言自明:

  • 仅在第一个月就生成了超过 7000 万张图像
  • 迄今为止,Firefly 已为 超过 20 billion 个资产的创建 提供支持
  • 集成 Adobe Photoshop、Adobe Premiere Pro、Adobe Express 和 Adobe Illustrator 等 Adobe 创意套件

利用 TensorRT 实现高效部署

为了应对扩展扩散模型的挑战,Adobe 使用了高性能深度学习推理优化器 NVIDIA TensorRT。NVIDIA H100 GPU 上的最新 FP8 量化实现了以下功能:

  • 减少显存占用 :FP8 可显著降低显存带宽,同时加速 Tensor Core 运算
  • 节省推理成本 :相同工作负载所需的 GPU 更少
  • 无缝模型可移植性 :TensorRT 支持 PyTorch、TensorFlow 和 ONNX,提高了部署效率

TensorRT 可在 PyTorch 和 TensorFlow 等各种框架中优化和部署模型,使其成为 Adobe 用例的理想选择。优化过程涉及几个关键步骤:

第 1 步:ONNX 导出

Adobe 之所以选择 ONNX (Open Neural Network Exchange) ,是因为它具有通用性和易于导出的特点。这一决策允许在研究和部署之间无缝共享代码,从而消除了耗时的重新实现。

第 2 步:TensorRT 实现

该团队实施了 TensorRT,专注于使用 FP8 和 BF16 的混合精度。这种方法显著减少了权重和激活函数的内存占用,从而降低了内存带宽并加速了 Tensor Core 运算。FP8 格式可大幅减少权重和激活函数的内存占用。E4M3 FP8 格式可以表示为:

x = (-1) - s \cdot 2e^{-\text{bias}} \cdot (1 + f)

其中:

  • s 是符号位 ( 1 位)
  • e 是指数 ( 4 位)
  • f 是分数 ( 3 位)
  • 对于 E4M3,bias 为 7

这种格式允许的可表示值范围约为 1.52 × 10-2 到 4.48 × 102。之所以选择 E4M3 而不是 E5M2,是因为它允许更精细的精度,而不是更高的激活峰值。这种权衡最适合 前向推理 ,而反向传播可以从更大的值范围中受益。

第 3 步:量化技术

Adobe 使用 NVIDIA TensorRT Model Optimizer PyTorch API 实现了训练后量化。PyTorch API 允许使用现有的评估流程进行内置研究。虽然 TensorRT Model Optimizer 提供的 PyTorch FP8 仿真无法反映 FP8 执行的实际性能提升,但它可以在不需要网络导出的情况下快速进行质量评估。

Diagram of an AWS-based Adobe Firefly video generator AI pipeline. It starts with a prompt input and flows through different encoding modules: Text encoder for text-to-video, LDM encoder for image-to-video, and Latent encoder for video-to-video. The encoded input is processed by a Diffusion DiT model (inference stage) and then passed through a latent decoder and an upsampling module. Finally, video postprocessing runs on the CPU to generate a 5-second, 24 fps, 1080p video. Each pipeline stage is color-coded to indicate hardware usage: green for A100, dark green for H100, and blue for CPU.
图 2。在 AWS 上运行的 Adobe Firefly 视频生成器 AI 管道

使用 NVIDIA Nsight Deep Learning Designer 识别瓶颈

借助 NVIDIA Nsight Deep Learning Designer ,工程师查明了扩散管道中的关键瓶颈,包括:

  • Scaled Dot Product Attention (SDPA) 是导致过度延迟的主要计算瓶颈。
  • ONNX 分析允许映射内核执行时间,这表明高分辨率图像和视频扩散模型效率低下。
Screenshot of GPU profiling tool analyzing the baseline ONNX model. The top section shows a timeline view with 'Inference Stream 0' highlighting kernel executions, including 'gemm' and 'mha' layers. Below it are performance metrics such as Compute Warps in Flight, PCIe Throughput, SMs Active, Tensor Active, and VRAM Throughput. The bottom section displays a table of network metrics, listing layer names, inference times, precisions, and input dimensions. To the right, a bar graph visualizes the average latency per layer type, with 'kgen' and 'gemm' layers contributing most to latency.
图 3。GPU 分析工具分析基准 ONNX 模型

通过隔离这些性能问题,Adobe 团队对 Transformer 主干进行了微调,以提高执行速度并降低内存消耗。

使用量化 Diffusers 克服部署挑战

量化扩散器的部署可能非常复杂,需要仔细调整模型参数和量化设置。然而,TensorRT 背后的完整生态系统(包括 NVIDIA Deep Learning SDK 和 TensorRT Model Optimizer)帮助 Adobe 克服了这些挑战。

Adobe 的工程师实施了评估和提高量化质量的技术,包括分布分析和使用 TensorRT Model Optimizer 进行自动量化。

量化

量化使用缩放系数 S 将全精度浮点值 BF16 映射到 FP8 表示:

q = \text{clip}(\text{round}(Sx), q_{\text{min}}, q_{\text{max}})

其中:

  • q 是量化的 FP8 值
  • S 是每个 Tensor 或每个通道的规模系数
  • q_{\text{min}}, q_{\text{max}} 定义 FP8 动态范围,例如: 对于 E4M3:大约 [−448,448] 对于 E5M2:大约 [−57344,57344]
  • clip 可确保值保持在可表示的 FP8 范围内

去量化

要恢复近似的全精度值:

x' = q \cdot S

其中:

  • x' 是重建的 floating-point 值
  • q 是已存储的 FP8 值
  • S 与量化期间使用的比例系数相同

缩放系数选择

扩展系数 S 通常通过以下方式确定:

  • 基于 Max 的扩展: S = \frac{\max(|X|)}{q_{\max}}
  • 每个 Tensor 扩展 :整个 Tensor 的一个扩展

误差分析 (量化噪音)

量化误差通常建模为:

\epsilon = x - x' = x - q \cdot S

如果正确扩展,这将遵循 uniform distribution。

处理 FP8 格式

TensorRT 支持 E4M3 和 E5M2 FP8 格式:

  • E4M3 (1 个符号位、4 个指数位、3 个尾数位):在较小的动态范围内提供更高的精度。它可以表示大约在 ±[1.52×10⁻², 448] 范围内的值,包括 NaN。
  • E5M2 (1 个符号位、5 个指数位、2 个尾数位):以更低的精度提供更宽的动态范围。它可以表示范围大致为 ±[5.96×10⁻⁸, 57344] 的值,包括 ±inf 和 NaN。

选择取决于 precision 和 dynamic range 之间的权衡。

AI 工作负载的可扩展性和成本优势

Adobe Firefly 在 AWS 上的部署在优化性能和确保无缝可扩展性方面发挥了关键作用。通过利用 AWS 高性能云基础设施,该团队能够更大限度地提高效率、降低延迟并提高大规模 AI 工作负载的成本效益。

使用 TensorRT 优化部署后,扩散延迟降低了 60%,总体拥有成本降低了 40%,显著节省了成本,并提高了 Adobe 创意应用程序的可扩展性。通过减少扩散模型推理所需的计算资源,Firefly 能够用更少的 GPU 为更多用户提供服务,从而降低成本并提高效率。

This chart compares the runtime of a diffusion model backbone inference across different precision configurations. Using TensorRT with BF16 precision yields a 1.6x speedup over the PyTorch baseline, while combining BF16 and FP8 precision in TRT further improves performance to a 2.5x speedup. These optimizations demonstrate the significant acceleration achieved by leveraging lower precision formats in inference workloads.
图 4. 扩散主干推理性能。与 PyTorch 基准相比,搭载 BF16 和 FP8 的 NVIDIA TensorRT 可将运行速度提高 2.5 倍

未来步骤

优化扩散模型部署对于让更多用户访问这些强大的模型至关重要。随着 Adobe 不断突破创意 AI 的界限,从 Firefly 的开发和部署中汲取的经验教训将塑造未来的创新。快速开发、战略技术决策和不懈优化的结合为生成式 AI 领域树立了新的标准。NVIDIA 很高兴能继续与 Adobe 和 AWS 合作,利用扩散模型和 Deep Learning 不断突破极限。

有关更多信息,请查看 NVIDIA TensorRT 文档 ,并观看 NVIDIA GTC 会议: Quantize Large Transformer Diffusion Models 以改善端到端延迟并节省推理成本

 

标签