Related Resources
代理式 AI/生成式 AI

生成式 AI 超分辨率利用可扩展的低计算模型加速天气预报

2025年 11月 10日
By , and

随着AI在天气和气候预测模型中的广泛应用，NVIDIA Earth-2平台提供了基于GPU优化的软件堆栈，以及用于加速解决方案的库和工具。降尺度是一项将粗分辨率（25公里尺度）天气数据精细化的任务，可帮助国家气象服务（NMS）机构生成更高空间分辨率的预测结果，从而为农业、能源、交通和防灾等领域提供更精准的决策支持与规划依据。

传统的动态降采样在高分辨率和大范围空间域处理大规模数据集时，计算成本较高。 CorrDiff，一种基于生成式 AI 的降采样模型，能够规避传统数值方法的计算瓶颈，利用基于图像块的 多扩散 策略，实现优异的降采样效果，并可扩展至大陆乃至全球尺度。相比传统数值方法，该 AI 驱动的解决方案在显著提升计算效率与可扩展性的同时，大幅降低了计算资源消耗。

CorrDiff 已在全球多种应用场景中得到广泛采用，充分展现了其在依赖高精度天气信息领域的通用性与影响力。

  • The Weather Company (TWC) 为农业、能源和航空等行业提供支持。
  • G42 用于提升中东地区对烟雾和沙尘暴的预测能力。
  • Tomorrow.io 则致力于增强多种规模风暴的预报精度，包括可能影响铁路运营的野火天气及强阵风预测。

在这篇博客文章中，我们展示了已集成到 Earth-2 堆栈中两个工具——NVIDIA PhysicsNeMo 和 NVIDIA Earth2Studio——的 CorrDiff 训练与推理性能优化及增强功能。这些优化可将训练和推理的基准速度提升五十多倍，从而显著提升计算效率与模型响应能力。

  • 在 3000 GPU 小时内完成基于补丁的全球模型训练。
  • 将多数国家或地区规模的训练耗时缩短至约 100 GPU 小时。
  • 针对美国本土（CONUS）的训练，耗时低于 1000 GPU 小时。
  • 通过微调与定制化训练，推动公里级人工智能气象模型的广泛应用。
  • 国家或地区级别的推理仅需 GPU 秒级算力，全球范围推理则在 GPU 分钟级内完成。
  • 以经济高效的方式生成大规模集成数据，支持高分辨率的概率化天气预测。
  • 实现对公里级气象数据的交互式探索与分析。

CorrDiff：训练和推理

Coarse-resolution global weather data at 25 km scale is used to predict the mean µ first using a regression model, which is then stochastically corrected using an Elucidated Diffusion Model (EDM) r, together producing the probabilistic high-resolution 2 km-scale regional forecast. Bottom right: diffusion model is conditioned with the coarse-resolution input to generate the residual r after a few denoising steps. Bottom left: the score function for diffusion is learned based on the UNet architecture. Coarse-resolution global weather data at 25 km scale is used to predict the mean µ first using a regression model, which is then stochastically corrected using an Elucidated Diffusion Model (EDM) r, together producing the probabilistic high-resolution 2 km-scale regional forecast. Bottom right: diffusion model is conditioned with the coarse-resolution input to generate the residual r after a few denoising steps. Bottom left: the score function for diffusion is learned based on the UNet architecture.
图 1：CorrDiff 训练和采样工作流程

图1展示了CorrDiff在生成式降采样中的训练与采样流程。在扩散模型的训练过程中，利用预训练的回归模型生成条件均值，并将其作为扩散模型训练的输入。有关CorrDiff的背景信息与技术细节，可参考CorrDiff的相关论文、PhysicsNeMo 文档及源代码。

为何要优化 CorrDiff？ 

扩散模型需要大量计算资源，因为其依赖迭代采样过程，每个降噪步骤都涉及多次神经网络运算，导致推理过程耗时且成本高昂。训练成本同样较高，因为降噪网络必须适应多种噪声水平。为提升性能，需进行针对性优化。

  • 简化核心操作（例如融合内核、采用混合精度、利用 NVIDIA CUDA 图形技术等），
  • 并通过减少降噪步骤数量以及采用优化的时间积分方案来提升采样效率。

CorrDiff 采用 EDM 架构，其中组归一化、激活函数和卷积等计算开销较大的操作，可通过高级软件库和内核融合技术进行优化。

CorrDiff 采用两阶段工作流（回归和校正），通过缓存回归输出，减少冗余计算，并在多个扩散步骤中分摊计算成本，从而提升效率。

加速 CorrDiff

在下图中，我们总结了各项优化措施，与 CONUS 区域相比，这些优化使训练和推理成本降低了 50 倍以上。图 2 和图 3 展示了每次连续优化在基准测试中实现的累计加速比。后续章节将详细阐述每一项优化的具体内容。

Patch-based CorrDiff training speedup vs. FP32 baseline with runtime shown on log scale. Patch-based CorrDiff training speedup vs. FP32 baseline with runtime shown on log scale.
图2。针对CONUS区域的训练，展示了基于图像块的CorrDiff训练加速效果. 每个图像块包含约20万像素（448 × 448）。绿色条形表示相对于fp32基准的加速比（左侧y轴），灰色实线表示每个分块元素在对数刻度下的采样运行时间（右侧y轴）
Bar and line chart titled “Patch-based CorrDiff Inference Speedup per Sample for CONUS-scale (~2M pixels, 1056×1792).” Green bars show the inference speedup compared to the FP32 baseline (using the left y-axis), while a solid gray line indicates sampling runtime per batch element on a logarithmic scale (right y-axis). The chart visualizes improved performance and reduced runtime for various precision modes. Bar and line chart titled “Patch-based CorrDiff Inference Speedup per Sample for CONUS-scale (~2M pixels, 1056×1792).” Green bars show the inference speedup compared to the FP32 baseline (using the left y-axis), while a solid gray line indicates sampling runtime per batch element on a logarithmic scale (right y-axis). The chart visualizes improved performance and reduced runtime for various precision modes.
图3展示了基于补丁的CorrDiff在CONUS尺度（200万像素，1056×1792）下每个样本的推理加速效果。绿色条形表示相对于fp32基准的加速比（左侧y轴），灰色实线表示每个批量元素在对数刻度下的采样运行时间（右侧y轴）

优化的 CorrDiff：实现方式

CorrDiff 在 NVIDIA H100 GPU 上的基准性能如下：采用 FP32 精度，batch size = 1，patch size = 1（以绝对时间计）。

  • 回归前处理：1204 毫秒
    • 区域范围：大小为 1056 × 1792 像素的 CONUS 区域
    • 输入通道：25 公里分辨率下的 [“u500”、“v500”、“z500”、“t500”、“u850”、“v850”、“z850”、“t850”、“u10m”、“v10m”、“t2m”、“tcwv”]
    • 输出通道：2 公里分辨率下的 [“refc”、“2t”、“10u”、“10v”]
  • 前向扩散：155 毫秒
    • 区域范围：大小为 448 × 448 像素的空间补丁
    • 输入通道：25 公里分辨率下的 [“u500”、“v500”、“z500”、“t500”、“u850”、“v850”、“z850”、“t850”、“u10m”、“v10m”、“t2m”、“tcwv”]
    • 输出通道：2 公里分辨率下的 [“refc”、“2t”、“10u”、“10v”]
  • 反向扩散：219 毫秒

该基准虽然有效，但受限于高昂的回归模型前向计算成本以及低效的数据转置操作。

Figure 4 shows an NVIDIA Nsight Systems training performance profile of the baseline patch-based CorrDiff model. The chart displays the runtime distribution across different training stages, highlighting that the regression forward stage dominates total iteration time, indicating it as the primary computational bottleneck during training. Figure 4 shows an NVIDIA Nsight Systems training performance profile of the baseline patch-based CorrDiff model. The chart displays the runtime distribution across different training stages, highlighting that the regression forward stage dominates total iteration time, indicating it as the primary computational bottleneck during training.
图4。训练性能分析：NVIDIA Nsight Systems 性能分析结果显示，在基于补丁的基准 CorrDiff 训练阶段中，前向传播占据了单次迭代的大部分运行时间。

CorrDiff 训练优化的关键之处

为大幅加速 CorrDiff 的训练过程，在 NVIDIA B200 上实现 53.86 倍加速，在 H100 上实现 25.51 倍加速，我们推出了一系列性能优化措施，概述如下。

优化 1：启用 AMP-BF16 进行训练
原始训练采用 FP32 精度，而本方案引入了支持 BF16 的自动混合精度（AMP）技术，在保障数值稳定性的前提下有效降低了内存占用并提升了吞吐量，训练速度达到基准的 2.03 倍。

优化 2：采用多迭代补丁以分摊回归成本
在早期基于补丁的训练流程中，每个 448 × 448 的补丁样本在用于扩散模型训练时，都需要对整个 1056 × 1792 的 CONUS 区域进行回归模型推理。这导致扩散模型的训练吞吐量受限于回归模型的推理速度。

我们通过缓存回归输出，并在每个时间步的16个扩散块中重复利用这些输出以提升效率。这种方法不仅实现了更广泛的空间覆盖，还更有效地分摊了回归计算成本，使运行速度较基准提升了12.33%。

优化 3：使用 Apex GroupNorm 消除数据转置
原始训练流程采用默认的 NCHW 内存布局，这会在卷积操作前后引发昂贵的隐式内存转置。为提升效率，将模型和输入张量切换为 NHWC（通道在最后）格式，以匹配 cuDNN 的推荐内存布局。然而，PyTorch 自带的 GroupNorm 操作不支持通道在最后的格式，若直接使用会导致额外的数据转置开销。为避免这一开销并保持数据在高效格式中处理，我们采用 NVIDIA Apex GroupNorm 替代原生 PyTorch 版本。该改进有效消除了数据转置操作，使整体训练速度相比基准提升了 16.71× speedup over the baseline。

优化 4：融合 GroupNorm 与 SiLU 激活函数
通过 Apex 将 GroupNorm 与 SiLU 激活函数融合为单个内核，减少了内核启动次数和全局内存访问开销。该优化提升了 GPU 利用率，使运行速度较基准提升了 17.15%。

优化 5：在 Apex Group 中扩展了通道维度的支持，以规范 CorrDiff 层中使用 Apex 原本不支持的通道维度。
通过增强对这些通道维度的兼容性，实现了所有层的融合，使性能相比基准提升了 19.74%速度提升。

优化 6：通过图形编译实现内核融合
我们采用 torch.compile 将剩余的元素级运算（如加法、乘法、sqrt、exp 等）进行融合。该方法优化了执行调度，减少了全局内存访问次数，同时降低了 Python 层面的开销，使运行速度较基准提升了 25.51%.

优化 7：在 NVIDIA Blackwell GPU 上采用经过专门优化的 Apex GroupNorm V2
使用 Apex GroupNorm V2，优化为 NVIDIA Blackwell GPU，相比 B200 上的基准提升了 53.86× 速度提升，相比 H100 优化工作流程提升了 2.1×。

NVIDIA Nsight Systems profile showing runtime distribution of optimized patch-based CorrDiff training stages, where regression forward cost was amortized across multiple diffusion training steps. NVIDIA Nsight Systems profile showing runtime distribution of optimized patch-based CorrDiff training stages, where regression forward cost was amortized across multiple diffusion training steps.
图5：基于补丁优化的CorrDiff训练配置文件

训练吞吐量 

我们将 CorrDiff 在 NVIDIA Hopper 架构上的基准训练吞吐量，与表 1 中 Hopper 和 Blackwell 架构上的优化版本进行了对比。优化后的实现提升了两种架构的效率，其中 Blackwell 的性能提升尤为显著。

注意：回归指回归的前向传播过程，扩散指前向扩散过程，总计包含（前向回归 + 前向扩散 + 反向扩散）的综合成本。

GPU 版本 精度 回归 (毫秒/补丁) 扩散 (毫秒/补丁) 总运行时 (毫秒/补丁) 吞吐量 (补丁/秒)
H100 基准 FP32 1204.0 374.0 1578.0 0.63
H100 优化 BF16 10.6 951.2 961.8 16.2
B200 优化 BF16 4.7 424.5 429.2 234.1
表 1。CorrDiff 训练吞吐量对比

光速分析

为评估优化后的 CorrDiff 工作流程与硬件性能上限之间的差距，我们在 H100 和 B200 GPU 上开展了光速（SOL）分析。该分析通过衡量 GPU 资源的利用效率，为可实现的性能提供了理论上限的估算。

估算 SOL 的步骤：

  1. 识别低利用率内核：
    我们重点关注 DRAM 读写带宽低于 60% 且 Tensor Core 利用率低于 60% 的内核。这类内核既不受内存带宽限制，也不受计算能力限制，可能存在其他性能瓶颈。
  2. 评估每个内核的优化潜力：
    针对每个低利用率内核，估算其在理想状态下的加速空间，即当 DRAM 带宽或 Tensor Core 活动达到饱和时可能实现的性能提升。
  3. 汇总整体加速效果：
    假设所有低利用率内核均被优化至理想性能水平，计算由此带来的端到端整体加速比。
  4. 计算 SOL 效率：
    若运行时占比最高的前 10 个主要内核各自提升至其理论性能上限，则将此时达到的 SOL 占理论最大 SOL 的比例，作为可实现峰值性能的估计指标。

使用该框架，我们优化后的 CorrDiff 工作流程在 H100 上达到了预计 SOL 的 63%，在 B200 上达到 67%。这表明 GPU 利用率已处于较高水平，同时仍为未来的内核级优化保留了提升空间。

为进一步评估系统效率，我们对内核性能进行了可视化分析，结果如图5和图6所示。图中每个点代表一个内核，其横轴为NVIDIA Tensor核心利用率 (x-axis)和组合的DRAM读/写带宽利用率 (y-axis)。点的大小表示该内核在总运行时间中所占的比例，从而突出对整体性能具有显著影响的关键操作。

靠近上边缘或右边缘的内核通常经过良好优化，能够充分高效地利用计算或内存资源。相比之下，位于左下方象限的核函数在两方面均未得到充分利用，因而具有较大的优化潜力。该可视化图有助于清晰了解运行时的分布情况，并为提升 GPU 效率指明改进方向。

Scatter plot showing kernel distribution by Tensor Core utilization (x-axis) and DRAM bandwidth utilization (y-axis) for baseline patch-based CorrDiff. Each dot’s size represents the kernel’s share of total runtime. Most kernels are underutilized in both Tensor Core and DRAM bandwidth. Scatter plot showing kernel distribution by Tensor Core utilization (x-axis) and DRAM bandwidth utilization (y-axis) for baseline patch-based CorrDiff. Each dot’s size represents the kernel’s share of total runtime. Most kernels are underutilized in both Tensor Core and DRAM bandwidth.
图 6：B200 上基于 FP32 补丁的 CorrDiff 基准测试内核利用率

图6展示了基准CorrDiff实现在Tensor Core利用率和DRAM带宽利用率方面的内核分布情况。在FP32精度的未优化工作流程中，超过95%的时间消耗在低利用率的内核上，这些内核的DRAM（读+写）和Tensor Core利用率均低于60%。

多数运行时主导内核集中在原点附近，表明其 DRAM 和 Tensor Core 的利用率普遍较低。仅有少数内核靠近图示的上边界或右边界，显示出明显的内存或计算瓶颈。在基于 B200 的 SOL 系统中，未经优化的美国 CorrDiff 工作流程仅占 1.23%。

Scatter plot showing kernel distribution by Tensor Core utilization (x-axis) and DRAM bandwidth utilization (y-axis) for optimized patch-based CorrDiff. Each dot’s size represents the kernel’s share of total runtime. Most kernels show higher utilization of Tensor Cores or DRAM bandwidth. Scatter plot showing kernel distribution by Tensor Core utilization (x-axis) and DRAM bandwidth utilization (y-axis) for optimized patch-based CorrDiff. Each dot’s size represents the kernel’s share of total runtime. Most kernels show higher utilization of Tensor Cores or DRAM bandwidth.
图 7：在 B200 上优化的基于补丁的 BF16 CorrDiff 内核利用率分析

图7展示了优化实现中内核在Tensor Core利用率和DRAM带宽利用率上的分布情况。在采用AMP-BF16训练的优化工作流程中，位于左上角或右下角区域的内核占比较高，表明其性能表现和GPU利用率较为理想。SOL在B200上对CorrDiff进行了优化，当前优化水平达到67%。尽管整体工作流程已取得显著改进，部分内核仍存在进一步加速的空间。

CorrDiff 推理优化

许多训练阶段的优化方法同样适用于推理阶段。我们进一步提出了更多针对推理的优化策略，以显著提升整体性能。

采用基于分块的多扩散方法
优化后的 CorrDiff 对重叠的空间分块进行降噪后再聚合。最初，27.1% 的总运行时间消耗在 im2col 的折叠与展开操作上。通过预计算每个像素的重叠次数，并结合 torch.compile() 加速剩余的折叠与展开步骤，成功消除了 im2col 带来的性能瓶颈，使整体速度提升了 7.86 倍。

采用确定性 Euler 采样器（12 步）
采用确定性 Euler 采样器（12 步）替代原始随机采样器（在 Heun 求解器和二阶校正中使用 18 个降噪步骤），在不降低输出质量的前提下，将降噪步骤减少至 12 步。这一优化使 Hopper 和 Blackwell 架构的推理速度进一步提升了 2.8%。使用 12 步确定性采样器后，H100 和 B200 的最终推理速度分别达到原来的 21.94 倍和 54.87 倍。

博客中描述的一些优化方法通常也适用于扩散模型，而另一些则专门针对基于补丁的方法。因此，这些模型可以迁移到 PhysicsNeMo 中的其他模型，并用于开发超越天气降尺度的各类解决方案。

入门指南

在 PhysicsNeMo 中训练与推理 CorrDiff 模型：基于 PhysicsNeMo CorrDiff 文档的指南

  • 如需使用优化的代码库进行训练，请参照 CorrDiff 库的 readme 文件中的说明，并在所选训练 YAML 配置文件的 training.perf 部分中设置以下选项：
fp_optimizations: amp-bf16
use_apex_gn: True
torch_compile: True
profile_mode: False
  • 如需使用优化后的代码库运行推理，请参照 CorrDiff 仓库的 README 说明，并在所选生成配置的 generation.perf 部分中设置以下选项：
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 的构建，或在加载容器后于本地完成构建。
    • 具体步骤为：克隆 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 训练方法，并在 PhysicsNeMo 和 Earth2Studio 中执行优化后的工作流。

视频1：在CONUS区域可视化基于补丁的CorrDiff降采样，速度提升55倍

 

 

 

标签

代理式 AI/生成式 AI | 模拟/建模/设计 | Advanced Technical | Deep dive | Earth-2 | HPC/科学计算 | PhysicsNeMo | 气候/天气/海洋建模 | 精选

关于作者

Alicia Sui 是 NVIDIA Earth-2 团队的一名开发者技术工程师，专注于开发和优化 GPU 加速的 AI 解决方案，推动 AI 在科学领域的发展。Alicia 拥有学士学位。加州大学圣地亚哥分校的数学 CS 和认知科学专业以及理科硕士学位。计算数据科学博士学位。
Akshay Subramaniam 是 NVIDIA 的高级 AI 开发者工程师。他一直致力于在 GPU 上加速 AI 和 ML 应用，以及解决物理学和深度学习的交叉问题。
Karthik Kashinath 是开发者技术中 HPC + AI 的首席工程师和科学家。他是 NVIDIA 地球 2 计划的联合负责人。他的工作使用机器学习来模拟高分辨率的天气和气候数据，构建数字孪生，并提供可操作的气候信息。 Karthik 拥有马德拉斯理工学院、斯坦福大学和剑桥大学的工程学位。