CUDA: New Features and Beyond

Stephen Jones, CUDA Architect | GTC 2024

目录

加速计算的核心议题

演示文稿首先指出,加速计算并不仅仅关乎性能(performance)。其核心在于能源(ENERGY),更具体地说是每瓦性能(Performance per Watt)。

数据中心与系统功耗的现状

能源消耗的增长趋势

北美主要市场数据中心总库存量(以兆瓦计)呈指数级增长,从几年前的1330兆瓦增长到预计的8253兆瓦。这表明了对计算能力和相应电力需求的急剧增加。

Page 5: 展示了北美主要市场数据中心总库存量(兆瓦)从2013年到2023年的增长曲线,显示出显著的上升趋势。
Page 5: 展示了北美主要市场数据中心总库存量(兆瓦)从2013年到2023年的增长曲线,显示出显著的上升趋势。

能源成本

随着电力需求的增长,主要市场的电力租用成本(以美元/千瓦/月计)也居高不下,中位数在$125到$200之间波动,这直接影响了数据中心的运营成本。

Page 6: 展示了北美主要市场数据中心月度租金中位数(美元/千瓦),数据在$125至$200之间波动。
Page 6: 展示了北美主要市场数据中心月度租金中位数(美元/千瓦),数据在$125至$200之间波动。

数据中心规模

一个典型的中型数据中心占地面积为20,000至100,000平方英尺,容纳2,000至10,000台服务器,功耗在5至20兆瓦之间。

系统功耗对比

不同计算系统的功耗差异巨大:
- 高性能台式机的最大系统功耗:1.5kW - 3kW
- 笔记本电脑的最大系统功耗:250W

晶体管发展的挑战:功耗效率

晶体管密度的指数增长

自1970年以来,晶体管密度一直在呈指数级增长,遵循摩尔定律的趋势。同时,从2004年的90nm工艺节点开始,台积电(TSMC)的光刻工艺密度也实现了超过100倍的提升。

Page 9: 两张图表分别展示了自1970年以来晶体管密度的对数尺度增长,以及台积电光刻工艺密度相对于90nm工艺的演进。
Page 9: 两张图表分别展示了自1970年以来晶体管密度的对数尺度增长,以及台积电光刻工艺密度相对于90nm工艺的演进。

功耗效率的下降

然而,一个关键问题是晶体管的功耗效率正在降低。虽然密度(Density)呈指数级增长,但功耗(Power)的增长速度虽然较慢,却也在持续上升,这意味着每个晶体管的能效并未同等比例提升。

Page 10: 图表显示,虽然晶体管密度(橙线)急剧增加,但功耗(蓝线)也在增加,表明能效提升速度跟不上密度提升速度。
Page 10: 图表显示,虽然晶体管密度(橙线)急剧增加,但功耗(蓝线)也在增加,表明能效提升速度跟不上密度提升速度。

功耗的两个主要来源

在计算中,我们主要在两个方面消耗电力:
1. 数据移动 (Data movement)
2. 计算 (Computation)

浮点运算与Tensor Core的能效优化

浮点运算的功耗成本分析

为了优化功耗,需要关注计算部分的能效,特别是浮点运算。

Page 12: 左侧图表展示了不同浮点数格式(从Double到FP8)的数值范围和精度构成。右侧表格列出了不同精度和操作的单位浮点运算能量消耗(Energy per FLOP)。
Page 12: 左侧图表展示了不同浮点数格式(从Double到FP8)的数值范围和精度构成。右侧表格列出了不同精度和操作的单位浮点运算能量消耗(Energy per FLOP)。

上图右侧的表格量化了不同精度下矩阵乘法操作的相对能耗(以FP32 FMA为基准1.0x):
- FP64 FMA: 2.5x
- FP32 FMA: 1.0x
- FP16 FMA: 0.5x
- FP64 Tensor Core MMA: 1.5x
- FP16 Tensor Core MMA: 0.12x
- FP8 Tensor Core MMA: 0.06x
- Int8 Tensor Core MMA: 0.04x

一个重要的观察是:浮点运算的功耗通常与尾数(mantissa)长度的平方成正比。

传统的融合乘加(FMA)指令在不同精度下的功耗差异显著。例如,FP64的功耗是FP32的2.5倍,而FP16的功耗仅为FP32的一半。

Page 13: 表格中高亮了FP64、FP32和FP16 FMA操作的能耗,以说明标准浮点运算的功耗成本。
Page 13: 表格中高亮了FP64、FP32和FP16 FMA操作的能耗,以说明标准浮点运算的功耗成本。

Tensor Cores的能效优势

与标准FMA操作相比,使用Tensor Cores进行矩阵乘法累加(MMA)操作可以显著提高能效。例如,FP16 Tensor Core MMA的能耗(0.12x)远低于FP16 FMA(0.5x)。总体而言,Tensor Cores比常规操作的能效高出1.5倍至4倍

Page 14: 表格中高亮了使用Tensor Core的各种精度操作,突显其相对于标准FMA的低能耗特性。
Page 14: 表格中高亮了使用Tensor Core的各种精度操作,突显其相对于标准FMA的低能耗特性。

通过比较,可以看出即便是FP64 Tensor Core MMA(1.5x),其能耗也远低于FP64 FMA(2.5x),而FP32 FMA的能耗(1.0x)则介于两者之间,这为在不同应用场景中选择合适的计算单元和数据精度以平衡性能与能效提供了依据。

Page 15: 通过高亮对比FP32 FMA和FP64 Tensor Core MMA,展示了不同计算路径下的能耗差异。
Page 15: 通过高亮对比FP32 FMA和FP64 Tensor Core MMA,展示了不同计算路径下的能耗差异。

张量核心 (Tensor Cores) 的功耗效率比普通操作高出1.5至4倍。下图比较了不同数值精度下的数值范围和精度,并列出了不同操作在矩阵乘法中每FLOP的能耗。通常,浮点运算的功耗与尾数长度的平方成正比。

  • FP16 Tensor Core MMA 的能耗仅为 FP32 FMA 的 0.12倍。
  • Int8 Tensor Core MMA 的能耗仅为 FP32 FMA 的 0.04倍。
Page 16
Page 16

功耗效率对比

与标准的FP64 FMA操作相比,使用FP16张量核心进行矩阵混合精度运算(MMA)可实现高达20倍的功耗效率提升。

  • FP64 FMA: 2.5x 相对能耗
  • FP16 Tensor Core MMA: 0.12x 相对能耗
Page 17
Page 17

利用低能耗晶体管解决高能耗问题

可以通过使用低精度张量核心的LU分解算法,来获得符合IEEE-754 FP64标准的精确结果。该方法采用迭代求精的策略,将计算密集型部分放在低精度的FP16上执行,从而大幅提升能效。

算法流程
1. 转换精度: 将输入的64位矩阵 A⁶⁴ 转换为16位矩阵 A¹⁶。
2. LU分解: 对 A¹⁶ 进行LU分解,得到初始解 x₀。这部分是计算密集型操作,在FP16上执行。
3. 转换回高精度: 将解 x₀ 转换回64位精度。
4. 计算残差: 在FP64精度下计算残差 r = b - Ax。
5. 求解修正量: 使用通用最小残差法(GMRES)求解线性方程组 L'U'Δc = L'rᵢ,得到修正量 Δc。
6. 修正解: 在FP64精度下更新解 xᵢ₊₁ = xᵢ + cᵢ。
7. 迭代: 重复步骤4-6,直到x足够精确。

最终输出的结果符合IEEE-754 FP64精度标准。

Page 18
Page 18

在FP16张量核心上实现完全FP64精度

实验结果表明,使用基于FP16张量核心的LU分解算法,通过迭代求精,可以在几次迭代后达到FP64的精度要求。如下图所示,残差(residual)在经过4次迭代后,就达到了FP64的精度水平。

Page 19
Page 19

兼顾功耗与性能

采用FP16张量核心和迭代求精的混合精度方法,可以在解决复杂矩阵LU分解问题(如此处的ZGESV,n=44000)时,同时实现功耗和性能的显著提升。

  • 性能 (TFLOP/s): 从FP64的41.1 TFLOP/s 提升到混合精度的 160.1 TFLOP/s,性能提升 3.9倍
  • 能效 (GFLOPS/Watts): 从FP64的74 GFLOPS/Watts 提升到混合精度的 425 GFLOPS/Watts,能效提升 5.7倍

下图展示了在GH200上运行时,FP64原生方法(蓝色曲线)与FP16+FP64混合精度方法(绿色曲线)的平均功耗对比。

Page 20
Page 20

在非原生硬件上实现高性能FP64 GEMM

基于Ootomo等人的研究方法,可以应用整数张量核心来执行双精度通用矩阵乘法(DGEMM),从而在没有原生FP64张量核心的硬件(如NVIDIA L40S)上实现高性能FP64计算。

实验显示,在L40S上,使用该方法进行Cholesky分解,性能比原生FP64提升了 7倍,并达到了A100 FP64张量核心性能的 54%

NVIDIA L40S GPU规格:
- GPU架构: NVIDIA Ada Lovelace
- GPU内存: 48GB GDDR6 with ECC
- 内存带宽: 864 GB/s
- FP64 TFLOPS: 1.4
- FP32 TFLOPS: 91.6
- TF32 Tensor Core TFLOPS: 183
- bfloat16 Tensor Core TFLOPS: 362
- FP16 Tensor Core TFLOPS: 362
- FP8 Tensor Core TFLOPS: 733

Page 21
Page 21

性能与能效分析

在L40S上,与原生FP64相比,模拟的FP64(使用整数张量核心)不仅性能更高,能效也显著提升。

  • 性能加速比: 随着矩阵规模 MNK 的增大,性能加速比稳步提升,最高可达7.7x。
  • 能效比: 模拟FP64的能效比原生FP64高出 8.6倍。尽管如此,其效率(0.7x)仍低于A100的原生张量核心。
Page 22
Page 22

编程模型、库与工具

张量核心编程抽象与融合

NVIDIA提供了不同层次的编程抽象库来使用张量核心,以平衡生产力(Productivity)和控制力(Control)。

  • cuBLAS: 用于GPU加速线性代数运算的API库。GEMM操作会自动使用张量核心以获得最大性能。提供额外的GEMM扩展以支持低精度和混合精度。
  • cuBLASLt: 用于GEMM的高级API。提供多种混合精度和epilogue(收尾操作)选项,以及可编程的、AI训练的启发式算法用于核函数选择和性能优化。
  • CUTLASS: 用于张量核心的C++和Python抽象。允许在核函数内联构建整个核函数或调用张量核心。对张量核心的配置、精度和执行有完全的控制权。
Page 23
Page 23

这些编程抽象可以根据其API类型(主机API vs. 设备API)和侧重点(生产力 vs. 控制力)进行分类。

  • 主机API (Host APIs):

    • cuBLAS: 侧重生产力。
    • cuBLASLt: 侧重控制力。
  • 设备API (Device APIs):

    • CUTLASS: 侧重控制力。
Page 24
Page 24

为了填补设备API在生产力方面的空白,NVIDIA引入了cuBLASDx。

  • cuBLASDx: 设备端GEMM API,具有自动启发式算法和张量核心配置。支持将GEMM与自定义的epilogue、其他GEMM或FFT融合在同一个CUDA核函数内。它简洁易用。
Page 25
Page 25

cuFFTDx, cuBLASDx: 用于FFT和MMA的数学设备扩展

这些库旨在帮助用户构建自己的高性能MMA(矩阵乘加)核函数。

  • cuBLASDx:
    • 功能: 一个易于使用的设备库,用于在核函数内部联优化的、张量核心加速的矩阵乘法(GEMM)。
    • 特性:
      - 块级GEMM操作。
      - 自动配置张量核心,无需手动调整或启发式设置,只需指定GPU架构。
      - 可与其他设备扩展库互操作。

下图展示了在一个融合操作 F = (A*B + C)*D + F 中,cuBLASDx相比于Thrust库的性能优势。

Page 26
Page 26
  • cuFFTDx:
    • 功能: 一个C++头文件库,用于在用户核函数内部联高性能的快速傅里叶变换(FFT)计算。
    • 特性:
      - 可根据核函数需求配置FFT块。
      - 当数据能装入共享内存/寄存器时,性能最佳。
      - 可与其他设备扩展库互操作。

下图展示了在H100上进行1D FP32卷积时,使用cuFFTDx(1个核函数)相比于cuFFT(3个核函数)的性能提升。

Page 27
Page 27

核函数融合与JIT编译

核函数融合(Kernel Fusion)是将多个独立的计算步骤合并到单个CUDA核函数中的技术,以减少开销并提高性能。

  • 非融合方式:

    • FP32->FP16转换 -> 矩阵乘法 -> 激活函数
    • 端到端成本:3次核函数启动、3次从内存加载数据、3次将数据存回内存。
  • 单一融合核函数:

    • 将上述三个步骤合并。
    • 端到端成本:1次核函数启动、1次从内存加载数据、1次将数据存回内存。
Page 28
Page 28

当预处理、内联函数和后处理的组合非常多时,静态地为每种组合编写一个融合核函数会导致“特化爆炸”(Specialization Explosion)。例如,4个预处理选项、4个内联函数选项和4个后处理选项会产生64种可能的组合。

解决方案: 在运行时根据需要动态构建融合,使用运行时核函数生成和即时编译(JIT)技术。

Page 29
Page 29

NVIDIA持续投入以减少在线和离线编译的时间。

  • 新的 "-minimal" 标志: 通过排除极少使用的、编译时开销大的特性(如纹理和表面),来减少前端编译时间。
  • 性能: 比默认的 "-O3" 优化级别开销低1.6倍。
  • 适用场景: 非常适合需要重复JIT编译的程序,例如在DNN中处理动态形状的场景。

下图展示了NVRTC(NVIDIA运行时编译库)编译一个"hello world"程序的耗时。自CUDA 11.8以来,编译速度提升了 6倍

Page 30
Page 30

Python 生态系统与框架

低层级的功能可能对高层级的开发产生巨大的影响。这张图展示了CUDA开发者的生态系统,它呈一个倒置的金字塔形状,越往上层,开发者数量越多。

Page 31 - 倒置的CUDA开发者金字塔
Page 31 - 倒置的CUDA开发者金字塔
  • 框架/SDK/DSL (Frameworks / SDKs / DSLs): 位于金字塔顶层,拥有广泛的内部和外部生态系统,服务于最大量的开发者。
  • 加速库 (Accelerated Libraries): 例如 cuBLAS, cuFFT, cuDNN 等,为特定领域的计算提供高性能实现。
  • 系统库与工具 (System Libraries & Utilities): 负责通信、存储、网络和系统管理等功能。
  • 设备内核 (Device Kernels): 使用原生语言和编程模型直接在GPU设备上执行的计算核心。
  • 系统软件 (System Software): 位于最底层,包括驱动程序、运行时、编译器和工具,是整个生态系统的基础。

为了让Python开发者能充分利用CUDA生态系统,NVIDIA正在构建一个更完整的体验。这个更宽广的倒置金字塔展示了针对Python生态系统的扩展。

Page 32 - 更宽广的倒置Python开发者金字塔
Page 32 - 更宽广的倒置Python开发者金字塔
  • 框架/SDK/DSL: 顶层保持不变,但对Python的支持更广泛。
  • 加速库: 提供更多带有Python绑定的库,以及支持统一内存的主机库。
  • 系统库与工具: 提供可定制的内存分配器。
  • 设备内核: 为机器学习应用中的数值类型提供支持。
  • 系统软件: 为设备端库提供支持。

nvmath-python 库

nvmath-python旨在以高性能的方式,提供对CUDA数学库功能的便捷Python访问。

Page 33 - nvmath-python 架构图
Page 33 - nvmath-python 架构图

其核心特性包括:
* 最小化前期准备时间: 可以快速地从Python中访问新的CUDA库特性。
* 与核心数值计算包的互操作性: 支持与NumPy, SciPy, CuPy, PyTorch, scikit-learn, scikit-image等库的协同工作。
* Pythonic的库接口: 为CUDA加速库提供了符合Python风格的接口。
* 平台无关性: 可在CPU (x86 & Arm) 和 Grace-Hopper 架构上运行。
* 访问所有扩展和峰值性能特性:
* 支持缩减精度和混合精度计算。
* 提供设备API。
* 支持功能定制(例如cuFFT的内联回调)。

CUTLASS

CUTLASS为编程NVIDIA GPU的张量核心(Tensor Cores)提供了一套分层抽象的C++模板库。其层次结构从原子层到设备层,提供了不同粒度的控制。

Page 34 - CUTLASS分层抽象模型
Page 34 - CUTLASS分层抽象模型

CUTLASS提供了Python接口,以便轻松地将其集成到深度学习框架中(例如PyTorch)。用户可以通过简洁的Python代码定义GEMM计划,并编译生成可直接在PyTorch中使用的扩展。

Page 35 - CUTLASS Python接口示例
Page 35 - CUTLASS Python接口示例

Python开发者工具

NVIDIA提供了用于详细检查基于Python的CPU和GPU执行的工具。这些工具支持Python调用栈采样、GIL追踪和动态代码范围注释,帮助开发者识别性能瓶颈。

Page 36 - Python开发者工具性能分析视图
Page 36 - Python开发者工具性能分析视图

通过Nsight Systems,开发者可以对目标应用程序进行动态性能分析,而无需更改源代码。通过JSON文件配置,即可在性能分析视图中高亮显示特定函数的执行范围。

Page 37 - Python动态性能分析注释配置与效果
Page 37 - Python动态性能分析注释配置与效果

Warp:面向Python的可微分空间计算

Warp是一个为Python设计的框架,专注于可微分空间计算。它允许使用纯Python语法编写JIT编译的CUDA内核,并支持自动微分,可与PyTorch、JAX等框架互操作。

Page 38 - Warp框架功能与应用
Page 38 - Warp框架功能与应用

Warp利用即时编译(JIT)技术实现反向模式自动微分,将Python代码的抽象语法树(AST)转换为可执行的CUDA代码。

Page 39 - Warp的JIT编译流程
Page 39 - Warp的JIT编译流程

一个应用示例是利用Warp进行神经应力场模拟,通过学习降阶神经基来表示应力场,实现了高达100,000倍的维度降低和10倍的模拟速度提升。

Page 40 - Warp在神经应力场模拟中的应用
Page 40 - Warp在神经应力场模拟中的应用

Legate:一个隐式并行框架

Legate是一个为实现隐式并行而设计的框架,它允许开发者使用顺序API来操作任意规模的系统。应用程序发出的顺序API调用,通过Legate运行时被映射到底层异构计算资源。

Page 41 - Legate框架架构
Page 41 - Legate框架架构

cuNumeric是基于Legate框架构建的,提供了NumPy API的隐式并行实现。用户编写标准的NumPy代码,无需任何修改即可扩展到数千个GPU上运行。

Page 42 - cuNumeric与Legate架构
Page 42 - cuNumeric与Legate架构

cuNumeric在模板计算基准测试中展现了出色的扩展性,其吞吐量随着GPU数量的增加几乎呈线性增长。

Page 43 - cuNumeric扩展性能图
Page 43 - cuNumeric扩展性能图

JAX:可微分计算框架

JAX是一个用于可微分计算的Python框架,其底层依赖XLA编译器和运行时。

Page 44 - JAX框架的架构分层
Page 44 - JAX框架的架构分层

Legate可以通过一个低层级插件与JAX集成,为仅支持SPMD的JAX带来了易于使用的流水线并行能力。

Page 45 - Legate与JAX集成架构
Page 45 - Legate与JAX集成架构

初步性能结果显示,在训练OPT 175B模型时,Legate JAX在不同规模的GPU集群上均表现出与当前最先进技术相当的竞争力。

Page 46 - Legate JAX 架构与性能
Page 46 - Legate JAX 架构与性能

性能分析与多节点通信

Nsight Systems 多节点分析

Nsight Systems 提供了多节点分析功能,支持从数千个节点收集报告,横向扩展处理,并以Jupyter Notebook、Parquet、CSV等多种格式生成集群范围内的结论。

Page 47 - Nsight Systems 多节点分析功能
Page 47 - Nsight Systems 多节点分析功能

分析工具生成的笔记本可以在Nsight Systems自身或Jupyter Lab中加载,以热力图或折线图的形式提供灵活的可视化与分析方式。

Page 48 - Nsight Systems 与 Jupyter Lab 中的工具笔记本
Page 48 - Nsight Systems 与 Jupyter Lab 中的工具笔记本

多节点热力图工具能够将整体GPU活动分解为通信、纯计算等不同视图,帮助开发者进行细致的性能瓶颈分析。

Page 49 - 多节点热力图工具示例
Page 49 - 多节点热力图工具示例

NVSHMEM 与 GPUDirect

NVSHMEM 是一种支持在 GPU 内核中进行细粒度通信和计算的技术。它为GPU间的通信提供了一个抽象层,屏蔽了底层物理传输(如NVLink, PCI, InfiniBand等)的复杂性。

Page 50 - NVSHMEM 概念与架构
Page 50 - NVSHMEM 概念与架构

NVSHMEM 位于软件栈的底层,为上层应用和库提供高性能通信支持。

Page 51 - NVSHMEM 在软件栈中的位置
Page 51 - NVSHMEM 在软件栈中的位置

在 GPUDirect 技术出现之前,GPU 数据通过网络发送需要CPU的深度参与,数据路径为 GPU -> CPU -> NIC,CPU成为了瓶颈。

Page 52 - GPUDirect 出现前的数据传输流程
Page 52 - GPUDirect 出现前的数据传输流程

异构计算与Grace Hopper新架构

CPU-GPU协同计算的演进与瓶颈

英伟达CEO黄仁勋(Jensen Huang)指出,一个关键的认识是,并行计算是对常规计算的扩展,而非取代。典型的CUDA程序分为在CPU上运行的主机代码和在GPU上运行的设备代码。

Page 61
Page 61

传统的异构计算系统由独立的CPU和GPU芯片组成,它们通过PCIe总线连接。

Page 63
Page 63

这种架构存在显著的性能瓶颈。CPU和GPU各自拥有独立的物理内存,其内存带宽差异巨大。PCIe总线的带宽(128 GB/s)远低于GPU的内存带宽(≤4900 GB/s),限制了数据传输速率。

Page 64
Page 64

Grace-Hopper:统一的异构处理器

为了解决PCIe瓶颈,NVIDIA 推出了 Grace Hopper 超级芯片,它将 Grace CPU 和 Hopper GPU 集成在一个模块上,通过高速的 NVLink-C2C 互连技术连接,构成一个统一的单一异构处理器。

  • Grace CPU:高性能Arm架构CPU。
  • Hopper GPU:强大的数据并行处理器。
  • NVLink-C2C:提供高达 900 GB/s 的双向带宽,远超PCIe Gen5。

这一设计极大地提升了CPU和GPU之间的通信效率。

Page 67
Page 67

Grace Hopper 架构包含两个独立的内存系统,每个系统都为其对应的处理器进行了优化。

Page 68
Page 68
  • CPU 内存系统:
    • 内存类型: CPU LPDDR5X
    • 优化目标: 为低延迟和深度缓存层次结构进行优化,适合执行链表操作等延迟敏感的代码。
Page 69
Page 69
  • GPU 内存系统:
    • 内存类型: GPU HBM3 (高带宽内存)
    • 优化目标: 为高吞吐量和高带宽缓存进行优化,适合执行向量规约等数据密集型任务。
Page 71
Page 71

统一地址空间与自动数据迁移

尽管存在两个物理上独立的内存系统,Grace Hopper 提供了 单一地址空间 的特性。这意味着CPU和GPU可以随时访问对方的内存,无需显式的数据拷贝操作。

Page 72
Page 72

为了进一步优化性能,统一内存系统(Unified Memory)支持 基于访问模式的自动数据迁移。当系统检测到GPU频繁访问位于CPU内存中的数据时,它会自动将这些数据迁移到GPU的HBM3内存中。

Page 73
Page 73

将数据迁移到GPU内存系统可以极大地优化GPU的访问性能。从GPU访问其本地HBM3内存的速度比通过NVLink-C2C访问CPU内存的速度快将近10倍。

Page 74
Page 74

即使数据被迁移到了GPU的HBM3内存中,CPU仍然可以通过NVLink-C2C访问这些数据,但会经历更高的延迟和更低的带宽。

Page 75
Page 75

Grace Hopper 应用案例

实现强耦合系统的 CPU+GPU 统一执行

下图展示了在 ALPS 上的 CSCS ICON 耦合大气和海洋模型(64 个 GPU)的性能。该模型耦合了 10km 分辨率的大气模型和 5km 分辨率的海洋模型。图中比较了不同系列的运行配置下,每日可模拟的天数,展示了Grace Hopper架构在强耦合科学计算应用中的优势。

Page 76
Page 76

监督式微调 (Supervised Fine Tuning)

监督式微调涉及对由多层 Transformer 堆叠而成的预训练语言模型进行调整。

Page 77
Page 77

在模型的前向传播过程中,每个 Transformer 层都会生成中间张量(activations)。

Page 78
Page 78

为了节省GPU内存,一种策略是在前向传播中丢弃中间张量,然后在反向传播时重新计算它们。

Page 79
Page 79
Page 80
Page 80

一种更高效的策略是利用Grace Hopper的内存层次结构。在前向传播过程中,将较大的中间张量异步地卸载到 CPU 内存中。

Page 81
Page 81

在反向传播过程中,当需要这些张量时,再从 CPU 内存中将它们调回 GPU。这种方法避免了重计算的开销。

Page 82
Page 82

在 Megatron LM 中对一个 10B 参数的 MoE 模型进行微调的实验表明,与基线方法相比,采用异步卸载的方案实现了 20% 的加速。

Page 83
Page 83

GraphSAGE 模型训练

GraphSAGE 是一种图神经网络(GNN)模型,其工作流程包括邻域采样、特征聚合和预测。

Page 84
Page 84

在OGBN-Papers100M数据集上的测试显示,与DGX-H100相比,GH200平台实现了2倍的整体加速,其中特征提取部分获得了6.5倍的显著加速。这得益于GH200的统一内存架构,极大地提升了大规模图数据特征的访问效率。

Page 85
Page 85

高级编程模型:任务图 (Task Graphs)

任务图执行基础

如果一项工作可以预先声明并且需要被多次执行,那么使用任务图(Task Graph)将非常有价值。任务图的核心思想是“一次定义,多次启动”。通过预先声明整个工作流并一次性提交,与传统的流式启动相比,可以将 CPU 的启动时间开销降低 7-8 倍。

Page 86
Page 86

任务图的属性与演进

任务图是由依赖关系连接起来的一系列操作。图中的每个节点都是一个任务,任务类型包括执行、控制、资源和子图。

Page 87
Page 87

任务图具有以下关键属性,使其功能强大且灵活:

  • 可重用 (Re-usable):图可以被生成一次,然后重复执行多次,非常适合迭代计算场景。
Page 88
Page 88
  • 循环 (Cyclic):“迭代直至收敛”是一种常见的模式,任务图支持这种循环执行。
Page 89
Page 89
  • 动态 (Dynamic):任务图支持运行时的、数据依赖的执行路径。例如,节点 B 的后续节点可以是 C 或 D,具体取决于运行时的数据。
Page 90
Page 90
  • 分层 (Hierarchical):一个任务可以是整个子图,即图的图。
Page 91
Page 91

数据依赖执行与动态控制流

“迭代直至收敛”是一种非常普遍的计算模式,例如求解线性方程组的共轭梯度算法。

function ConjugateGradient(A, b, x):
    ...
    do
        ...
    while(residual > 1e-8)
    ...
end

在传统实现中,循环的终止条件 while(residual > 1e-8) 在CPU上判断。这意味着每次迭代后,都需要将 residual 的值从GPU复制回CPU并进行同步,这会成为严重的性能瓶颈。

Page 98
Page 98

为了消除该瓶颈,可以通过使用新的“条件节点”(Conditional Nodes),将整个 while 循环(包括控制流)都移到GPU上,并转换为一个单一的任务图。

Page 100
Page 100

这样,整个求解过程可以通过一次图启动在GPU上运行至完成,避免了每次迭代的CPU-GPU交互,从而显著提升性能。

Page 101
Page 101

条件图节点 (Conditional Graph Nodes)

条件图节点是一种新型的图节点,它包含一个子图,该子图在满足 if()while() 条件为真时运行。

Page 106
Page 106
  • 主要特性:

    • 图结构保留: 条件节点只是另一种类型的图节点,因此原始的图结构得以保留。
    • IF 条件: "if" 节点内的子图在运行时,如果其条件为真则执行。如果条件为假,则跳过该节点。
    • WHILE 条件: "while" 循环与 "if" 类似,但在其子图执行完成后会重新评估条件。这允许 "while" 的循环体自身来设置其终止条件。
    • 嵌套: 条件节点可以任意深度地嵌套。
    • 多路分支: 可以实现多个 if 并行,功能类似于 switch 语句。多个并行的条件节点可以根据各自的条件独立执行,最终汇合。
  • 可用性:
    条件图节点功能从 CUDA 12.4 版本开始提供。

Page 107
Page 107

参考文献

Page 108
Page 108
  • S62162 深入了解数学库 (Deep Dive into Math Libraries)
  • S61198 CUTLASS: 一种针对 Hopper Tensor Cores 的高性能、灵活且可移植的方法 (CUTLASS: A Performant, Flexible, and Portable Way to Target Hopper Tensor Cores)
  • S63345 Warp: 利用 Python 中的可微分 GPU 计算推进模拟 AI (Warp: Advancing Simulation AI with Differentiable GPU Computing in Python)
  • S62262 Legate: 一个用于可组合、可扩展、加速库的高效编程框架 (Legate: A Productive Programming Framework for Composable, Scalable, Accelerated Libraries)
  • S62388 从您的数据中心和云应用中实现更高性能 (Achieving Higher Performance From Your Data Center and Cloud Application)
  • S61339 用于 HPC 和 AI 的多 GPU 编程模型 (Multi GPU Programming Models for HPC and AI)
  • S62559 加速和保护对大型数据集的 GPU 访问 (Accelerating and Securing GPU Accesses to Large Datasets)
  • S62129 大规模训练深度学习模型:NCCL 如何在 AI 数据中心网络上实现最佳性能 (Training Deep Learning Models at Scale: How NCCL Enables Best Performance on AI Data Center Networks)
  • S61598 使用 NVIDIA Grace 和 Arm 软件生态系统进行科学计算 (Scientific Computing With NVIDIA Grace and the Arm Software Ecosystem)
  • S61203 深入了解最新的 HPC 软件 (A Deep Dive into the Latest HPC Software)
  • S62275 Grace CPU 超级芯片的性能优化 (Performance Optimization for Grace CPU Superchip)
  • S62157 在 Alps 上利用 Grace Hopper 进行大规模早期科学研究 (Early Science with Grace Hopper at Scale on Alps)
  • S61159 Grace Hopper 超级芯片架构和深度学习应用的性能优化 (Grace Hopper Superchip Architecture and Performance Optimizations for Deep Learning Applications)
  • S61197 利用 cuGraph 加速的大规模图神经网络训练 (Large-Scale Graph GNN Training Accelerated With cuGraph)
  • CWE2240 NVIDIA 数学库 (NVIDIA Math Libraries)
  • CWE62915 使用 CUDA NVCC 编译器的 GPU 代码开发最佳实践 (Best Practices of GPU Code Development With CUDA NVCC Compiler)
  • CWE61532 您的 CUDA 工具箱里有什么?CUDA 分析、优化和调试工具 (What's in Your CUDA Toolbox? CUDA Profiling, Optimization, and Debugging Tools)
  • CWE63651 Legate: 一个用于可组合、可扩展、加速库的编程框架 (Legate Programming Framework for Composable, Scalable, Accelerated Libraries)
  • CWE61229 用于 HPC 和 AI 的 GPU 间通信技术和库 (Inter-GPU Communication Techniques and Libraries for HPC and AI)