观潮 · 模式预报与气象 AI 观察
有人把 WRF v4 用 JAX 重写跑在了 GPU 上：目标是省电，测出来的结论对自己不太成立
这条帖子的价值不在"WRF 有 GPU 版本了"，而在它把 GPU 加速模式这条路线的成本结构讲得很具体，包括失败的那部分。作者是业余做项目的，明确说自己在气象和 WRF 代码侧"几乎一无所知"，所以下面所有关于性能的数字都要按这个前提读。
背景动机很具体：作者想给加那利群岛做一个免费的、真正能用的天气预报，理由是当地地形陡峭、局地气候多变，他用 1 km WRF 网格做物理骨架搭了个物理–机器学习混合方案。机器学习那部分很快，但 WRF 物理部分把他那台老工作站压垮了，而这套东西需要每晚跑、未来最好还能只靠太阳能供电。所以他的真实需求是"比 WRF-v4 快得多、耗电少得多"，而不是"最快"。
---
项目是怎么做出来的
作者给出的答案是 agent swarm——一队 AI agent 协同完成。他自述由 Opus 4.8 Max 组成、并与 GPT 5.5 XHigh 对抗配合，由一个 Opus 4.8 XHigh 管理器制定和维护路线图。
技术路线上有个转折值得记：在最初一轮 kernel 较劲、并读了一些关于 nvfortran 移植及其可预期扩展性限制的论文之后，这些 agent 自主决定通过一个 XLA/JAX dycore kernel 把 WRF 从零完全重写。 之后是大约 2 到 3 天让 kernel 变快变稳，再加约 3 周、约 3000 次 sprint（agent 运行）、约 2000 次 commit，才到"核心元素都已实现、测试、文档化、稳定并做过显存优化"的状态。
这个过程的性质，作者自己说得很清楚：主要目的是学习并观察最新一批 AI 模型能否自组织起来处理这样一个出了名难缠的代码库。 他也提到后续还有 Fable 5 做了最后几天 kernel 的速度与显存优化。
---
目标达成情况：加速没达标，省电达标了
这是帖子里最有信息量的一段，作者没有粉饰。
目标 · 结果
相比他自己 CPU 上的 WRF，嵌套 1 km 加那利算例提速至少 2 倍 · **未达成**，加速比约 **135%**（对 12 核 CPU）
能耗降到 1/3 · **达成**
作者对为什么没达标的解释是：问题不在 kernel 不够优化，而是那个小网格配"tiny" RTX 5090 根本不适合这套架构。 具体两条：
第一，5090 和所有消费级 GPU 的 FP64 性能都极差。他试过把压力 kernel 的核心机制降到 FP32、并用各种计算技巧保持长期稳定，到目前为止都失败了。
第二，kernel 初始化开销对小网格计算影响很大。
他对大网格给了理论预期（明确是"理论上"，他本人测不了）：B200、GB300、乃至 NVL72-GB300 按设计应该都能支持，微 kernel 能建的并行数量没有上限；一个 NVL72-GB300 应该能在单系统内存里装下整个 1 km 分辨率的全球数据并以极高效率计算。大网格和真实 GPU 情形下每千瓦时的算力应该在 3 到 10 倍区间（GPU 与 CPU 的差距还在拉大）。
后续更新里，0.18.2 版本做了显著的显存使用改进，同算例下显存与内存占用现在大体与 WRF v4 持平，这抵消了消费级 GPU 相对 WRF-CPU 在加速比上的理论上限——因为现在能装更大的网格了，他把上限提到 RTX 5090 对 12 核 CPU 上 WRF v4 的 200% 到 300%，但明确写了没有测试。
---
验证做到哪一步
版本演进和验证范围作者交代得比较清楚。v0.17 还不完备，缺 40 多个冷门方案（scheme）。后续几版的记录：
- 身份测试（identity tests）他跑过的都位级一致（在可能的情况下），长期真实世界模拟中所有变量保持稳定并接近 WRF v4 解。
- 仓库里带一个"hardcore test"，模拟瑞士阿尔卑斯一个典型冬日，需要 32 GB 显存运行，他说结果看起来与 WRF v4 解一致（RAINNC 和 QVAPOR 的累积误差按其描述接受，以及一个正在飞行中的 kernel 层修复）。
- 加那利测试（他个人唯一需要的那个）与 WRF v4 数值全绿。 他还做了 30 次回报（hindcast），用 TOST 与数百个真实气象站比较，技巧与 WRF 相同，因此对他来说够用。
- v0.18.1 时称对 WRF v4 已功能完备；v0.23 时是一个全 FP32 版本、完成度约 60%；v0.3.0 用了一个新 kernel，作者称对消费级 GPU 快得多、对大 GPU 也快得多，但仍未功能完备。
一个诚实的自评："a fool with a tool is still only a fool" ——他把自己让 AI 去设计一个自己测不了的大系统优化的做法称为愚蠢。
---
后续转向与给别人的建议
作者明确说自己的个人目标只部分达成，正在转向生成式 AI 方法，训练一个降尺度生成模型，从 3 km 或 9 km 数据生成接近 1 km 技巧的原始场预报——理由是这对他的问题更合适，应该几分钟而不是几小时就能在工作站上跑完。
但他也说愿意再发一版纯物理的 WRF GPU 重写，或者把这个项目直接交出去，前提是有人有兴趣。
回复他的另一位用户（想跑尚普兰湖地区的 1 km 嵌套模式，因为 3 km HRRR 分辨不出湖周围地形）得到了几条具体建议：
- 文档是"agent 优先、agent 会替你做完一切"式的；作者说他已经基本不再让 agent 优化 kernel，一是没时间没 token，二是他（或者说那些 AI）已经为他的区域建好了一个不错的 U-Net 加时空 Transformer 代理模式。
- 先看 Google Weather Next 3，几周前发布，小区域可以免费或在 BigQuery 免费额度内拿到 0.05° 和 0.01°（温度和风辐照度）；作者说这是 AI 预报，"在我用的案例里远好于 ECMWF、GFS、AIFS 等等"。他强调要打败它是有难度的：只有在用完整的 WN3 cube 作输入、并对自己区域做了非常细致的地形调优时，他才用 WRF 1 km 超过它。
- 全 FP32 的 WRF 移植 v0.23 完成度约 60%，但作者当时没有 token 去完成；他说目前只有 Astra 和 Opus 5.5 看起来合适，这是"AI 当前能力边缘上的硬 kernel 重写"。
- 之后他完成了并发布了全 FP32 版本 0.3.0，说还有一个版本会再带来约 5% 的速度提升，之后就停止优化，因为它已充分满足他的用途。
FP64 这一条对业务模式用户有直接含义：如果这个说法成立，消费级显卡跑 WRF 的路主要被 FP64 卡住，而不是被 kernel 质量卡住。 作者关于 FP32 的多番尝试失败、以及最终通过重写 kernel 才拿到可用的 FP32 版本，也说明这不是一个改个精度参数就能绕过的开关。
---
作者承认的限制
这一段在帖子里是明写的，不需要转译。
验证的规模根本不够。 作者直接写：我没有时间、没有算力、没有任何资金去在那个"科学上或商业上有意思"的区域做大规模验证。他也说不知道这东西对这里到底有没有意义，或者是不是大家已经有更好的方案。
他本人不是这个领域的从业者。 他说得很直接：关于气象和 WRF 代码侧他几乎不知道自己做的是什么，这是个全新的领域，他的日常工作和这个完全无关。他能提供的最大价值是希望有人有资源和知识去做规模化测试并给出完整报告，好让 agent 们最终构建出 1.0.0 版本。
这个项目只是他某个业余项目的一个子项目，无资金、无时间。 他靠公司给的免费 token 和家里工作站的夜间算力完成。
性能预期明确标注为未测试。 3 到 10 倍每千瓦时的算力、大网格与真实 GPU 上的表现、200% 到 300% 的加速比上限，全部标注为理论上、无法测试。他还说消费级 GPU 永远不会接近那个最佳使用场景（大网格 + 原生大 FP64 的 GPU）。
FP64 到 FP32 的降精度尝试全部失败。 原话是"到目前为止各种计算技巧都失败了"。
对代理模式的评价来自单个区域、单个人的验证。 他描述的"技巧与 WRF 相同"的结论基于 30 次回报和数百个站点，用的是 TOST 检验；他没有给出技巧分数的具体数值。
---
一句话判断
这条帖子对业务用户的价值在两条否定性结论上：一是消费级显卡因为 FP64 太弱而不适合这套架构（大网格 + 真 GPU 才成立，且作者测不了），二是作者自己最终转向了降尺度生成模型，因为它"几分钟而不是几小时"。 如果你在评估要不要把 WRF 挪到 GPU 上，这两条比"135% 加速"这个数字更有决策价值——后者是针对一个作者明说不适合的硬件配置测出来的。