NVIDIA и Hugging Face объединили усилия: NeMo Automodel получает поддержку диффузионных моделей

# NVIDIA и Hugging Face объединили усилия: NeMo Automodel получает поддержку диффузионных моделей

Еще в июле NVIDIA и Hugging Face выпустили совместный пост, анонсировавший интеграцию опенсорсной библиотеки для обучения нейросетей NeMo Automodel с диффузионными моделями. В официальном сообщении упоминались передовые технологии, такие как FSDP2, тензорный параллелизм, контекстный параллелизм и мультинодовая оркестрация. Обещали, что дообучение моделей FLUX, Wan и HunyuanVideo можно будет масштабировать с одной видеокарты до сотен, не требуя конвертации точек проверки (чек-пойнтов).

Несмотря на впечатляющий список технологий, стоит отметить, что сами алгоритмы распределенного обучения разработаны давно и существовали в PyTorch и Megatron-Core. Тем не менее, их объединение в единую экосистему с открытым исходным кодом решает давнюю проблему инженеров: необходимость ручного управления «конструктором» при работе с GPU.

Проблема нестабильности в экосистеме Diffusers и DeepSpeed

До появления обновлений в NeMo Automodel процесс дообучения диффузионной модели на нескольких графических процессорах был сопряжен с рядом трудностей. Стандартная практика предполагала использование связки Diffusers, Accelerate и DeepSpeed ZeRO или FSDP1. Однако эти инструменты, изначально спроектированные как универсальная обвязка, часто конфликтовали при работе с спецификой диффузионных пайплайнов.

Инженеры столкнулись со следующими проблемами:

* Неэффективное использование памяти: При дообучении модели FLUX ControlNet с использованием DeepSpeed Stage-3 на 8 GPU объем памяти на каждой карте оставался практически таким же, как при работе на одном устройстве (примерно 42 ГБ), вместо ожидаемого шардирования весов. * Снижение производительности: Попытки решить проблему памяти переходом на Stage-3 приводили к резкому падению скорости обучения (до 16 секунд на итерацию), что делало процесс практически нецелесообразным. * Частые сбои: Совмещение чек-пойнтинга градиентов и DeepSpeed ZeRO3 часто вызывало критические ошибки и падения обучения на разных версиях Diffusers. * Замедление распределенного обучения: В ряде случаев обучение через мультикарточное запускалось медленнее, чем на одном GPU.

Суть проблемы заключалась в том, что инструменты шардирования (расщепления параметров модели) для диффузионных архитектур, содержащих сложные компоненты вроде VAE-энкодеров и декодеров, и нестандартные механизмы внимания (как в FLUX, Wan или HunyuanVideo), не всегда работали корректно «из коробки». Инженерам приходилось выбирать между ограничением по размеру модели (DDP) и ненадежным шардированием.

Как NeMo Automodel решает вопросы масштабирования

Разработчики, включая соавтора поста и мейнтейнера библиотеки Diffusers — Сайяка Пола — создали обновление, которое устраняет необходимость в ручном сборке сложной инфраструктуры. Библиотека NeMo Automodel была расширена для поддержки диффузионных моделей, что затрагивает код на стороне Hugging Face, а не только добавляет обвязку поверх существующих библиотек.

Ключевые особенности новой интеграции:

* Прямая работа с объектами Hugging Face: Библиотека взаимодействует непосредственно с классами модели и пайплайнов из Diffusers (например, WanTransformer3DModel). Это позволяет сохранять оригинальный чек-пойнт после обучения и сразу загружать его обратно в DiffusionPipeline для генерации без конвертации. * Унифицированная конфигурация: Различные виды параллелизма (FSDP2, тензорный, контекстный, HSDP, SP) теперь определяются в едином YAML-конфиге. Пользователю не нужно переписывать код для изменения масштаба с одной карты до кластера — достаточно изменить настройки сетки устройств (device mesh). * Фундамент на PyTorch DTensor: В основе лежит принцип SPMD (Single Program Multiple Data), где скрипт запускается на одной GPU или на всем кластере одинаково. Это позволяет легко комбинировать различные методы параллелизма. * Поддержка Flow Matching: Пока что инструмент ориентирован на модели с согласованием потока (flow matching), обучение проходит в латентном пространстве на основе заранее закэшированных выходов VAE, а также поддерживается группировка данных разного разрешения в один батч.

Технические детали и ограничения железа

Важно различать маркетинговые заявления и технические реалии. В основе параллелизма лежат NVIDIA-специфичные ускорители, такие как Transformer Engine, DeepEP и FlexAttention.

FSDP2 против FSDP1 Одной из ключевых технологий является FSDP2. В отличие от предыдущей версии, которая хранила параметры модели как один плоский тензор (FlatParameter), FSDP2 использует представление DTensor. Это позволяет тензорам быть размеченными по устройствам в кластере, что упрощает комбинирование с тензорным и контекстным параллелизмом. Кроме того, FSDP2 позволяет независимо настраивать точность вычислений для параметров, reductions и выходов.

Требования к видеопамяти и модели Несмотря на заявления о масштабируемости, практические требования к железу остаются высокими. Официальные бенчмарки сняты на связке из восьми видеокарт NVIDIA H100 (80 ГБ).

* Модель HunyuanVideo 1.3B на разрешении 720p требует от 60 до 80 ГБ видеопамяти. Даже на карте на 80 ГБ существует риск переполнения (OOM) из-за скачков памяти при инференсе. * Модель FLUX в облегченной FP8-версии потребляет 40–50 ГБ памяти.

Также стоит обратить внимание на путаницу в документации относительно параметров модели HunyuanVideo 1.5. В таблице указано 13B параметров, однако карточка модели на Hub подписана как 8.3B. Вероятно, в таблице использовались данные старой версии, в то время как актуальная версия от Tencent является более компактной.

Лицензирование и масштаб С точки зрения лицензирования, все материалы открыты под лицензией Apache 2.0, что разрешает свободное коммерческое использование. Однако заявка на масштаб «до сотен» карт пока основана на словах.

Публикованные цифры производительности получены на одном узле из 8 карт H100. Ни одного мультинодового бенчмарка NVIDIA в блоге не представила. Более того, в открытом бэклоге задач на GitHub указано, что проверка мультинодовой производительности именно для диффузионного обучения на моделях типа HunyuanVideo является задачей на будущее. Это означает, что механизм работы уже существует и задокументирован, но его реальная эффективность на больших масштабах еще не измерялась.

Заключение

Инструмент NeMo Automodel предлагает значимое облегчение для инженеров, которые ранее страдали от необходимости вручную конвертировать веса между форматами Hugging Face и Megatron-Core. Возможность начинать эксперименты на AutoModel и сразу же передавать готовые модели в экосистему Megatron через официальный мост устраняет этот барьер.

Для команд, уже имеющих стабильную связку на Diffusers с Accelerate или DeepSpeed, переход может потребовать усилий, если текущая конфигурация работает без сбоев. Однако для старта новых проектов с нуля или для случаев, где регулярно возникают проблемы с масштабом и конвертацией весов, новая библиотека является мощным решением. Полное дообучение моделей такого уровня, как FLUX.1-dev, все равно требует значительных ресурсов (минимум 8 карт H100), но инфраструктурные боли вокруг шардирования и параллелизма теперь решены на уровне кода.

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

Мультинодное обучение и поддержка актуальных архитектур Blackwell (MXFP8) остаются в планах, но текущее обновление уже готово для развертывания на доступном железе Hopper и выше.