2026/5/13 22:44:43
网站建设
项目流程
怎么做网站弹窗,建筑装饰网站模板,杭州企业seo网站优化,如何制作一个二维码随着人工智能天气和气候预测模型迅速得到采用#xff0c;由某中心提供的Earth-2平台提供了一系列库和工具#xff0c;可利用GPU优化的软件栈来加速解决方案。降尺度#xff08;Downscaling#xff09;是一项精炼粗分辨率#xff08;25公里尺度#xff09;天气数据的任务由某中心提供的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为训练启用AMP-BF16原始训练方案使用FP32精度。此处为训练启用了带有BF16的自动混合精度AMP以减少内存使用并提高吞吐量同时不影响数值稳定性从而实现了相对于基线的2.03倍加速。优化2使用多迭代块分摊回归成本在原始的基于块的训练工作流程中用于扩散模型训练的每个448×448块样本都需要对整个1056×1792的CONUS区域进行回归模型推理。这导致扩散模型训练吞吐量受限于回归模型推理。我们通过缓存回归输出并在每个时间戳的16个扩散块中重复使用提高了效率。这提供了更广泛的空间覆盖同时更有效地分摊了回归成本从而实现了相对于基线的12.33倍加速。优化3使用Apex GroupNorm消除数据转置训练管道最初使用默认的NCHW内存布局这会在卷积前后触发代价高昂的隐式内存转置。将模型和输入张量切换为NHWC通道最后格式使其与cuDNN的首选布局对齐。然而PyTorch的GroupNorm操作不支持通道最后格式。为了防止转置并保持数据为通道最后格式以实现更高效的归一化内核我们将PyTorch GroupNorm替换为某中心的Apex GroupNorm。这消除了转置开销并实现了相对于基线的16.71倍加速。优化4融合GroupNorm与SiLU通过使用Apex将GroupNorm和SiLU激活融合到单个内核中减少了内核启动次数和全局内存访问次数。这提高了GPU利用率并实现了相对于基线的17.15倍加速。优化5在Apex GroupNorm中扩展通道维度支持一些CorrDiff层使用的通道维度不受Apex支持。我们扩展了对这些通道维度的支持为所有层解锁了融合功能。这将性能提升到相对于基线的19.74倍加速。优化6通过图编译进行内核融合我们使用torch.compile来融合剩余的元素级操作例如加法、乘法、sqrt、exp。这改善了调度减少了全局内存访问并削减了Python开销实现了相对于基线的25.51倍加速。优化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显示出最显著的增益。GPU版本精度回归 (ms/块)扩散 (ms/块)总运行时间 (ms/块)吞吐量 (块/秒)H100基线FP321204.0374.01578.00.63H100优化版本BF1610.60951.25961.8516.2B200优化版本BF164.73424.56429.2934.1表1. CorrDiff训练吞吐量比较“光速”Speed-of-Light分析为了评估优化后的CorrDiff工作流距离硬件性能上限有多近我们在H100和B200 GPU上进行了光速SOL分析。这通过评估GPU资源的使用效率提供了可达到性能的上限估计。估算SOL的步骤识别低利用率内核我们关注DRAM读/写带宽 60% 且 Tensor Core利用率 60%的内核。此类内核既不是内存受限也不是计算受限很可能是性能瓶颈。估算每个内核的潜力对于每个低利用率内核我们估算其在理想条件下的潜在加速比——即在充分利用DRAM带宽或Tensor Core活动的情况下。汇总总体加速比然后如果每个内核都优化到其理想性能我们计算假设的端到端加速比。计算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-bf16use_apex_gn:Truetorch_compile:Trueprofile_mode:Falseuse_fp16:Trueuse_apex_gn:Trueuse_torch_compile:Trueprofile_mode:Falseio_syncronous:True将profile_mode设置为False以获得优化性能因为NVTX注释会为torch.compile工作流引入图中断。要利用最新的Apex GroupNorm内核可以在PhysicsNeMo容器Dockerfile中构建Apex GroupNorm或者在加载PhysicsNeMo容器后在本地构建克隆Apex仓库并使用以下命令构建CFLAGS-g0NVCC_APPEND_FLAGS--threads 8\pipinstall\--no-build-isolation\--no-cache-dir\--disable-pip-version-check\--config-settings--build-option--group_norm.了解更多关于PhysicsNeMo中优化的CorrDiff训练的信息并在Earth2Studio中运行优化的工作流。视频1. 可视化在CONUS上基于块的CorrDiff降尺度并实现55倍加速更多精彩内容 请关注我的个人公众号 公众号办公AI智能小助手或者 我的个人博客 https://blog.qife122.com/对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号网络安全技术点滴分享