NVIDIA выпустила NeMo Automodel для дообучения диффузионных моделей Hugging Face Diffusers на сотнях GPU

NVIDIA и Hugging Face представили интеграцию открытой библиотеки NeMo Automodel (часть фреймворка NVIDIA NeMo) с библиотекой Diffusers. Цель, дать любой диффузионной модели в формате Diffusers на Hugging Face Hub промышленное распределённое обучение без конвертации чекпоинтов и без переписывания кода модели под новый инструмент. Интеграция полностью открыта по лицензии Apache 2.0 и описана в официальном руководстве по обучению Diffusers.

NeMo Automodel, это PyTorch-библиотека, построенная на DTensor, с двумя ключевыми принципами. Во-первых, она нативно работает с Hugging Face: достаточно указать путь к любой модели Diffusers на Hub, и обучение стартует, библиотека использует классы моделей Diffusers (например, WanTransformer3DModel) и пайплайны Diffusers (WanPipeline) для генерации, а чекпоинты свободно возвращаются в экосистему Diffusers. Во-вторых, один и тот же код работает на любом масштабе: переключение между схемами параллелизма, FSDP2, тензорный, экспертный, контекстный и конвейерный (pipeline) параллелизм, задаётся конфигурацией, а не переписыванием модели. Сейчас NeMo Automodel поддерживает только модели, обучаемые методом flow matching (согласования потоков), с обучением в латентном пространстве через предварительно закодированные VAE-представления и мультиразрешающей группировкой данных (multiresolution bucketing) для ускорения.

Готовые рецепты дообучения из коробки поставляются для открытых диффузионных моделей, включая FLUX.1-dev (12 млрд параметров), Wan 2.1, HunyuanVideo (13 млрд параметров) и FLUX.2.

Что конкретно даёт интеграция пользователям Diffusers: предобученные веса с Hub работают сразу, без отдельного «формата для обучения», дообученный чекпоинт напрямую загружается в DiffusionPipeline для инференса или публикуется обратно на Hub, а инструменты квантования, компиляции, LoRA-адаптеров и кастомных сэмплеров продолжают работать. Поддержка новой модели Diffusers добавляется в NeMo Automodel небольшим точечным изменением (обработчик предобработки данных и адаптер модели), а не полным новым скриптом обучения, остальной стек (FSDP2, группировка данных, чекпоинты, генерация) переиспользуется без изменений. Поддерживаются как полное дообучение (full fine-tuning), так и параметрически эффективное дообучение в стиле LoRA, можно выбирать между максимальным качеством на большом кластере и максимальной эффективностью на одном узле. Именно масштабируемость (шардинг FSDP2, тензорный/контекстный/конвейерный параллелизм, оркестрация нескольких узлов через SLURM, в перспективе, Kubernetes, и мультиразрешающая группировка) делает возможным обучение моделей такого размера, как FLUX.1-dev и HunyuanVideo.

В демонстрационном сценарии команда дообучила FLUX.1-dev с полным (full-transformer) дообучением на датасете из 78 карт Таро в стиле Rider-Waite, используя один узел с 8 GPU NVIDIA H100 80GB. Рекомендуемый способ установки, Docker-контейнер NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06) с предустановленными PyTorch и TransformerEngine; также доступна установка через pip или сборка из исходников. Обучение шло 200 шагов с чекпоинтами на 50, 100, 150 и 200 шаге; подписи к изображениям датасета содержали триггер-токен trtcrd. После обучения при генерации с этим токеном модель сохраняла запрошенный сюжет (например, «космонавт ухаживает за розовым садом на Марсе»), но приобретала кремово-красно-чёрную винтажную палитру, плотные чернильные контуры, плоские цветовые заливки и композицию в духе гадальной карты; без токена та же сцена оставалась фотореалистичной, то есть выученный эффект привязан именно к триггеру, а не подменяет базовую модель целиком. Отдельно показан пример дообучения Wan 2.1 на видео в стиле Ghibli (заметное изменение внешнего вида цветка по сравнению с базовой моделью) и применения LoRA-адаптера к Wan 2.1, которое придаёт видео характерный стиль Ghibli, особенно заметный в прорисовке глаз персонажей.

Производительность измерялась на одном узле с 8 GPU NVIDIA H100 80GB (полностью соединены через NVLink) отдельно для задач «текст-в-изображение» (512×512) и «текст-в-видео» (512×512, 49 кадров на клип), время шага и пиковая память приведены как среднее по трём установившимся окнам из 10 шагов.

В ближайшем релизе NeMo Automodel команда планирует добавить полностью типизированный Pythonic API поверх тех же рецептов диффузии, наряду с текущим YAML-конфигом, чтобы их было проще встраивать в существующий обучающий код, ноутбуки и рабочие процессы экспериментов.

Ключевые факты

NeMo Automodel, открытая (Apache 2.0) PyTorch-библиотека NVIDIA для обучения, интегрированная с форматом Diffusers без конвертации чекпоинтов
Поддерживает только модели с обучением методом flow matching: готовые рецепты дообучения для FLUX.1-dev (12 млрд параметров), Wan 2.1, HunyuanVideo (13 млрд параметров) и FLUX.2
Масштабирование задаётся конфигурацией, а не переписыванием кода: FSDP2, тензорный, экспертный, контекстный и конвейерный параллелизм; поддержаны и полное дообучение, и LoRA
В демонстрации FLUX.1-dev дообучили на 78 изображениях карт Таро (Rider-Waite) на одном узле с 8 GPU NVIDIA H100 80GB, используя триггер-токен trtcrd
В разработке, Pythonic API поверх текущих YAML-рецептов диффузионного обучения

Почему это важно

Диффузионные модели (FLUX.1-dev, Wan 2.1, HunyuanVideo и другие) стали основой открытых релизов для генерации изображений и видео, а библиотека Diffusers, стандартным интерфейсом для инференса и адаптации таких моделей. Но их обучение и дообучение требуют памяти-эффективного шардинга, кэширования латентных представлений, группировки по разрешениям и масштабирования от одной видеокарты до сотен, раньше эти возможности нужно было собирать самостоятельно или писать под каждую модель заново. Интеграция NeMo Automodel даёт готовый промышленный конвейер распределённого обучения для любой модели в формате Diffusers на Hub без конвертации чекпоинтов и без переписывания кода модели.

Кому это важно

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

Как это применить

Рекомендуемый способ, Docker-контейнер NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06) с предустановленными PyTorch и TransformerEngine; альтернативы, установка через pip (pip3 install nemo-automodel) или сборка из исходников с GitHub. Дообучение запускается по готовым YAML-конфигам из репозитория (например, для FLUX.1-dev) с переопределением нужных параметров через командную строку, датасет, разрешение, число шагов, путь к чекпоинтам; поддержаны и полное дообучение, и LoRA. Часть процесса, предкодирование датасета (VAE-латенты и текстовые эмбеддинги) отдельным шагом перед запуском самого обучения.

Можно ли доверять

Материал, официальный технический пост в блоге Hugging Face, написанный совместно с NVIDIA, с конкретными командами, конфигурационными файлами, номером версии Docker-образа и измеренными на конкретном железе (8× NVIDIA H100 80GB) показателями производительности. Библиотека полностью открыта по лицензии Apache 2.0, а описанный пример дообучения (карты Таро, триггер-токен trtcrd) воспроизводим по приведённым командам, это техническая документация к реальному релизу, а не маркетинговый анонс без деталей.

Риски и подводные камни

Сейчас NeMo Automodel поддерживает только диффузионные модели, обучаемые методом flow matching, не любую архитектуру. Оркестрация нескольких узлов пока работает только через SLURM, поддержка Kubernetes заявлена как будущая. Демонстрационное дообучение крупной модели (FLUX.1-dev) потребовало узла с 8 GPU NVIDIA H100 80GB, ресурс, доступный не всем командам. Pythonic API взамен текущего YAML-конфига пока не выпущен, а сам инструмент и интеграция, молодые, специфика конкретных моделей и стабильность API могут меняться.