Зачем NVIDIA скрестила NeMo Automodel с диффузионными моделями

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

Diffusers и DeepSpeed часто ломались на ровном месте

До появления NeMo Automodel дообучение диффузионной модели на нескольких GPU означало ручную сборку конструктора. Diffusers давал модель и пайплайн, поверх ставили Accelerate, а под капотом у Accelerate уже работал DeepSpeed ZeRO или FSDP1, отвечающие за шардинг. И этот конструктор регулярно ломался, причем в самых обычных сценариях, которые легко найти в трекере задач Diffusers на GitHub. При дообучении FLUX ControlNet с DeepSpeed Stage-3 на 8 GPU память на каждой карте оставалась такой же, как на одной GPU, примерно 42 GB что на одной, что на восьми, хотя весь смысл ZeRO Stage-3 в том, чтобы шардировать веса между картами и снижать нагрузку на каждую. При дообучении FLUX.1-dev через dreambooth-скрипт с Accelerate и DeepSpeed на нескольких GPU обучение стабильно падало, и автор тикета не смог самостоятельно разобраться с ошибкой.

При обучении FLUX на 4×A100 80GB с DeepSpeed Stage 2 возникал OOM. Единственным обходным путем оказался переход на Stage 3. Но скорость просела примерно до 16 секунд на итерацию, и шардирование «решало» проблему с памятью ценой скорости, которая делала обучение практически нецелесообразным. По отдельности чек-пойнтинг градиентов и DeepSpeed ZeRO3 нормально работали, но вместе, начиная с определенной версии Diffusers, они стабильно приводили к падению обучения и вызывали критические ошибки. Получается, инженерам приходилось выбирать не между хорошим и плохим вариантом, а между разными формами боли: либо мириться с тем, что шардирование не дает обещанной экономии памяти, либо терять скорость ради стабильности, либо ловить случайные падения на стыке фич.

Проблема была не только в стабильности, но и в самом смысле масштабирования. У нескольких пользователей мультикарточное обучение через accelerate launch --multi_gpu оказывалось заметно медленнее, чем на одной карте. То есть распределенное обучение только ухудшало процесс.

Похоже, дело не в том, что инженеры Diffusers и Accelerate делают что-то неправильно. Скорее эти инструменты изначально проектировались как универсальная обвязка поверх любых PyTorch-моделей. А DeepSpeed ZeRO как техника разрабатывался прежде всего под LLM, где архитектура трансформера более однородная. У диффузионных пайплайнов, судя по всему, есть своя специфика, которую такая обвязка плохо учитывает «из коробки»: отдельные VAE-энкодер и -декодер со своими требованиями к памяти, нестандартные архитектуры внимания (FLUX, Wan и HunyuanVideo устроены по-разному), необходимость держать в памяти шум и промежуточные состояния диффузионного процесса.

В результате инженерам приходилось выбирать между двумя так себе вариантами. Первый — DDP. Он предсказуем, но не шардирует веса и ограничен по размеру модели, которую можно обучить на имеющихся картах. Второй — DeepSpeed или FSDP1. Теоретически они снимают потолок по памяти, но с заметной вероятностью либо не дают обещанной экономии, либо падают с ошибками. И чтобы понять, в чем дело, приходится самостоятельно разбираться в исходниках интеграции. Именно поэтому объем видеопамяти конкретной карты кажется главным лимитирующим фактором. Если шардирование ненадежно, то есть только два рабочих пути обучить модель больше той, что помещается на одну карту: либо вручную обходить проблемы DeepSpeed одну за другой, либо просто покупать карту побольше.

NeMo Automodel научился работать с диффузионными моделями через Diffusers

Соавтором поста стал Сайяк Пол, один из мейнтейнеров библиотеки Diffusers, и это неспроста. Если бы интеграцию делали только со стороны NVIDIA, то к ней можно было бы отнестись как к внешней надстройке. А участие мейнтейнера говорит о том, что изменения затрагивают и код на стороне Hugging Face, а не только обвязку NVIDIA поверх чужой библиотеки.

Формально ребята взяли уже существующую библиотеку NeMo Automodel и расширили ее на диффузионные модели. До этого она закрывала только языковые и мультимодальные модели вроде Llama, Qwen, DeepSeek и других. Проблема была в том, что из-за роста интереса к дообучению диффузионных моделей стало не хватать эффективного по памяти шардинга, кеширования латентов, группировки данных разного разрешения в один батч и конфигураций, которые плавно масштабируются с одной GPU до сотен.

Устанавливается все как обычная опенсорсная библиотека. Можно взять готовый Docker-контейнер nvcr.io/nvidia/nemo-automodel:26.06, там уже собраны PyTorch и Transformer Engine. Можно просто написать pip3 install nemo-automodel. Никакого закрытого SDK и доступа по запросу тут нет. Весь код и все конфигурации открыты под лицензией Apache 2.0, включая диффузионный модуль.

Работа напрямую с объектами Hugging Face

В конфиге указываете ID модели из Diffusers на Hub, и NeMo Automodel сам подхватывает нужные классы модели и пайплайна, например WanTransformer3DModel и WanPipeline. Библиотека не создает отдельный слой поверх Diffusers, а работает прямо с ее объектами. Поэтому чек-пойнт после обучения не нужно конвертировать, он сразу загружается обратно в DiffusionPipeline для генерации.

Один конфиг — любой масштаб

А параллелизм здесь ставится настройкой. Хотите FSDP2, хотите тензорный, экспертный или контекстный параллелизм — просто переключаете в конфиге. Модель трогать не нужно, масштаб сам меняется.

AutoModel пока ограничивается поддержкой только моделей с согласованием потока (flow matching). Обучение проходит в латентном пространстве, по заранее закешированным выходам VAE. А еще есть группировка по разрешению, которая нужна, чтобы в один батч попадали разных размеров данные. Так что это инструмент под конкретное семейство моделей, и туда как раз входят FLUX, Wan и HunyuanVideo.

FSDP2 вместо FlatParameter, и своя цена в VRAM

Все виды параллелизма собраны в единый YAML-конфиг вместо отдельных кастомных скриптов под каждый случай. В основе лежит PyTorch DTensor и принцип SPMD (single program multiple data). Это значит, что один и тот же скрипт запускается и на одной GPU, и на целом кластере, а масштаб задается только сменой device mesh в конфиге. Переписывать код под каждый новый масштаб не нужно.

Кроме FSDP2, тензорного и контекстного параллелизма, в документации есть еще два вида: HSDP, hybrid sharded data parallel, и SP, sequence parallel. А ускоряют все это NVIDIA-специфичные кернели: Transformer Engine, DeepEP, FlexAttention.

Здесь имеет смысл остановиться на разнице между FSDP1 и FSDP2, потому что это не ребрендинг. FSDP1 хранил параметры модели как один большой уплощенный тензор на GPU, который называется FlatParameter. Из-за плоской структуры FSDP1 плохо сочетался с другими видами параллелизма, их сложно было совместить в одном обучении. FSDP2 устроен иначе: параметры хранятся как DTensor, то есть тензор, размеченный по устройствам в кластере. Такое представление легко комбинируется с тензорным, контекстным и пайплайн-параллелизмом. Плюс у FSDP2 можно отдельно настраивать точность вычислений на разных этапах: param_dtype, reduce_dtype и output_dtype задаются независимо друг от друга. Похоже, именно эта техническая база и позволяет NeMo Automodel предлагать принцип «один конфиг — любой масштаб» вместо набора фич, конфликтующих между собой.

Это удобство имеет свою цену. Официальные бенчмарки NVIDIA сняты на связке из восьми видеокарт H100 80GB, и это фактический минимум для моделей такого масштаба. Оригинальный HunyuanVideo, 13B параметров, на разрешении 720p требует 60–80 GB видеопамяти, и даже на карте 80GB из-за скачков при инференсе есть риск OOM. FLUX даже в облегченной FP8-версии на дообучении съедает 40–50 GB.

А еще в анонсе NVIDIA есть путаница, которая может дорого стоить тому, кто решит планировать железо по этим цифрам. В таблице поддерживаемых моделей HunyuanVideo 1.5 указан как модель на 13B параметров. Но карточка этой же модели, hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v, в том же посте подписана уже как 8B. Разница не случайная: настоящая HunyuanVideo 1.5 от Tencent — компактная модель на 8.3B параметров, которая в потребительских картах работает и на 14 GB, заметно легче тяжелого оригинального HunyuanVideo на 13B. Похоже, в таблицу с бенчмарками просто перекочевала цифра от старой версии, а карточку модели уже успели поправить. Так что если планировать VRAM под конкретную модель по таблице конфигураций NeMo Automodel, я бы на месте читателя проверила параметры модели отдельно.

С лицензией все чисто, а вот масштаб пока на словах

Прежде чем говорить, кому подходит NeMo Automodel, я бы отделила факты от маркетинга.

С лицензией все чисто. И NeMo Automodel, и весь диффузионный модуль внутри него лицензированы под Apache License 2.0. Это открытый код без ограничений на коммерческое использование, и здесь вопросов нет.

А вот с «любым масштабом» есть вопросики. Код действительно не привязан к одной конкретной карте, база на PyTorch DTensor работает везде, где есть CUDA. Но производительность — это уже совсем другой разговор. Ускоренные кернели вроде Transformer Engine и DeepEP собираются под конкретные поколения железа, и часть фич попросту недоступна на другом железе. FP8 через Transformer Engine работает начиная с Hopper, то есть с H100 и H200. А MXFP8, микроскейлинг FP8 с блочным масштабированием вместо потензорного у Hopper, вообще эксклюзив Blackwell, то есть B200 и GB200. При этом все опубликованные в блоге цифры производительности сняты на одном узле из 8×H100. Ни одного мультинодового бенчмарка NVIDIA не привела, хотя заголовок анонса обещает масштаб «до сотен» карт.

Согласно аналитике SemiAnalysis, крупномасштабные продакшен-тренировки на GB200 NVL72 на момент публикации все еще редкость, экосистема софта под Blackwell продолжает дозревать, и на масштабе есть нюансы с надежностью. Даже сама NVIDIA для демопримера в блоге, дообучения FLUX.1-dev на датасете карт Таро, использовала 8×H100, а не GB200, причем с явно выключенным флагом transformer_engine_fp8 в команде запуска. То есть даже для собственной демонстрации NVIDIA не форсировала точность до FP8.

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

Где фокус?

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

NeMo Automodel убирает эту конвертацию как обязательный шаг. Можно начать с простого AutoModel-пути на HF-моделях, а когда экспериментов станет больше, перейти на Megatron-Core через Megatron Bridge. Мост берет готовый HF-чек-пойнт из AutoModel и конвертирует его в Megatron-формат через AutoBridge.from_hf_pretrained(). Никакой ручной пересборки весов. Путь официально задокументирован и описан как как интуитивно понятный и простой. Дообучил модель или прогнал быстрый эксперимент на AutoModel, указал Megatron Bridge на получившийся HF-чек-пойнт, и дальше можно работать с ним в экосистеме Megatron.

Кому переезжать, а кому сидеть на месте

NeMo Automodel одинаково поддерживает LoRA-файнтюн на одной карте и полное дообучение на кластере в рамках одной логики конфигов. Это подтверждено официально, эталонный пример NVIDIA с FLUX.1-dev построен на полном дообучении. Но порог входа все равно высокий. Полное дообучение даже для FLUX.1-dev на 12B параметров требует восьми H100 80GB — это минимум, который показала в демо NVIDIA. Так что выбор между LoRA и полным дообучением на практике зависит от того, сколько карт у вас есть.

Если у команды уже есть стабильная DIY-связка на Diffusers с Accelerate или DeepSpeed, которая не падает и устраивает по скорости, менять ее на новый инструмент ради самого факта новизны смысла нет. Но если проект начинается с нуля, или регулярно упирается в масштаб, или просто надоело тратить время на ручную конвертацию чек-пойнтов между HF и Megatron-Core, вот тут NeMo Automodel снимает боль. А мультинод и MXFP8 на GB200 для диффузии пока только в планах NVIDIA, реальных тестов на нескольких узлах компания еще не показала.