生成式AI超分辨率加速可扩展、低计算天气预测

随着人工智能天气和气候预测模型迅速得到采用，由某中心提供的Earth-2平台提供了一系列库和工具，可利用GPU优化的软件栈来加速解决方案。降尺度（Downscaling）是一项精炼粗分辨率（25公里尺度）天气数据的任务，它使国家气象服务机构能够以足够精细的空间分辨率，为农业、能源、交通和防灾提供高分辨率预测，以支持可操作的决策和规划。

传统的动力降尺度方法成本极其高昂，尤其是在高分辨率、大范围空间域上运行大型集合时。CorrDiff是一个生成式AI降尺度模型，它绕过了传统数值方法的计算瓶颈，实现了最先进的结果，并采用基于块（patch）的多重扩散方法，可扩展到大陆和全球范围。与传统的数值方法相比，这种基于AI的解决方案在显著降低计算成本的同时，实现了效率和可扩展性的大幅提升。

CorrDiff已在全球范围内应用于多种场景，证明了其在需要精细尺度天气信息的各领域的多功能性和影响力：

• 某机构的天气预报服务，用于支持农业、能源和航空业。
• 某机构用于改善中东地区的烟雾和沙尘暴预测。
• 某机构用于增强一系列风暴尺度预测，包括火灾天气预报和扰乱铁路运营的阵风预报。

本文展示了为CorrDiff训练和推理进行的性能优化和增强，这些优化已集成到Earth-2软件栈中的两个工具：某中心的PhysicsNeMo和某中心的Earth2Studio。这些优化实现了超过50倍的训练和推理基线加速，从而能够：

• 在3000 GPU小时以内，完成基于块的整个行星尺度训练。
• 将大多数国家尺度的训练降低到O(100) GPU小时量级。
• 在美国本土（CONUS）区域进行训练，耗时低于1000 GPU小时。
• 支持微调和定制训练，使公里尺度的AI天气预测更加普及。
• 国家尺度推理在GPU秒级完成，行星尺度推理在GPU分钟级完成。
• 经济高效地生成大型集合，用于高分辨率概率预测。
• 交互式探索公里尺度数据。

CorrDiff：训练和推理

图1说明了用于生成式降尺度的CorrDiff训练和采样工作流程。在扩散训练期间，使用预训练的回归模型生成条件均值，作为训练扩散模型的输入。关于CorrDiff的背景和详细信息，请参阅CorrDiff论文、PhysicsNeMo文档和源代码。

为何优化CorrDiff？

扩散模型是资源密集型的，因为它们依赖于迭代采样，每个去噪步骤都涉及多次神经网络计算。这使得推理耗时且成本高昂。训练也更加昂贵，因为去噪器必须在所有噪声水平上进行训练。优化其性能需要：

• 简化核心操作（例如，融合内核、使用混合精度、使用某中心的CUDA图等）。
• 通过减少去噪步骤数和使用最优时间积分方案来改进采样过程。

CorrDiff使用EDM架构，其中几个计算昂贵的操作，例如组归一化、激活函数和卷积，可以使用高级软件包和内核融合进行优化。

CorrDiff还使用两阶段流水线（回归和校正），通过缓存回归输出并在多个扩散步骤中分摊成本，从而最小化冗余计算。

加速后的CorrDiff

在下图中，我们总结了实现训练和推理成本超过50倍加速的各种优化。图2和图3总结了每项连续优化相对于基线实现的累积加速因子。每项优化的详细信息在后续章节中提供。

图2. CONUS尺度训练中，基于块的CorrDiff训练每个块的加速。每个块约20万像素（448×448）。绿色条形图绘制了相对于fp32基线性能的加速比（左y轴），灰色实线绘制了每个块元素的采样运行时间（对数刻度，右y轴）。

图3. CONUS尺度推理中，基于块的CorrDiff每个样本的加速。样本约200万像素（1056×1792）。绿色条形图绘制了相对于fp32基线性能的加速比（左y轴），灰色实线绘制了每个批次元素的采样运行时间（对数刻度，右y轴）。

优化的CorrDiff：如何实现

在NVIDIA H100 GPU上使用FP32精度、批大小=1、块大小=1的CorrDiff基线性能（绝对时间）如下：

• 回归前向传播：1204ms
  • 区域：美国本土（CONUS），大小为1056 × 1792像素
  • 输入通道：25公里分辨率下的 "u500", "v500", "z500", "t500", "u850", "v850", "z850", "t850", "u10m", "v10m", "t2m", "tcwv"
  • 输出通道：2公里分辨率下的 "refc", "2t", "10u", "10v"
• 扩散前向传播：155ms
  • 区域：空间块，大小为448 x 448像素
  • 输入通道：25公里分辨率下的 "u500", "v500", "z500", "t500", "u850", "v850", "z850", "t850", "u10m", "v10m", "t2m", "tcwv"
  • 输出通道：2公里分辨率下的 "refc", "2t", "10u", "10v"
• 扩散反向传播：219ms

虽然有效，但这个基线受到昂贵的回归模型前向传递和低效的数据转置的限制。

图4. 训练性能剖析：NVIDIA Nsight Systems剖析图，显示了基线基于块的CorrDiff训练各阶段的运行时间分布，回归前向传播主导了总迭代时间。

关键CorrDiff训练优化

为了实现CorrDiff训练的实质性加速，最终在NVIDIA B200上达到53.86倍、在H100上达到25.51倍的加速比，我们引入了以下一系列性能优化：

1. 优化1：为训练启用AMP-BF16 原始训练方案使用FP32精度。此处，为训练启用了带有BF16的自动混合精度（AMP），以减少内存使用并提高吞吐量，同时不影响数值稳定性，从而实现了相对于基线的2.03倍加速。
2. 优化2：使用多迭代块分摊回归成本 在原始的基于块的训练工作流程中，用于扩散模型训练的每个448×448块样本都需要对整个1056×1792的CONUS区域进行回归模型推理。这导致扩散模型训练吞吐量受限于回归模型推理。 我们通过缓存回归输出并在每个时间戳的16个扩散块中重复使用，提高了效率。这提供了更广泛的空间覆盖，同时更有效地分摊了回归成本，从而实现了相对于基线的12.33倍加速。
3. 优化3：使用Apex GroupNorm消除数据转置 训练管道最初使用默认的NCHW内存布局，这会在卷积前后触发代价高昂的隐式内存转置。将模型和输入张量切换为NHWC（通道最后）格式，使其与cuDNN的首选布局对齐。然而，PyTorch的GroupNorm操作不支持通道最后格式。为了防止转置并保持数据为通道最后格式以实现更高效的归一化内核，我们将PyTorch GroupNorm替换为某中心的Apex GroupNorm。这消除了转置开销，并实现了相对于基线的16.71倍加速。
4. 优化4：融合GroupNorm与SiLU 通过使用Apex将GroupNorm和SiLU激活融合到单个内核中，减少了内核启动次数和全局内存访问次数。这提高了GPU利用率，并实现了相对于基线的17.15倍加速。
5. 优化5：在Apex GroupNorm中扩展通道维度支持 一些CorrDiff层使用的通道维度不受Apex支持。我们扩展了对这些通道维度的支持，为所有层解锁了融合功能。这将性能提升到相对于基线的19.74倍加速。
6. 优化6：通过图编译进行内核融合 我们使用torch.compile来融合剩余的元素级操作（例如加法、乘法、sqrt、exp）。这改善了调度，减少了全局内存访问，并削减了Python开销，实现了相对于基线的25.51倍加速。
7. 优化7：在NVIDIA Blackwell上使用Apex GroupNorm V2 使用为NVIDIA Blackwell GPU优化的Apex GroupNorm V2，在B200上实现了相对于基线的53.86倍加速，以及在H100优化工作流上的2.1倍加速。

图5：优化后的基于块的CorrDiff训练剖析图

训练吞吐量

我们在表1中比较了在NVIDIA Hopper架构上的基线CorrDiff与在Hopper和Blackwell架构上优化版本的训练吞吐量。优化实现在两种架构上都实现了效率提升，其中Blackwell显示出最显著的增益。

表1. CorrDiff训练吞吐量比较

"光速"（Speed-of-Light）分析

为了评估优化后的CorrDiff工作流距离硬件性能上限有多近，我们在H100和B200 GPU上进行了"光速"（SOL）分析。这通过评估GPU资源的使用效率，提供了可达到性能的上限估计。

估算SOL的步骤：

1. 识别低利用率内核： 我们关注DRAM读/写带宽 < 60% 且 Tensor Core利用率 < 60%的内核。此类内核既不是内存受限也不是计算受限，很可能是性能瓶颈。
2. 估算每个内核的潜力： 对于每个低利用率内核，我们估算其在理想条件下的潜在加速比——即在充分利用DRAM带宽或Tensor Core活动的情况下。
3. 汇总总体加速比： 然后，如果每个内核都优化到其理想性能，我们计算假设的端到端加速比。
4. 计算SOL效率： 最后，如果前10个运行时间主导的内核分别提升到其理论最大值，我们估计可实现理论最大SOL的比例。

在此框架下，优化后的CorrDiff工作流在H100上达到了估计SOL的63%，在B200上达到了67%。这表明GPU利用率很高，同时仍为未来的内核级改进留有空间。

为了进一步评估效率，我们如图5和图6所示可视化内核性能。每个点代表一个内核，根据NVIDIA Tensor Core利用率（x轴）和组合的DRAM读/写带宽利用率（y轴）绘制。点的大小反映了其在总运行时间中的占比，突出了性能关键操作。

位于顶部或右侧边缘附近的内核通常优化良好，因为它们充分利用了计算或内存资源。相比之下，位于左下象限的内核两者都未充分利用，代表了进一步优化的最佳机会。这种可视化提供了运行时分布的清晰图像，有助于确定GPU效率可以改进的地方。

图6：在B200上的基线FP32基于块CorrDiff内核利用率分布

图6显示了基线CorrDiff实现的内核在Tensor Core利用率和DRAM带宽利用率方面的分布。在使用FP32精度的未优化工作流中，>95%的时间花费在低利用率内核上，其中DRAM利用率（读取+写入）和张量核心利用率均低于60%。大多数运行时间主导的内核聚集在原点附近，显示出非常低的DRAM和Tensor Core利用率。只有少数内核位于靠近上边界或右边界的区域，那里的内核变得明显是内存受限或计算受限。未优化的美国CorrDiff工作流在B200上仅达到SOL的1.23%。

图7：在B200上的优化后BF16基于块CorrDiff内核利用率分布

图7显示了优化后实现的内核在Tensor Core利用率和DRAM带宽利用率方面的分布。在使用AMP-BF16训练的优化工作流中，更高比例的内核位于左上或右下边缘附近，表明具有良好的性能和GPU利用率。优化后的CorrDiff现在在B200上达到了SOL的67%。尽管优化工作流整体有所改进，但一些内核仍有进一步加速的潜力。

CorrDiff推理优化

许多训练优化也可应用于推理。我们提出了几项针对推理的优化，以最大化性能。

• 优化的多重扩散 CorrDiff使用基于块的多重扩散方法，其中重叠的空间块被去噪然后聚合。最初，总运行时间的27.1%花费在im2col折叠/展开操作上。通过为每个像素预计算重叠计数，并使用torch.compile()加速剩余的折叠/展开步骤，完全消除了im2col瓶颈，从而实现了7.86倍的加速。
• 确定性欧拉采样器（12步） 原始的随机采样器使用带有Heun求解器和二阶修正的18个去噪步骤。通过启用使用欧拉求解器（无二阶修正）的确定性采样器，我们将去噪步骤数减少到12，而不影响输出质量。这一改变在Hopper和Blackwell上都提供了额外的约2.8倍加速。使用12步确定性采样器的最终加速在H100上为21.94倍，在B200上为54.87倍。

本文描述的几项优化也适用于一般的扩散模型，有些则是基于块的方法所特有的。因此，这些优化可以移植到PhysicsNeMo中的其他模型，并用于开发天气降尺度之外的解决方案。

开始使用

• 在PhysicsNeMo中训练/推理CorrDiff： 请参阅PhysicsNeMo CorrDiff文档。

fp_optimizations: amp-bf16
use_apex_gn: True
torch_compile: True
profile_mode: False

use_fp16: True
use_apex_gn: True
use_torch_compile: True
profile_mode: False
io_syncronous: True

将profile_mode设置为False以获得优化性能，因为NVTX注释会为torch.compile工作流引入图中断。

要利用最新的Apex GroupNorm内核，可以在PhysicsNeMo容器Dockerfile中构建Apex GroupNorm，或者在加载PhysicsNeMo容器后在本地构建：

1. 克隆Apex仓库并使用以下命令构建：

CFLAGS="-g0" NVCC_APPEND_FLAGS="--threads 8" \
pip install \
  --no-build-isolation \
  --no-cache-dir \
  --disable-pip-version-check \
  --config-settings "--build-option=--group_norm" .

了解更多关于PhysicsNeMo中优化的CorrDiff训练的信息，并在Earth2Studio中运行优化的工作流。

视频1. 可视化在CONUS上基于块的CorrDiff降尺度，并实现55倍加速