Ai2的OlmoEarth平台将卫星推理转化为分布式系统问题

当艾伦人工智能研究所（Ai2）着手构建一个地理空间推理平台时，它遇到了其他机器学习系统很少遇到的障碍：卫星图像并不规整。一次微调任务可能移动数太字节的数据并运行数小时。输入包含多个光谱波段、传感器类型、时间步长和提供商，每个都有不同的投影和分辨率。云层遮挡像素。缺失波段是常态。输出是一张地图，这意味着每个预测必须与其相邻预测保持在同一地理网格上。

OlmoEarth平台（在Hugging Face上的一篇近期博文中详细介绍）是Ai2对这些约束的答案。它专为需要监测森林砍伐、作物健康或野火风险的政府和非政府组织设计，范围从国家到大陆尺度。该平台将推理首先视为分布式系统问题，其次才是建模问题。

三个阶段，与硬件匹配

核心见解在于GPU时间昂贵，大多数推理作业的时钟时间花在数据获取和准备上，而不是模型本身。平台将每个作业分为三个阶段，每个阶段分配给最适合的硬件。

数据获取和预处理运行在高I/O的CPU实例上。它们从Sentinel-2和Landsat等来源获取、重投影、对齐和归一化图像，然后以优化快速GPU加载的格式写入。推理阶段运行在GPU上，执行模型的前向传播，并将最小处理的输出直接流式传输到blob存储。后处理阶段，将逐窗口输出拼接成Zarr或GeoTIFF等格式的无缝栅格，再回到CPU。

多进程数据加载器确保每个GPU持续被馈送。这种划分防止系统最昂贵的硬件在较慢操作完成时空闲。

数据获取的瓶颈

获取正确的像素比处理它们更难。平台必须查询多个提供商的STAC目录以找到覆盖给定区域和时间范围的场景。对于光学图像，目标是最少云层的场景；对于合成孔径雷达，相关的极化通道很重要。一个大的作业可能生成数千个元数据查询，远超过ESA或Microsoft Planetary Computer等外部服务设计能同时处理的量。

Ai2通过维护自己的卫星图像索引来规避此问题，该索引通过AWS开放数据源的SNS通知和每隔几分钟轮询上游目录来更新。运行时，平台选择最佳可用来源，并针对云优化格式（如COG或Zarr）执行窗口读取，仅检索给定分区所需的字节，而不是下载整个场景。同一个索引也支持注释工具，通过相同的窗口读取系统从任何索引场景提供瓦片。

行星级的弹性

在该平台目标覆盖的规模下，失败不是异常，而是设计输入。提供商变慢或离线。元数据指向的图像可能缺失一个波段。云层覆盖导致可用观测太少。任务崩溃。

OlmoEarth平台将每个任务视为幂等和可重入的。对于每个地理分区和阶段，它提供一个运行带有任务参数的Docker容器的虚拟机。运行时执行工作，返回结果，然后关闭。如果失败，系统自动重试，并在可用时回退到其他提供商。一个单独的监控进程检测到已停滞的运行器并重新启动其任务。

具体成果：155倍加速，成本不到一分钱

该平台迄今为止最引人注目的演示是一张覆盖整个北美的野火风险地图。在峰值时，运行使用了大约19,600个CPU和994个GPU并行，网络吞吐量超过168 GB/s。这种并行级别将估计的4,737小时串行计算减少到大约30.5小时的挂钟时间，实现了155倍加速。

成本：每平方公里几美分。该平台已经可以在大约一天内对大陆大小的区域进行推理，处理数十太字节的图像。这些经济效益还得到2026年中发布的OlmoEarth v1.1模型补充，该模型通过更智能的令牌合并技术，将计算成本比原始v1降低了多达3倍，同时保持了相似的性能。

当然，并行不是无限的。平台为每个作业提供了几个可调参数：扇出受限于云配额；输出分辨率在数据量和细节之间权衡；模型大小在GPU时间和准确性之间权衡；缓存原始图像在重复运行同一区域时以存储换取速度。正确的设置取决于任务和预算。

下一步计划

Ai2的路线图根据合作伙伴反馈，包括由索引中新图像触发的自动模型运行；变化检测和警报，将森林砍伐或洪水作为推送通知呈现，而不是需要人工检查的栅格；以及代理工具，降低非专业用户执行以往需要经验丰富的ML研究员才能完成的工作的门槛。团队还在开发专门的嵌入模型；全局规模的预计算嵌入可以取代许多任务中对原始图像的全前向传播，使工作负载更快更便宜。

平台目前运行在Google Cloud上，但架构设计支持多个云以及部署在合作伙伴自己的计算环境中。对于从事保护、粮食安全、灾害响应和气候工作的组织来说，这种灵活性很重要。

地理空间基础模型仍是一项新兴技术，这些团队需要了解地球所需的东西与其预算和技术资源允许之间的差距依然巨大。Ai2押注将推理首先视为基础设施问题有助于缩小差距。