ИИ ускоряет прогнозы погоды супер-разрешением

Платформа NVIDIA Earth-2 предлагает библиотеки и инструменты для ускорения решений в области прогнозирования погоды и климата с помощью оптимизированного для GPU программного стека. Процесс детализации, который уточняет данные о погоде с грубого разрешения (масштаб 25 км), позволяет национальным метеорологическим службам предоставлять высокоточные прогнозы для сельского хозяйства, энергетики, транспорта и подготовки к катастрофам на пространственном разрешении, достаточном для принятия обоснованных решений и планирования.

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

CorrDiff набрал популярность по всему миру для разных задач, показывая свою универсальность и влияние в сферах, где нужны детальные данные о погоде:

• The Weather Company (TWC) для поддержки отраслей сельского хозяйства, энергетики и авиации.
• G42 для улучшения прогнозов смога и пыльных бурь на Ближнем Востоке.
• Tomorrow.io для повышения точности прогнозов на масштабе штормов, включая предсказания пожарной погоды и порывов ветра, влияющих на работу железных дорог.

В этой статье демонстрируются оптимизации производительности и улучшения для обучения и вывода CorrDiff, внедренные в два инструмента стека Earth-2: NVIDIA PhysicsNeMo и NVIDIA Earth2Studio. Достигнутый ускорение более чем в 50 раз по сравнению с базовыми показателями для обучения и вывода позволяет:

• Масштабировать обучение на основе патчей для всей планеты менее чем за 3000 GPU-часов.
• Снижать затраты на обучение для большинства стран до порядка 100 GPU-часов.
• Обучать модель для континентальной части США (CONUS) менее чем за 1000 GPU-часов.
• Проводить тонкую настройку и специализированное обучение, делая ИИ-прогнозы на км-масштабе доступными для всех.
• Выполнять вывод для стран за секунды на GPU, для планеты — за минуты на GPU.
• Генерировать большие ансамбли экономично для вероятностного прогнозирования высокого разрешения.
• Интерактивно исследовать данные на километровой шкале.

CorrDiff: обучение и вывод

Рисунок 1 показывает рабочий процесс обучения и сэмплирования CorrDiff для генеративной детализации. Во время обучения диффузии предобученная регрессионная модель генерирует условное среднее значение, которое используется как вход для обучения модели диффузии. Для фона и деталей о CorrDiff смотрите публикацию CorrDiff, документацию PhysicsNeMo и исходный код.

Зачем оптимизировать CorrDiff?

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

• Упрощение ключевых операций (например, слияние ядер, использование смешанной точности, применение NVIDIA CUDA graphs и т.д.).
• Улучшение процесса сэмплирования за счет снижения числа шагов denoising и использования оптимальных схем интеграции по времени.

CorrDiff использует архитектуру EDM, где несколько вычислительно затратных операций, таких как групповую нормализацию, функции активации и свертки, можно оптимизировать с помощью продвинутых пакетов и слияния ядер.

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

Ускоренный CorrDiff

На следующих рисунках суммированы различные оптимизации, приводящие к ускорению более чем в 50 раз как для затрат на обучение, так и для вывода на домене CONUS. Рисунки 2 и 3 показывают кумулятивные факторы ускорения по сравнению с базой для каждой последовательной оптимизации. Детали каждой оптимизации приведены в следующих разделах.

Оптимизированный CorrDiff: как это достигается

Базовая производительность CorrDiff на GPU NVIDIA H100 с точностью FP32, размер батча = 1, размер патча = 1 (в абсолютном времени) составляла следующее:

• Прямой проход регрессии: 1204 мс
  • Домен: CONUS размером 1056 × 1792 пикселей
  • Входные каналы: [“u500”, “v500”, “z500”, “t500”, “u850”, “v850”, “z850”, “t850”, “u10m”, “v10m”, “t2m”, “tcwv”] с разрешением 25 км
  • Выходные каналы: [“refc”, “2t”, “10u”, “10v”] с разрешением 2 км
• Прямой проход диффузии: 155 мс
  • Домен: пространственный патч размером 448 x 448 пикселей
  • Входные каналы: [“u500”, “v500”, “z500”, “t500”, “u850”, “v850”, “z850”, “t850”, “u10m”, “v10m”, “t2m”, “tcwv”] с разрешением 25 км
  • Выходные каналы: [“refc”, “2t”, “10u”, “10v”] с разрешением 2 км
• Обратный проход диффузии: 219 мс

Хотя базовая версия работала эффективно, она страдала от дорогих прямых проходов регрессионной модели и неэффективных транспонирований данных.

Ключевые оптимизации обучения CorrDiff

Чтобы добиться значительного ускорения обучения CorrDiff, достигшего 53.86x на NVIDIA B200 и 25.51x на H100, внедрили серию оптимизаций производительности, описанных ниже.

Оптимизация 1: Включение AMP-BF16 для обучения
Исходный рецепт обучения использовал точность FP32. Здесь включили автоматическую смешанную точность (AMP) с BF16 для снижения потребления памяти и повышения пропускной способности без ущерба для численной стабильности, что привело к ускорению 2.03x по сравнению с базой.

Оптимизация 2: Распределение затрат регрессии с помощью многоитерационного патчинга
В исходном рабочем процессе обучения на основе патчей каждый сэмпл патча 448×448 для обучения модели диффузии требовал вывода регрессионной модели для полного домена CONUS 1056×1792. Это создавало bottleneck для пропускной способности обучения модели диффузии из-за вывода регрессии.

Эффективность повысили за счет кэширования выходов регрессии и повторного использования их для 16 патчей диффузии на временную метку. Это обеспечило более широкое пространственное покрытие, лучше распределяя затраты на регрессию, и дало ускорение 12.33x по сравнению с базой.

Оптимизация 3: Устранение транспонирований данных с помощью Apex GroupNorm
Конвейер обучения изначально использовал макет памяти NCHW, что вызывало дорогие неявные транспонирования перед и после сверток. Переход модели и входных тензоров к формату NHWC (channels-last) выравнивает их с предпочтительным макетом cuDNN. Однако операции GroupNorm в PyTorch не поддерживают формат channels-last. Чтобы избежать транспонирований и сохранить данные в формате channels-last для более эффективных ядер нормализации, заменили PyTorch GroupNorm на NVIDIA Apex GroupNorm. Это устранило overhead транспонирований и дало ускорение 16.71x по сравнению с базой.

Оптимизация 4: Слияние GroupNorm с SiLU
Слияние GroupNorm и активации SiLU в одно ядро с помощью Apex снизило число запусков ядер и доступов к глобальной памяти. Это повысило утилизацию GPU и обеспечило ускорение 17.15x по сравнению с базой.

Оптимизация 5: Расширение поддержки размерностей каналов в Apex GroupNorm
Некоторые слои CorrDiff используют размерности каналов, не поддерживаемые Apex. Расширили поддержку для этих размерностей, разблокировав слияние для всех слоев. Это улучшило производительность до ускорения 19.74x по сравнению с базой.

Оптимизация 6: Слияние ядер через компиляцию графа
Применение torch.compile позволило слить оставшиеся элементарные операции (например, сложение, умножение, sqrt, exp). Это улучшило планирование, снизило доступы к глобальной памяти и уменьшило overhead Python, дав ускорение 25.51x по сравнению с базой.

Оптимизация 7: Apex GroupNorm V2 на NVIDIA Blackwell
Использование Apex GroupNorm V2, оптимизированного для GPU NVIDIA Blackwell, дало ускорение 53.86x по сравнению с базой на B200 и 2.1x по сравнению с оптимизированным рабочим процессом на H100.

Пропускная способность обучения

Сравнение пропускной способности обучения базового CorrDiff на NVIDIA Hopper с оптимизированными версиями на Hopper и Blackwell приведено в таблице 1. Оптимизированная реализация показывает улучшения эффективности на обеих архитектурах, с наибольшими приростами на Blackwell.

Примечание: Регрессия относится к прямому проходу регрессии. Диффузия — к прямому проходу диффузии. Общее время включает комбинированные затраты (прямой проход регрессии + прямой проход диффузии + обратный проход диффузии).

Анализ Speed-of-Light

Чтобы оценить, насколько оптимизированный рабочий процесс CorrDiff приближается к потолку производительности оборудования, провели анализ Speed-of-Light (SOL) на GPU H100 и B200. Это дает оценку верхней границы достижимой производительности, оценивая, насколько эффективно используются ресурсы GPU.

Шаги для оценки SOL:

1. Выявление ядер с низкой утилизацией:
   Фокус на ядрах с утилизацией пропускной способности DRAM чтение/запись менее 60% и утилизацией Tensor Core менее 60%. Такие ядра не ограничены ни памятью, ни вычислениями, что делает их вероятными bottlenecks производительности.
2. Оценка потенциала на ядро:
   Для каждого ядра с низкой утилизацией оценивают потенциальное ускорение в идеальных условиях — полном использовании пропускной способности DRAM или полной активности Tensor Core.
3. Агрегация общего ускорения:
   Затем вычисляют гипотетическое end-to-end ускорение, если каждое ядро оптимизировано до идеальной производительности.
4. Вычисление эффективности SOL:
   Наконец, оценивают долю теоретического максимума SOL как долю пиковой производительности, достижимой, если топ-10 доминирующих по времени ядер individually доведены до теоретического максимума.

С помощью этой схемы оптимизированный рабочий процесс CorrDiff достигает 63% от оцененного SOL на H100 и 67% на B200. Это говорит о сильной утилизации GPU при оставшемся запасе для будущих улучшений на уровне ядер.

Для дальнейшей оценки эффективности визуализируют производительность ядер, как показано на рисунках 6 и 7. Каждая точка представляет ядро, размещенное по утилизации NVIDIA Tensor Core (ось x) и комбинированной утилизации пропускной способности DRAM чтение/запись (ось y). Размер точки отражает долю в общем времени выполнения, выделяя критичные для производительности операции.

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

Рисунок 6 показывает распределение ядер по утилизации Tensor Core и пропускной способности DRAM для базовой реализации CorrDiff. В неоптимизированном рабочем процессе с точностью FP32 более 95% времени тратится на ядра с низкой утилизацией, где как утилизация DRAM (чтение + запись), так и Tensor Core ниже 60%.

Большинство доминирующих по времени ядер группируются недалеко от начала координат, показывая очень низкую утилизацию DRAM и Tensor Core. Только небольшое число ядер лежит ближе к верхней или правой границе, где ядра явно ограничены памятью или вычислениями. Неоптимизированный рабочий процесс CorrDiff для США составлял всего 1.23% от SOL на B200.

Рисунок 7 показывает распределение ядер в оптимизированной реализации по утилизации Tensor Core и пропускной способности DRAM. В оптимизированном рабочем процессе с обучением AMP-BF16 большая доля ядер ближе к верхнему левому или нижнему правому краю, указывая на хорошую производительность и утилизацию GPU. Оптимизированный CorrDiff теперь достигает 67% от SOL на B200. Несмотря на общие улучшения, некоторые ядра все еще могут быть ускорены дальше.

Оптимизации вывода CorrDiff

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

Оптимизированный multi-diffusion
CorrDiff использует подход multi-diffusion на основе патчей, где пересекающиеся пространственные патчи denoising'ятся и агрегируются. Изначально 27.1% общего времени тратилось на операции im2col folding/unfolding. Предварительный расчет счетчиков пересечений для каждого пикселя и применение torch.compile() для ускорения оставшихся шагов folding/unfolding полностью устраняет bottleneck im2col, приводя к ускорению 7.86x.

Детерминистический сэмплер Euler (12 шагов)
Исходный стохастический сэмплер использовал 18 шагов denoising с солвером Heun и коррекцией второго порядка. Включение детерминистического сэмплера на основе солвера Euler (без коррекции второго порядка) снизило число шагов denoising до 12 без влияния на качество выхода. Это изменение дало дополнительное ускорение ~2.8x на Hopper и Blackwell. Итоговое ускорение с 12-шаговым детерминистическим сэмплером — 21.94x на H100 и 54.87x на B200.

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

Начало работы

Обучение/вывод CorrDiff в PhysicsNeMo: Документация CorrDiff в PhysicsNeMo

• Для обучения с оптимизированным кодом следуйте инструкциям в readme репозитория CorrDiff и установите следующие опции в разделе training.perf в выбранной YAML-конфигурации обучения:

fp_optimizations: amp-bf16
use_apex_gn: True
torch_compile: True
profile_mode: False

• Для вывода с оптимизированным кодом следуйте инструкциям в readme репозитория CorrDiff и установите следующие опции в разделе generation.perf в выбранной конфигурации генерации:

use_fp16: True
use_apex_gn: True
use_torch_compile: True
profile_mode: False
io_syncronous: True

• Установите profile_mode в False для оптимизированной производительности, поскольку аннотации NVTX вводят разрывы графа в рабочем процессе torch.compile.
• Для использования последних ядер Apex GroupNorm либо соберите Apex GroupNorm в Dockerfile контейнера PhysicsNeMo, либо соберите локально после загрузки контейнера PhysicsNeMo
  • Клонируйте репозиторий Apex и соберите с помощью:

CFLAGS="-g0" NVCC_APPEND_FLAGS="--threads 8" \
pip install \
--no-build-isolation \
--no-cache-dir \
--disable-pip-version-check \
--config-settings "--build-option=--group_norm" .

Узнайте больше об оптимизированном обучении CorrDiff в PhysicsNeMo и запускайте оптимизированные рабочие процессы в Earth2Studio.