Как запустить ИИ‑прогноз погоды AIFS 2.0 от ECMWF на своём железе

Европейский центр среднесрочных прогнозов погоды (ECMWF) выпустил открытые веса модели искусственного интеллекта AIFS Single 2.0, которая способна предсказывать погоду на 15 дней вперёд. Однако открытые веса — это только половина дела: для запуска модели требуется специфическое окружение и GPU, что отпугивает многих энтузиастов. В этой статье мы покажем, как запустить AIFS Single 2.0 через Hugging Face Jobs или локально, обойти зависимость от flash-attn и получить собственный прогноз погоды даже на CPU или Apple MPS.

Что такое AIFS Single 2.0 и почему это важно

AIFS (Artificial Intelligence Forecasting System) — это модель машинного обучения, разработанная ECMWF для численного прогноза погоды. В отличие от традиционных физических моделей, AIFS обучается на огромных массивах исторических данных и затем делает прогнозы, используя нейросети. Версия Single 2.0 — это однопрогоночная модель, которая выдаёт детерминированный прогноз на 15 суток. Она основана на архитектуре трансформера и обучена на данных ERA5, что позволяет ей конкурировать с классическими моделями, такими как IFS.

Важность AIFS в том, что она делает высокотехнологичные прогнозы доступными для широкого круга пользователей. Если раньше прогноз погоды высокой точности был прерогативой крупных метеослужб с суперкомпьютерами, то теперь любой энтузиаст с достаточно мощным GPU может получить прогноз, сравнимый по качеству с официальными. Однако на практике запуск AIFS сопряжён с трудностями: требуются определённые версии библиотек, поддержка flash-attention и наличие GPU с достаточным объёмом памяти.

Как запустить AIFS Single 2.0 через Hugging Face Jobs

Hugging Face предлагает сервис Jobs, который позволяет запускать модели в облачной среде с предустановленным окружением. Это самый простой способ начать работу с AIFS, особенно если у вас нет мощного локального GPU. Для запуска через Hugging Face Jobs достаточно создать новый Job, выбрать подходящий GPU (рекомендуется A100 или V100 с минимум 16 ГБ памяти), указать репозиторий модели и скрипт запуска.

Основная сложность — зависимость от flash-attn, которая не всегда устанавливается корректно. В статье на Хабре автор предлагает обходной путь: использовать версию модели без flash-attn, либо установить её вручную. На Hugging Face Jobs окружение можно настроить через Docker-образ, где вы заранее установите все зависимости. Это позволяет избежать ошибок, связанных с несовместимостью библиотек.

Какие требования к GPU для запуска AIFS 2.0?

Для запуска AIFS Single 2.0 рекомендуется GPU с минимум 16 ГБ видеопамяти, например, NVIDIA A100 или V100. Однако, как показывает практика, можно обойтись и менее мощными видеокартами, если использовать оптимизации и урезать некоторые параметры. Например, автор статьи на Хабре успешно запускал модель на MacBook Pro с чипом M1 Pro, используя MPS, что говорит о том, что требования не такие уж жёсткие. Главное — правильно настроить окружение и обойти зависимость от flash-attn, которая может быть проблемой на некоторых GPU.

Предыстория и контекст

ECMWF начала развивать направление ИИ-прогнозирования несколько лет назад, видя успехи таких моделей, как GraphCast от Google DeepMind и Pangu-Weather от Huawei. AIFS стала ответом ECMWF на этот вызов. Первая версия AIFS была представлена в 2023 году, а в 2024 году вышла версия 2.0, которая показала улучшенную точность и скорость. Открытые веса модели — это часть политики ECMWF по демократизации доступа к метеорологическим данным и инструментам.

Однако открытые веса не означают, что модель легко запустить. Архитектура AIFS использует механизм внимания, который требует специализированных оптимизаций, таких как flash-attention, для эффективной работы на GPU. Без этих оптимизаций модель может работать крайне медленно или вообще не запуститься на некоторых видеокартах. Это создаёт барьер для многих потенциальных пользователей.

Чем отличается AIFS Single 2.0 от предыдущих версий?

AIFS Single 2.0 отличается от предыдущих версий прежде всего тем, что это однопрогоночная модель, то есть она выдаёт один детерминированный прогноз, а не ансамбль. Это делает её более лёгкой и быстрой по сравнению с ансамблевыми версиями, которые требуют многократных прогонов. Также в версии 2.0 улучшена архитектура: увеличено количество слоёв и параметров, что повышает точность, но и требования к вычислительным ресурсам. По сравнению с первой версией, 2.0 показывает лучшие результаты на коротких и средних сроках прогноза.

Технические подробности запуска

Для локального запуска AIFS Single 2.0 вам понадобится Python 3.10 или выше, PyTorch с поддержкой CUDA (или MPS для Mac), и библиотека transformers от Hugging Face. Основная проблема — flash-attn, которая требует специфической компиляции под вашу видеокарту. Автор статьи на Хабре предлагает решение: использовать версию модели, которая не зависит от flash-attn, либо установить её из исходников с правильными флагами.

Если у вас нет GPU, вы можете запустить модель на CPU, но это будет медленно. Например, прогноз на 15 дней может занять несколько часов на CPU, в то время как на GPU — несколько минут. Для Apple MPS (Metal Performance Shaders) ситуация лучше: модель можно запустить с использованием MPS-бэкенда, но некоторые операции могут не поддерживаться, и потребуется адаптация кода.

На примере из статьи: автор запускал модель на MacBook Pro с чипом M1 Pro, используя MPS. Для этого пришлось вручную заменить некоторые операции, связанные с flash-attention, на стандартные. В результате прогноз на 15 дней занял около 10 минут, что приемлемо для личного использования.

Как обойти зависимость от flash-attn в AIFS 2.0?

Зависимость от flash-attn — одна из главных проблем при запуске AIFS Single 2.0. Чтобы её обойти, можно использовать версию модели, которая не требует flash-attn, либо установить flash-attn вручную с учётом особенностей вашей системы. В некоторых случаях помогает отключение flash-attn в коде модели и замена его на стандартную реализацию attention из PyTorch. Это может снизить производительность, но позволит запустить модель на более широком спектре устройств, включая CPU и MPS.

Кого затронет и как

Прежде всего, AIFS Single 2.0 интересна разработчикам метеорологических приложений, исследователям, а также энтузиастам, которые хотят иметь собственный прогноз погоды без обращения к коммерческим сервисам. Для них открывается возможность создавать локальные системы прогнозирования, адаптированные под конкретные нужды, например, для сельского хозяйства, логистики или просто для личного использования.

В России и СНГ эта тема также актуальна, так как доступ к зарубежным метеорологическим сервисам может быть ограничен, а локальные модели позволяют получить независимый прогноз. Кроме того, AIFS может быть использована в образовательных целях: студенты и преподаватели могут изучать, как работают нейросетевые прогнозы погоды.

Однако есть и ограничения: для полноценного использования модели нужны значительные вычислительные ресурсы, особенно для ансамблевых версий. Это может отпугнуть обычных пользователей, но для профессионалов это вполне решаемая задача.

Что будет дальше

ECMWF продолжает развивать AIFS, и можно ожидать появления новых версий с улучшенной точностью и скоростью. Возможно, в будущем модель будет оптимизирована для работы на более широком спектре устройств, включая мобильные. Также вероятно, что сообщество разработает более простые инструменты для запуска AIFS, снижая порог входа.

Уже сейчас есть примеры интеграции AIFS в открытые проекты, такие как Open-Meteo, что делает прогнозы доступными через API. Это может привести к появлению новых сервисов, которые используют ИИ-прогнозы для различных отраслей.

Итог

AIFS Single 2.0 от ECMWF — это мощный инструмент для прогнозирования погоды, который теперь доступен каждому благодаря открытым весам. Несмотря на технические сложности, запустить модель можно даже на CPU или Apple MPS, если следовать инструкциям. Это открывает новые возможности для разработчиков и исследователей, а также делает высокоточные прогнозы доступными для всех. Следите за развитием AIFS и экспериментируйте — возможно, вы создадите свой собственный сервис прогноза погоды.