OlmoEarth：地球级地理空间推理平台

OlmoEarth 模型是我们的地球观测基础模型系列，在约 10 TB 的多模态卫星数据上进行了预训练。政府机构、非政府组织和其他以使命驱动的组织已经在采用 OlmoEarth，用于包括森林砍伐监测、粮食安全和野火风险等应用。

在 Ai2，我们知道如何训练和发布强大的开源模型，对于拥有强大工程团队的组织来说，一个开源模型就足够让他们运行了。但是环保领域的大多数组织——最有能力应用这些模型的那些——没有基础设施或工程团队来管理完整的生命周期：标签数据、微调模型和运行大规模推理。我们在 Skylight 和 EarthRanger 等平台上运营了十多年，这是世界各地用户每天都依赖的软件，所以它必须每天都能正常工作。那段经历教会了我们什么是产生影响：以适当的时间和地点成本效益地运行模型、监控性能、将原始输出转化为可操作的见解，以及验证这些输出是否推动了合作伙伴想要的结果。

这就是为什么我们构建了 OlmoEarth Platform：从微调和评估到大规模推理的地理空间模型基础设施。

这种规模的推理呈现出自身的一系列挑战。卫星图像必须在多个提供商之间查找和访问，跨投影和分辨率对齐，并高效处理。然后结果必须缝合成地理上一致的地图，同时基础设施从分布式计算的常规故障中恢复。

目前，该平台可以在大约一天内跨大陆级别的区域运行推理，处理数十 TB 的图像，成本为每平方公里几分钱。开发它意味着面对一系列工程挑战，其他从事大规模地理空间系统工作的人也可能会遇到。这篇文章讨论了这些挑战和我们找到的解决方案。

为什么卫星推理很有挑战性

大多数 ML 模型接受几 MB 的数据并在一秒钟内产生结果——想想 LLM 处理一段文本或计算机视觉模型分析智能手机的照片。地球观测推理的规模完全不同——为了最大限度地提高性能而微调基础模型的单个工作可以处理数 TB 数据并运行数小时。输入可能跨越多个光谱波段、传感器类型和时间步长，覆盖大地理区域。它们可能来自多个提供商，每个都使用不同的投影和分辨率，可能包括缺失或被云遮挡的观测。输出本身就是一张地图，所以每个预测都必须与周围区域保持相同的投影和坐标网格精确对齐。

即使获取数据也可能是一个重大挑战。预测工作经常花费比运行模型本身更多的时间来下载和准备图像，这使得高效的数据管道至关重要。这些管道必须处理高容量 I/O，同时提供重新投影和重新采样图像所需的计算。

为任务选择合适的硬件

由于数据获取和准备通常占据推理工作运行时间的主导，将该工作分配给 GPU 会导致系统最昂贵的硬件做更适合 CPU 的任务。因此，我们将每个工作分为三个阶段，每个阶段与不同的硬件配置相匹配：

数据获取和预处理（CPU，高 I/O）：获取、重新投影、对齐和规范化图像，然后以针对推理期间快速加载的格式写入。

推理（GPU）：运行模型的前向传递并将最少处理的输出直接写入存储。

后处理（CPU）：将每窗口的输出缝合在一起，应用掩码或重新缩放，并以用户友好的格式（如 Zarr、GeoTIFF 或 GeoJSON）导出。

OlmoEarth Platform 将这些阶段分布在许多机器上，同时保持 GPU 充分利用。多进程数据加载器不断为每个 GPU 提供数据，而已完成的输出直接流式传输到 blob 存储。

一个请求，数百个 workers，以及数千个进程

OlmoEarth Run 是平台用于大规模推理工作的执行层，它将每个工作覆盖的地理区域划分为大小适合个人计算实例（workers）的分区，然后将这些分区进一步细分为 OlmoEarth 模型处理的更小窗口。因为每个窗口可以在单独的前向传递中独立处理，一张地图的一部分的工作不需要等待另一部分。

实际上，州级别的区域可能变成大约一百个分区，而大陆级别的运行可能变成数千个。相邻分区略有重叠，我们在组装输出时调和这种重叠，以确保最终栅格中没有出现接缝。

由于分区是独立的，同一阶段可以同时跨数千个计算实例运行。我们最近使用这种方法生成了覆盖整个北美的野火风险地图。在峰值时，该运行同时使用了大约 19,600 个 CPU 和 994 个 GPU，网络吞吐量超过 168 GB/s。这种并行程度将估计的 4,737 小时串行计算减少到约 30.5 小时的挂钟时间——155 倍的加速。

但是并行度不是无限的。更多的 workers 会冲击云配额，所以并行性是每次运行的一个可调参数，是我们可以在各个工作上调整的几个参数之一。输出分辨率权衡数据量和计算与细节；模型大小权衡 GPU 时间和准确性；缓存原始图像权衡存储和在同一区域上重复运行的速度。正确的设置取决于手头的任务——和预算。

查找和获取合适的像素

给定地理区域和时间范围，平台首先必须确定哪些卫星场景应该馈送给模型。这意味着要识别在所有提供商中捕获了什么、在哪里以及何时，这些提供商有不同的目录、格式和发布延迟。选择标准还取决于来源：对于 Sentinel-2 等光学图像，我们通常希望获得可用的云层最少的场景，而对于合成孔径雷达，可用的极化通道可能更重要。

我们尽可能依赖公开 STAC 目录和开放标准。但是大型推理工作可以同时生成数千个元数据查询——远远超过外部服务（如 ESA 或 Microsoft Planetary Computer 的 STAC API）设计用来并发处理的数量。

为了避免使这些服务不堪重负，OlmoEarth Platform 维护自己的元数据索引，随着新图像的发布而更新。对于通过 AWS Open Data 托管的数据集，我们对每个新场景都会收到一个 SNS 通知。当提供商不提供更改流时，我们每隔几分钟轮询其上游索引。结果是，我们对外部服务的请求遵循新发布的稳定步伐，而不是大型推理工作产生的突然爆发。

每个索引条目存储场景元数据以及指向底层像素可用的每个位置的指针。在运行时，平台选择最佳源并对 COG 或 Zarr 等云优化格式执行窗口读取，仅检索给定分区所需的字节，而不是下载整个场景。

该索引还支持我们的注释工具。因为它维护指向云优化格式的 Sentinel-1、Sentinel-2、Landsat 和 NISAR 图像的指针，我们可以通过相同的窗口读取系统从任何索引场景提供瓦片，无需构建单独的摄取管道。

使该工作流最容易的提供商共享了三个特性，我们建议作为发布地球观测数据的最佳实践：当新图像可用时基于队列的通知、在主要云平台上的存储且没有定制的速率限制或可用性瓶颈、以及支持范围读取的云优化格式。

大规模故障处理

OlmoEarth Platform 设计用于自动从故障恢复。对于阶段和地理分区内的每个任务，它动态配置运行我们的 runner Docker 容器的虚拟机。runner 检索任务参数、执行工作、返回结果并关闭。由于每个任务都是可重入且幂等的，间歇性故障可以通过重新运行来安全处理。

在这种规模下，故障是预期的：提供商可能很慢或暂时不可用；元数据可能表明图像存在，即使所需波段或窗口缺失；云覆盖可能留下太少可用观测；或任务可能直接崩溃。平台通过任务跟踪、自动重试、在可用时回退到备用提供商，以及区分可重试和致命错误来响应。单独的监控过程检测已停止或停滞的 runner，并重新启动其任务。

我们的路线图由合作伙伴确定的差距和他们最需要的功能塑造而成。在我们正在研究的领域中：

自动化模型运行。提前安排推理工作或在图像索引在感兴趣区域注册新场景时触发。

变化检测和警报。当他们监测的景观发生变化时通知用户，以便森林砍伐或洪水等事件表现为警报，而不是必须有人去寻找和检查的栅格。

Agent 工具和界面。Agent 可以降低使用地理空间模型的门槛，从数据策划和特征工程到识别改进微调模型的方法。我们希望任何技术水平的用户都能做以前需要经验丰富的 ML 研究人员才能做的工作。

更快的模型。与我们的研究团队合作，开发更高效的架构，以减少每个窗口的 GPU 时间。

更多模态。将新的卫星传感器和数据源添加到 OlmoEarth 模型和支持它们的图像索引。我们目前专注于合并天气数据（ERA-5）和提供更多环境因素细节的卫星。

嵌入。开发专门的嵌入模型并在全球范围内预计算嵌入。对于许多任务，针对这些嵌入运行推理可以替代对原始图像的完整前向传递，使工作负载更快且成本更低。微调和直接推理对于具有挑战性任务的最大性能仍然很重要，但嵌入可以打开更广泛范围的高效应用。

在任何地方运行。OlmoEarth Run 仅需要能够运行 Docker 镜像并访问 blob 存储的虚拟机。我们目前在 Google Cloud 上运营它，但该架构设计用于支持多个云和在合作伙伴自己的账户和计算环境内的部署。

这仅仅是开始。地理空间基础模型，特别是其操作化，仍然是一项新兴技术，许多最有能力使用它们的组织——那些从事保护、粮食安全、灾难应对和气候工作的组织——从未能够访问这样的基础设施。这些团队需要理解地球的内容与他们的预算和技术资源允许他们做的内容之间的差距仍然很大。

我们正在构建 OlmoEarth 来帮助弥补这个差距。