Платформа OlmoEarth: как Ai2 масштабирует анализ спутниковых данных на всю планету

Институт Allen Institute for AI (Ai2) анонсировал платформу OlmoEarth — специализированную инфраструктуру для масштабного применения геопространственных моделей (inference). Система позволяет обрабатывать десятки терабайт мультимодальных спутниковых данных и создавать карты континентального масштаба всего за один день. Стоимость такой обработки составляет доли цента за квадратный километр.

Исторически сложилось так, что создание качественных моделей — лишь часть решения проблемы. Базовые модели OlmoEarth были предварительно обучены на 10 терабайтах спутниковых снимков для задач мониторинга вырубки лесов, продовольственной безопасности и оценки риска лесных пожаров. Однако большинство некоммерческих и правительственных организаций в экологической сфере не обладают инженерными ресурсами для управления полным жизненным циклом машинного обучения: разметкой данных, тонкой настройкой (fine-tuning) и запуском моделей на больших площадях.

Главная сложность работы со спутниковыми снимками заключается не в самих вычислениях, а во вводе-выводе данных. В отличие от текстовых моделей (LLM), где на вход подается несколько килобайт текста, геопространственный анализ требует загрузки терабайтов информации. Снимки поступают от разных провайдеров, имеют разные проекции, разрешения и часто скрыты облаками. На практике подготовка данных занимает больше времени, чем непосредственно работа нейросети.

Чтобы оптимизировать процесс, инженеры Ai2 разделили выполнение каждой задачи на три этапа, привязанных к конкретному оборудованию. Сбор и предварительная обработка данных выполняются на центральных процессорах (CPU), так как эта задача требует высокой пропускной способности ввода-вывода. Непосредственно применение модели (forward pass) происходит на графических ускорителях (GPU). Постобработка, включающая склейку фрагментов карты и экспорт в форматы Zarr или GeoTIFF, снова возвращается на CPU. Это позволяет избежать простоя дорогостоящих GPU.

Масштабирование достигается за счет дробления географических регионов на независимые фрагменты. Например, при создании карты риска лесных пожаров для всей Северной Америки платформа одновременно задействовала 19 600 CPU и 994 GPU. Пропускная способность сети превысила 168 гигабайт в секунду. Благодаря такому распараллеливанию задача, которая потребовала бы 4737 часов последовательных вычислений, была выполнена за 30,5 часов реального времени.