Stephen Jones, CUDA Architect | GTC 2024
演示文稿首先指出,加速计算并不仅仅关乎性能(performance)。其核心在于能源(ENERGY),更具体地说是每瓦性能(Performance per Watt)。
北美主要市场数据中心总库存量(以兆瓦计)呈指数级增长,从几年前的1330兆瓦增长到预计的8253兆瓦。这表明了对计算能力和相应电力需求的急剧增加。
随着电力需求的增长,主要市场的电力租用成本(以美元/千瓦/月计)也居高不下,中位数在$125到$200之间波动,这直接影响了数据中心的运营成本。
一个典型的中型数据中心占地面积为20,000至100,000平方英尺,容纳2,000至10,000台服务器,功耗在5至20兆瓦之间。
不同计算系统的功耗差异巨大:
- 高性能台式机的最大系统功耗:1.5kW - 3kW
- 笔记本电脑的最大系统功耗:250W
自1970年以来,晶体管密度一直在呈指数级增长,遵循摩尔定律的趋势。同时,从2004年的90nm工艺节点开始,台积电(TSMC)的光刻工艺密度也实现了超过100倍的提升。
然而,一个关键问题是晶体管的功耗效率正在降低。虽然密度(Density)呈指数级增长,但功耗(Power)的增长速度虽然较慢,却也在持续上升,这意味着每个晶体管的能效并未同等比例提升。
在计算中,我们主要在两个方面消耗电力:
1. 数据移动 (Data movement)
2. 计算 (Computation)
为了优化功耗,需要关注计算部分的能效,特别是浮点运算。
上图右侧的表格量化了不同精度下矩阵乘法操作的相对能耗(以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的一半。
与标准FMA操作相比,使用Tensor Cores进行矩阵乘法累加(MMA)操作可以显著提高能效。例如,FP16 Tensor Core MMA的能耗(0.12x)远低于FP16 FMA(0.5x)。总体而言,Tensor Cores比常规操作的能效高出1.5倍至4倍。
通过比较,可以看出即便是FP64 Tensor Core MMA(1.5x),其能耗也远低于FP64 FMA(2.5x),而FP32 FMA的能耗(1.0x)则介于两者之间,这为在不同应用场景中选择合适的计算单元和数据精度以平衡性能与能效提供了依据。
张量核心 (Tensor Cores) 的功耗效率比普通操作高出1.5至4倍。下图比较了不同数值精度下的数值范围和精度,并列出了不同操作在矩阵乘法中每FLOP的能耗。通常,浮点运算的功耗与尾数长度的平方成正比。
与标准的FP64 FMA操作相比,使用FP16张量核心进行矩阵混合精度运算(MMA)可实现高达20倍的功耗效率提升。
可以通过使用低精度张量核心的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精度标准。
实验结果表明,使用基于FP16张量核心的LU分解算法,通过迭代求精,可以在几次迭代后达到FP64的精度要求。如下图所示,残差(residual)在经过4次迭代后,就达到了FP64的精度水平。
采用FP16张量核心和迭代求精的混合精度方法,可以在解决复杂矩阵LU分解问题(如此处的ZGESV,n=44000)时,同时实现功耗和性能的显著提升。
下图展示了在GH200上运行时,FP64原生方法(蓝色曲线)与FP16+FP64混合精度方法(绿色曲线)的平均功耗对比。
基于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
在L40S上,与原生FP64相比,模拟的FP64(使用整数张量核心)不仅性能更高,能效也显著提升。
NVIDIA提供了不同层次的编程抽象库来使用张量核心,以平衡生产力(Productivity)和控制力(Control)。
这些编程抽象可以根据其API类型(主机API vs. 设备API)和侧重点(生产力 vs. 控制力)进行分类。
主机API (Host APIs):
设备API (Device APIs):
为了填补设备API在生产力方面的空白,NVIDIA引入了cuBLASDx。
cuFFTDx, cuBLASDx: 用于FFT和MMA的数学设备扩展
这些库旨在帮助用户构建自己的高性能MMA(矩阵乘加)核函数。
下图展示了在一个融合操作 F = (A*B + C)*D + F 中,cuBLASDx相比于Thrust库的性能优势。
下图展示了在H100上进行1D FP32卷积时,使用cuFFTDx(1个核函数)相比于cuFFT(3个核函数)的性能提升。
核函数融合与JIT编译
核函数融合(Kernel Fusion)是将多个独立的计算步骤合并到单个CUDA核函数中的技术,以减少开销并提高性能。
非融合方式:
FP32->FP16转换 -> 矩阵乘法 -> 激活函数单一融合核函数:
当预处理、内联函数和后处理的组合非常多时,静态地为每种组合编写一个融合核函数会导致“特化爆炸”(Specialization Explosion)。例如,4个预处理选项、4个内联函数选项和4个后处理选项会产生64种可能的组合。
解决方案: 在运行时根据需要动态构建融合,使用运行时核函数生成和即时编译(JIT)技术。
NVIDIA持续投入以减少在线和离线编译的时间。
下图展示了NVRTC(NVIDIA运行时编译库)编译一个"hello world"程序的耗时。自CUDA 11.8以来,编译速度提升了 6倍。
低层级的功能可能对高层级的开发产生巨大的影响。这张图展示了CUDA开发者的生态系统,它呈一个倒置的金字塔形状,越往上层,开发者数量越多。
为了让Python开发者能充分利用CUDA生态系统,NVIDIA正在构建一个更完整的体验。这个更宽广的倒置金字塔展示了针对Python生态系统的扩展。
nvmath-python 库
nvmath-python旨在以高性能的方式,提供对CUDA数学库功能的便捷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++模板库。其层次结构从原子层到设备层,提供了不同粒度的控制。
CUTLASS提供了Python接口,以便轻松地将其集成到深度学习框架中(例如PyTorch)。用户可以通过简洁的Python代码定义GEMM计划,并编译生成可直接在PyTorch中使用的扩展。
Python开发者工具
NVIDIA提供了用于详细检查基于Python的CPU和GPU执行的工具。这些工具支持Python调用栈采样、GIL追踪和动态代码范围注释,帮助开发者识别性能瓶颈。
通过Nsight Systems,开发者可以对目标应用程序进行动态性能分析,而无需更改源代码。通过JSON文件配置,即可在性能分析视图中高亮显示特定函数的执行范围。
Warp:面向Python的可微分空间计算
Warp是一个为Python设计的框架,专注于可微分空间计算。它允许使用纯Python语法编写JIT编译的CUDA内核,并支持自动微分,可与PyTorch、JAX等框架互操作。
Warp利用即时编译(JIT)技术实现反向模式自动微分,将Python代码的抽象语法树(AST)转换为可执行的CUDA代码。
一个应用示例是利用Warp进行神经应力场模拟,通过学习降阶神经基来表示应力场,实现了高达100,000倍的维度降低和10倍的模拟速度提升。
Legate:一个隐式并行框架
Legate是一个为实现隐式并行而设计的框架,它允许开发者使用顺序API来操作任意规模的系统。应用程序发出的顺序API调用,通过Legate运行时被映射到底层异构计算资源。
cuNumeric是基于Legate框架构建的,提供了NumPy API的隐式并行实现。用户编写标准的NumPy代码,无需任何修改即可扩展到数千个GPU上运行。
cuNumeric在模板计算基准测试中展现了出色的扩展性,其吞吐量随着GPU数量的增加几乎呈线性增长。
JAX:可微分计算框架
JAX是一个用于可微分计算的Python框架,其底层依赖XLA编译器和运行时。
Legate可以通过一个低层级插件与JAX集成,为仅支持SPMD的JAX带来了易于使用的流水线并行能力。
初步性能结果显示,在训练OPT 175B模型时,Legate JAX在不同规模的GPU集群上均表现出与当前最先进技术相当的竞争力。
Nsight Systems 多节点分析
Nsight Systems 提供了多节点分析功能,支持从数千个节点收集报告,横向扩展处理,并以Jupyter Notebook、Parquet、CSV等多种格式生成集群范围内的结论。
分析工具生成的笔记本可以在Nsight Systems自身或Jupyter Lab中加载,以热力图或折线图的形式提供灵活的可视化与分析方式。
多节点热力图工具能够将整体GPU活动分解为通信、纯计算等不同视图,帮助开发者进行细致的性能瓶颈分析。
NVSHMEM 与 GPUDirect
NVSHMEM 是一种支持在 GPU 内核中进行细粒度通信和计算的技术。它为GPU间的通信提供了一个抽象层,屏蔽了底层物理传输(如NVLink, PCI, InfiniBand等)的复杂性。
NVSHMEM 位于软件栈的底层,为上层应用和库提供高性能通信支持。
在 GPUDirect 技术出现之前,GPU 数据通过网络发送需要CPU的深度参与,数据路径为 GPU -> CPU -> NIC,CPU成为了瓶颈。
英伟达CEO黄仁勋(Jensen Huang)指出,一个关键的认识是,并行计算是对常规计算的扩展,而非取代。典型的CUDA程序分为在CPU上运行的主机代码和在GPU上运行的设备代码。
传统的异构计算系统由独立的CPU和GPU芯片组成,它们通过PCIe总线连接。
这种架构存在显著的性能瓶颈。CPU和GPU各自拥有独立的物理内存,其内存带宽差异巨大。PCIe总线的带宽(128 GB/s)远低于GPU的内存带宽(≤4900 GB/s),限制了数据传输速率。
为了解决PCIe瓶颈,NVIDIA 推出了 Grace Hopper 超级芯片,它将 Grace CPU 和 Hopper GPU 集成在一个模块上,通过高速的 NVLink-C2C 互连技术连接,构成一个统一的单一异构处理器。
这一设计极大地提升了CPU和GPU之间的通信效率。
Grace Hopper 架构包含两个独立的内存系统,每个系统都为其对应的处理器进行了优化。
尽管存在两个物理上独立的内存系统,Grace Hopper 提供了 单一地址空间 的特性。这意味着CPU和GPU可以随时访问对方的内存,无需显式的数据拷贝操作。
为了进一步优化性能,统一内存系统(Unified Memory)支持 基于访问模式的自动数据迁移。当系统检测到GPU频繁访问位于CPU内存中的数据时,它会自动将这些数据迁移到GPU的HBM3内存中。
将数据迁移到GPU内存系统可以极大地优化GPU的访问性能。从GPU访问其本地HBM3内存的速度比通过NVLink-C2C访问CPU内存的速度快将近10倍。
即使数据被迁移到了GPU的HBM3内存中,CPU仍然可以通过NVLink-C2C访问这些数据,但会经历更高的延迟和更低的带宽。
下图展示了在 ALPS 上的 CSCS ICON 耦合大气和海洋模型(64 个 GPU)的性能。该模型耦合了 10km 分辨率的大气模型和 5km 分辨率的海洋模型。图中比较了不同系列的运行配置下,每日可模拟的天数,展示了Grace Hopper架构在强耦合科学计算应用中的优势。
监督式微调涉及对由多层 Transformer 堆叠而成的预训练语言模型进行调整。
在模型的前向传播过程中,每个 Transformer 层都会生成中间张量(activations)。
为了节省GPU内存,一种策略是在前向传播中丢弃中间张量,然后在反向传播时重新计算它们。
一种更高效的策略是利用Grace Hopper的内存层次结构。在前向传播过程中,将较大的中间张量异步地卸载到 CPU 内存中。
在反向传播过程中,当需要这些张量时,再从 CPU 内存中将它们调回 GPU。这种方法避免了重计算的开销。
在 Megatron LM 中对一个 10B 参数的 MoE 模型进行微调的实验表明,与基线方法相比,采用异步卸载的方案实现了 20% 的加速。
GraphSAGE 是一种图神经网络(GNN)模型,其工作流程包括邻域采样、特征聚合和预测。
在OGBN-Papers100M数据集上的测试显示,与DGX-H100相比,GH200平台实现了2倍的整体加速,其中特征提取部分获得了6.5倍的显著加速。这得益于GH200的统一内存架构,极大地提升了大规模图数据特征的访问效率。
如果一项工作可以预先声明并且需要被多次执行,那么使用任务图(Task Graph)将非常有价值。任务图的核心思想是“一次定义,多次启动”。通过预先声明整个工作流并一次性提交,与传统的流式启动相比,可以将 CPU 的启动时间开销降低 7-8 倍。
任务图是由依赖关系连接起来的一系列操作。图中的每个节点都是一个任务,任务类型包括执行、控制、资源和子图。
任务图具有以下关键属性,使其功能强大且灵活:
“迭代直至收敛”是一种非常普遍的计算模式,例如求解线性方程组的共轭梯度算法。
function ConjugateGradient(A, b, x):
...
do
...
while(residual > 1e-8)
...
end
在传统实现中,循环的终止条件 while(residual > 1e-8) 在CPU上判断。这意味着每次迭代后,都需要将 residual 的值从GPU复制回CPU并进行同步,这会成为严重的性能瓶颈。
为了消除该瓶颈,可以通过使用新的“条件节点”(Conditional Nodes),将整个 while 循环(包括控制流)都移到GPU上,并转换为一个单一的任务图。
这样,整个求解过程可以通过一次图启动在GPU上运行至完成,避免了每次迭代的CPU-GPU交互,从而显著提升性能。
条件图节点 (Conditional Graph Nodes)
条件图节点是一种新型的图节点,它包含一个子图,该子图在满足 if() 或 while() 条件为真时运行。
主要特性:
if 并行,功能类似于 switch 语句。多个并行的条件节点可以根据各自的条件独立执行,最终汇合。可用性:
条件图节点功能从 CUDA 12.4 版本开始提供。