Grace Hopper Superchip Architecture and Performance Optimizations for AI Applications

Matthias Jouanneaux, Vishal Mehta, NVIDIA DevTech Compute | GPU Technology Conference/March 20th, 2024

目录

• NVIDIA Grace Hopper 超级芯片
• Grace Hopper 性能预览
• Grace Hopper CUDA 特性
• 在 Grace Hopper 上优化应用
• 微调应用
• 大语言模型推理
• 推荐系统工作流
• 数据库查询
• 结论
• GTC'24 核心性能优化技术

NVIDIA Grace Hopper 超级芯片

NVIDIA Grace Hopper 超级芯片的设计允许 GPU 以 CPU 内存的速度访问 CPU 内存。

上图展示了 Grace Hopper 超级芯片的架构，其核心组件包括：
- Grace CPU: 配备高达 480 GB 的 LPDDR5X 内存，提供最高 512 GB/s 的带宽。
- Hopper GPU: 配备 96 GB HBM3 或 144 GB HBM3e 内存，提供最高 4.9 TB/s 的带宽。
- NVLINK C2C: 一种 900 GB/s 的芯片间互联技术，用于连接 Grace CPU 和 Hopper GPU。
- 高速 I/O: 包括 4x 16x PCIe-5 接口，提供 512 GB/s 的带宽。
- NVLINK 网络: 通过 18x NVLINK 4 接口（900 GB/s）支持硬件一致性，可扩展至多达 256 个 GPU。

Grace Hopper 性能预览

Grace Hopper 通过提升 GPU 利用率，显著改善了 AI 应用的性能。

• AI 推理应用 (AI Inference Applications):

  • 带持久性 KV 缓存的推理 (Inference with Persistent KV Cache): 吞吐量提升至 x86+H100 80G 的 1.6 倍。
  • 带模型流的推理 (Inference with Model Streaming): 吞吐量提升至 4.9 倍。
  • 推荐系统工作流 (Recommender Workflows): 吞吐量提升至 1.2 倍。

• AI 微调应用 (AI Fine Tuning Applications):

  • 监督式微调 10B 模型 (Supervised Fine Tuning 10B MoE, 8 GPUs): 吞吐量提升至 1.3 倍。
  • Llama 70B LoRA 微调 (Llama 70B LoRA Fine-Tuning, 8 GPUs vs 4 GH200): 吞吐量提升至 1.5 倍。

• 数据库应用 (DataBase Applications):

  • 向量数据库搜索 CAGRA-Q (VectorDB search CAGRA-Q): 吞吐量提升至 6.8 倍。
  • 数据库查询 (Database Queries): 吞吐量提升至 6.6 倍。

注：GH200: Grace CPU 480GB LPDDR5, H100 96GB HBM3. DGX H100: H100 80GB HBM3.

Grace Hopper CUDA 特性

内存分配器 (Memory Allocators)

仅设备内存管理 (cudaMalloc)

无论是标准的 PCIe Hopper 配置还是 Grace Hopper 配置，使用 cudaMalloc 分配的 GPU 内存都无法被 CPU直接访问。CPU 侧的指针无法解引用以访问 GPU 内存中的数据。

CUDA 统一内存 (cudaMallocManaged)

对于使用 cudaMallocManaged 分配的 CUDA 统一内存，PCIe Hopper 和 Grace Hopper 的行为是相同的。当 CPU 或 GPU 首次访问数据时，会发生页错误（Page fault），并将相应的内存页迁移到访问方。

Grace Hopper 统一内存 (Unified Memory)

Grace Hopper 通过地址转换服务（Address Translation Service, ATS）实现真正的统一内存。

系统页表（System Page Table）可以将 CPU 分配的内存地址转换为 CPU 或 GPU 的物理地址。这使得 Grace CPU 和 Hopper GPU 能够通过 NVLINK C2C 互联，实现对 CPU 驻留（CPU-resident）和 GPU 驻留（GPU-resident）内存的直接访问或远程访问。

系统分配内存 (malloc/mmap)

• PCIe Hopper: GPU 无法直接访问 CPU 内存，反之亦然。需要异构内存管理（Heterogeneous Memory Management, HMM）来通过页错误和迁移模拟此行为。
• Grace Hopper: CPU 和 GPU 都可以通过 C2C 互联直接访问对方的内存。使用 malloc() 分配的内存，在首次被 CPU 或 GPU 触碰时填充，或被迁移。

Grace Hopper 优化

统一内存的自动迁移 (CPU -> GPU)

对于最初位于 CPU 物理内存中的数据页（Page A），如果 Hopper GPU 对其进行频繁访问，系统会自动将该页迁移到 GPU 物理内存中，以提升后续访问的性能。而对于不频繁的访问，则会通过 NVLINK C2C 进行远程访问。

CUDA 内存调优 API

在 Grace Hopper 上，CUDA 内存调优 API 可用于 cudaMallocManaged 和系统分配的内存。

• Grace Hopper 批量传输 (Bulk Transfers): 使用 cudaMemPrefetchAsync 可以将 CPU 内存中的数据块异步预取到 GPU 内存中，适用于批量数据传输场景。
• Grace Hopper 直接访问 (Direct Access): 使用 cudaMemAdvise 并设置 cudaPreferredLocation，可以建议运行时将某块内存的首选位置设为 GPU，从而实现从 GPU 对 CPU 内存的直接、高速访问，而无需显式迁移。

NVLink-C2C 性能

cudaMemcpy 在 Grace Hopper 上的性能表现优于传统 x86 系统。

上图展示了 MGX Grace Hopper 480GB CPU 内存的 cudaMemcpy 性能：
- 固定内存 (Pinned Memory):
- x86+H100 H2D (主机到设备): 55.67 GB/s
- x86+H100 D2H (设备到主机): 54.42 GB/s
- Grace Hopper H2D: 381.7 GB/s
- Grace Hopper D2H: 297.7 GB/s

• 可分页内存 (Pageable Memory):
  • x86+H100 H2D: 9.5 GB/s
  • x86+H100 D2H: 14.5 GB/s
  • Grace Hopper H2D: 381 GB/s
  • Grace Hopper D2H: 295.8 GB/s

Grace Hopper 凭借 NVLink-C2C，在主机与设备间的数据传输带宽上实现了巨大飞跃，无论是固定内存还是可分页内存。

在 Grace Hopper 上优化应用

关键优化技术

针对 AI 应用，Grace Hopper 的关键优化技术主要有两种：

1. 异步批量传输 (Asynchronous Bulk cudaMemcpy):
   将 CPU 内存作为 NVLink C2C 上的扩展，通过异步卸载能力传输激活值（activations）和参数（parameters）。
2. 直接内存访问 (Direct Memory Access):
   将 CPU 内存显式地用作一个高带宽、大容量的缓存，通过 NVLink C2C 进行直接内存访问。

微调应用

有监督微调 (Supervised Fine Tuning, SFT)

感谢：Matthias Langer (DevTech Compute APAC)

概述

有监督微调是使用 Grace Hopper 统一内存、PyTorch 和 Megatron 进行的一种微调方法。其基本流程是：首先，使用网络规模的数据（Web-scale data）预训练一个大型语言模型（Pre-trained LLM）；然后，利用领域特定数据或自定义知识库（Domain-specific data / custom knowledge base）对该预训练模型进行微调，从而得到一个针对特定任务或领域的微调后的大型语言模型（Fine-tuned LLM）。

专家混合模型 (Mixture of Experts, MoE)

• 每个 Transformer 层由于包含多个前馈（Feed Forward）专家层，因此内存密集度很高。
• 持续预训练（Continual Pre-training）涉及使用新的数据集来增强已训练的模型。

下图展示了 LLM 模型中 Transformer 层的结构，其中包含一个注意力（Attention）模块和一个专家混合（Mixture of Experts）模块。MoE 模块通过 MoE 路由（MoE Routing）将输入分配给不同的专家（Expert 1, ..., Expert N），然后对专家的输出进行加权求和（Weighted Sum）。

技术：激活检查点 (Activation Checkpointing)

为了更有效地利用内存，可以采用重新计算激活值的技术。

• 问题：中间张量（Intermediate tensors）具有很高的内存需求。
• 解决方案：在前向传播过程中不存储所有中间激活值，而在反向传播过程中需要时再重新计算它们。
• 权衡：这种重新计算的方法虽然内存开销低，但会导致冗余计算和更高的能耗。

技术：异步激活卸载 (Asynchronous Activation Offloading)

这是一种将激活值异步卸载到 CPU 内存的技术，以减少 GPU 内存的占用。

• 前向传播 (Forward Pass)：在模型逐层进行前向计算时，生成的中间张量（激活值）会被异步地从 GPU 内存移动到 CPU 内存。通常，体积较小的张量可能会被保留在 GPU 上，而较大的张量则被卸载。
• 反向传播 (Backward Pass)：当模型进行反向传播计算梯度时，之前卸载到 CPU 内存的相应中间张量会被异步地回调（Recall）到 GPU 内存中，以供梯度计算使用。

在 Megatron LM 中实现异步卸载

• 在 Megatron LM 中的实现会暴露用于卸载优化的运行时参数。

• 关键参数包括：

  • --cpu-offload-activations
  • --cpu-offload-asynchronous
  • --cpu-offload-size-threshold 67108864

• 下图展示了一个约 10B 参数的 MoE 模型（MegatronLM 配置，用于 8 个 H100 GPU 系统）的临时张量分布。总大小为 251.2 GB，其中大小在 1000~1999 MB 之间的张量占了总量的 54.1%。

性能分析与对比

• 核心发现：在 Grace Hopper 平台上，激活卸载比在 x86 平台上更高效。

• 性能对比：

  • 在 Grace Hopper 上，卸载操作可以被计算掩盖（hidden behind compute），从而实现计算和数据传输的并行。
  • 在 x86+H100 平台上，卸载操作主要受数据传输（D2H/H2D）的限制。

• 结果：如下图所示，对于一个 10B MoE 模型的有监督微调，使用 Grace Hopper 的卸载（offload GraceHopper）方案比在 x86+H100 上使用重新计算（recompute x86+H100）的方案快 25%。

LoRA 微调 (LoRA Fine Tuning)

使用 LoRA 微调 Llama 2 70B

本节介绍使用 Grace Hopper 统一内存、PyTorch 和 Huggingface 对 Llama 2 70B 模型进行 LoRA 微调。

• 模型与设置:

  • Llama 2 70B，全 16 位精度（仅参数大小为 140GB）。
  • 采用 PEFT（Parameter-Efficient Fine-Tuning）中的低秩适应（Low-rank adaptation, LoRA）技术。
  • 减少可训练参数数量。
  • 减少优化器的内存占用。

• 数据集: SAMSum 数据集，包含类似即时通讯的对话。

下图展示了使用 Grace Hopper Unified Memory、PyTorch 和 Huggingface 进行微调的性能。该图比较了不同配置下的微调时间与吞吐量，所有数据均已对 DGX H100 8 GPU 进行归一化处理。图表中，越靠左上方性能越好。

• Grace Hopper 1x: 训练时间延长 1.36 倍，但每 GPU 吞吐量提高 1.46 倍。
• Grace Hopper 4x (IB400): 训练时间延长 4.7 倍，每 GPU 吞吐量提高 1.67 倍。

实验设置：7 个 epoch，每个 epoch 包含 368 个批次，全局批次大小为 40。

内存占用分析

• 任何给定的 Transformer 层仅使用总内存的一小部分（Llama 2 有 80 层，每层约占 1.25% 的内存）。
• 实现单 GPU 微调需要同时进行参数卸载和激活卸载。

下表对比了使用和不使用 LoRA 及检查点技术时，单个层和整个模型的参数、优化器状态以及激活值的内存占用（单位：GB）。

• 参数和优化器状态: LoRA 将整个模型的内存需求从 560 GB 降低到 140 GB。
• 激活值: 检查点技术将整个模型的激活值内存需求从 112 GB 降低到 6.7 GB。

LoRA 微调策略与技术

对 Llama 2 70B 进行 LoRA 微调的核心策略是利用快速且大容量的 CPU 内存。

• 主要技术:
  • 使用 PyTorch 的前向/后向钩子（forward/backward hooks）进行参数卸载。
  • 使用 PyTorch 的 pack/unpack 激活钩子进行激活卸载，并使用后向钩子进行激活预取（prefetching）。

下图的时间线展示了在前向和后向传播过程中，计算（Forward N, Backward N）与参数/激活的卸载（Offload）和预取（Prefetch）操作如何重叠，从而提高训练效率。

大语言模型推理

LLM 推理：持久化 KV 缓存

本节内容致谢 Thor Johnsen (TRT-LLM)。

LLM 推理中的注意力计算

在 LLM 推理中，使用 KV（键-值）缓存可以在生成新 token 时重用之前 token 的计算结果。

• 步骤 1：上下文阶段 (Context Phase)：处理输入的整个序列，计算 Queries、Keys 和 Values，生成 Attention Score，并将 Keys 和 Values 存入 GPU 内存中的 KV 缓存。
• 步骤 N：生成阶段 (Generation Phase)：对于新生成的每个 token，只需计算当前 token 的 Query，并与缓存中的 Keys 和 Values 进行计算，从而显著减少计算量。

LLM 推理中的多轮对话

为多个用户提供服务的推理服务器需要高效地管理内存和计算资源。

• 恢复空闲会话可能迫使服务器重新计算整个上下文。
• 在 x86 + H100 架构中，通过 PCIe 卸载上下文的带宽有限，因此重新计算上下文通常是性能更优的选择。
• 在 Grace Hopper 架构中，通过 NVLink C2C 进行上下文卸载的带宽高，性能优于重新计算。

使用持久化 KV 缓存的注意力计算

持久化 KV 缓存技术通过将先前序列的计算结果保存在 CPU 内存中，进一步优化了性能。这项技术利用了高带宽直接 CPU 访问的优势。

• 上下文信息首先被加载到 CPU 内存中的持久化 KV 缓存。
• 在需要时，这部分缓存可以快速传输到 GPU 内存中，供生成阶段使用。

KV 缓存卸载性能

• 与 x86+H100 相比，在 Grace Hopper 上从持久化 KV 缓存重新加载上下文，可将首个 token 的延迟降低多达 40%。
• TensorRT-LLM v0.9 版本已实现 KV 缓存重用和卸载功能。

下图比较了有无持久化 KV 缓存时的吞吐量（越高越好）：
- 在 x86+A100、x86+H100 和 Grace Hopper 三种平台上，“上下文重载 + 持久化 KV 缓存”的吞吐量均高于“首次推理”。
- Grace Hopper 在两种情况下都表现出最高的吞吐量。

LLM 推理：参数卸载

本节内容致谢 Weiliang Liu, Sahil Modi, Rajnish Aggarwal (TRT-LLM team), Martin Mehringer (DevTech Compute EMEA)。

TensorRT 中的参数卸载优化

• Grace Hopper 允许将一部分模型参数卸载到 Grace CPU 内存中。
• GH200 拥有高达 480GB 的 CPU 内存和 96GB 的 GPU 内存。内存分配示意见下图，大部分模型参数存储在 CPU 内存中，仅部分权重、激活值和 KV 缓存驻留在 GPU 内存中。

异步参数卸载技术

• TRT 10.0 将允许用户在构建推理引擎时启用卸载功能。
• 构建 TRT 执行引擎时，可以使用特定命令（trtexec ... --stronglyTyped --enableWeightStreaming --weightStreamingBudget=<budget_megabytes>M）来启用权重流式传输。
• 这项技术被称为异步参数卸载。

Grace Hopper 上改进的吞吐量

• 在 x86 平台上，由于较高的首个 token 延迟和端到端延迟，参数流的价值有限。
• 在 Grace Hopper 上，参数卸载使得以合理的性能运行更大的模型成为可能。

下图展示了在单个 Grace Hopper 上推理 Llama2 70B FP16 模型（140GB 权重）的性能：
- 在批次大小为 16 时，首个 token 延迟低于 1 秒。
- 在批次大小为 96 时，达到最大吞吐量。
- 图中比较了 40% 权重在 GPU 和 60% 权重在 GPU 两种情况下的性能曲线。

推荐系统工作流

本节内容致谢 Abdurrahman Yasar, Chirayu Garg (DevTech Compute NALA)。

延迟流水线

推荐系统的工作流通常包含以下步骤：
1. 检索用户特征 (Retrieval User Features) - 使用 Pandas
2. 双塔模型 (Two-Tower Model) - 使用 Onnx runtime
3. 近似最近邻图搜索 (Approximate Nearest Neighbor Graph) - 使用 Rapids RAFT
4. 检索物品特征 (Retrieval Item Features) - 使用 Pandas
5. DLRM 模型 (DLRM Model) - 使用 Onnx runtime

基准测试通过随机生成查询并逐一执行流水线中的每个组件来完成。

数据集属性

• 用户: 22.1 万

  • 40 个连续特征（如年龄、消费水平）
  • 44 个分类特征

• 物品: 85 万

  • 20 个连续特征（如价格、库存）
  • 16 个分类特征

• 交互: 400 万次

延迟分析

• 总体而言，Grace Hopper 的性能比 x86 + H100 好 1.23 倍。
• Grace CPU 在增强推荐系统中深度学习模型的 CPU 任务上表现更佳。

下图展示了端到端推荐系统的运行时分解：
- 图中比较了 x86+A100、x86+H100 和 Grace Hopper 的执行时间。
- 堆叠条形图清晰地显示了 CPU 运行时（如特征获取）和 GPU 运行时（如模型推理）的耗时。
- Grace Hopper 在 CPU 密集型任务（用户和物品特征获取）上显著减少了执行时间，从而降低了总延迟。

数据库查询

本节内容感谢 Rui Lan, Eyal Soha, Nikolay Sakharnykh (DevTech Compute NALA) 的贡献。

数据库查询：x86 与 Grace Hopper 对比

• 在 x86 平台上，PCIe 成为瓶颈。
• Grace Hopper 凭借 NVLink C2C 带宽能够有效支持数据库查询。

下图展示了两种架构在处理数据库查询时的主要区别。在 x86 平台，CPU 与 GPU 之间通过 64GB/s 的 PCIe 连接，这限制了数据从 CPU 内存到 GPU 内存的传输速度。而在 Grace Hopper 超级芯片中，Grace CPU 与 Hopper GPU 通过 900 GB/s 的 NVLink C2C 直接互连，极大地提升了数据交换带宽。

不同数据库查询的性能

下图展示了在不同类型的数据库查询中，Grace Hopper 相对于 x86+A100 和 x86+H100 的性能优势。其核心技术是高带宽直接 CPU 访问。

• 聚合查询 (Aggregation Query): Grace Hopper 的吞吐量远超 x86 平台。
• 带压缩的连接查询 (Join Query + compression): Grace Hopper 表现出显著的性能提升。
• 连接查询 (Join Query): Grace Hopper 同样表现出色。
• 复杂连接查询 (Complex Join Query): Grace Hopper 依然保持领先。

总体而言，由于消除了 PCIe 瓶颈，Grace Hopper 在各类数据密集型查询中均展现出数倍的性能提升。

结论

关键要点

Grace Hopper 为 AI 应用带来的核心优势可以总结为以下三点：

1. 通过异步卸载实现更高的 GPU 吞吐量：Grace CPU 可以将任务异步地卸载到 Hopper GPU，实现 CPU 和 GPU 的并行处理，从而提升整体吞吐量。
2. 通过统一内存为系统内存提供高带宽缓存：利用统一内存技术，Hopper GPU 可以高效地访问 Grace CPU 的系统内存，如同访问一个高带宽缓存，极大地便利了对海量数据集的处理。
3. 优化的 AI 软件栈与开发者生产力：提供从硬件（Grace CPU + Hopper GPU）到 CUDA，再到上层训练和推理框架的完整、优化软件栈，提升了开发效率。

GTC'24 核心性能优化技术

以下是在 GTC'24 上展示的相关核心性能优化技术演讲列表：

• CUDA 编程与性能优化入门 [S62191]
• CUDA 高级性能优化 [S62192]
• Grace CPU 超级芯片性能优化 [S62275]
• Grace Hopper 超级芯片架构及深度学习应用性能优化 [S61159]
• 用于 HPC 和 AI 的多 GPU 编程模型 [S61339]
• 更快处理更多数据：Python 和 C++ 中的 GPU 内存管理最佳实践 [S62550]
• 利用 Grace Hopper 的能力加速向量数据库搜索 [S62339]
• 从零到极致：通过逐步优化将服务吞吐量提升数十倍 [S62410]

推荐议程:
- S62550，今天下午3点，由 Mark Harris 主讲内存分配器。
- CWE61236，明天上午10点，由 DevTech 主讲 Grace Hopper，将对本次演讲中的问题进行更深入的解答。