NeMo Automodel

NeMo Automodel — это библиотека обучения от NVIDIA, нативно поддерживающая PyTorch DTensor, предназначенная для дообучения и предобучения диффузионных моделей в промышленных масштабах. Библиотека нативно интегрирована с Hugging Face: можно обучать любую модель в формате Diffusers из Hub без конвертации контрольных точек. Один и тот же YAML-рецепт и модифицируемый скрипт обучения работают на любой конфигурации — от 1 GPU до сотен узлов — с распределённым обучением FSDP2, загрузкой данных с мультiresolution-бакетизацией и обучением в предварительно закодированном латентном пространстве для максимальной загрузки GPU. Для обучения используется flow matching. Проект полностью открыт (лицензия Apache 2.0), поддерживается NVIDIA и активно развивается.

NeMo Automodel напрямую интегрируется с Diffusers. Он загружает предобученные модели из Hugging Face Hub с использованием классов моделей Diffusers и генерирует выходные данные с помощью DiffusionPipeline.

Типичный рабочий процесс включает установку NeMo Automodel (через pip или Docker), подготовку данных путём их кодирования в файлы .meta, настройку рецепта в формате YAML, запуск обучения с помощью torchrun и выполнение инференса с использованием полученного чекпоинта.

Поддерживаемые модели

Установка

Требования к оборудованию

Установите NeMo Automodel с помощью pip. Полный список способов установки (в том числе из исходного кода) описан в руководстве по установке NeMo Automodel.

pip3 install nemo-automodel

В качестве альтернативы используйте готовый Docker-контейнер, содержащий все зависимости.

docker pull nvcr.io/nvidia/nemo-automodel:26.02.00
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.02.00

Точки контроля теряются при завершении работы контейнера, если каталог с контрольными точками не смонтирован на хост. Например, добавьте -v /host/path/checkpoints:/workspace/checkpoints к команде docker run.

Подготовка данных

NeMo Automodel обучает диффузионные модели в латентном пространстве. Исходные изображения или видео необходимо предварительно обработать в файлы .meta, содержащие латенты VAE и текстовые эмбеддинги, перед началом обучения. Это позволяет избежать повторного кодирования на каждом шаге обучения.

Используйте встроенный инструмент предобработки для кодирования данных. Инструмент автоматически распределяет работу по всем доступным GPU.

python -m tools.diffusion.preprocessing_multiprocess video \
    --video_dir /data/videos \
    --output_dir /cache \
    --processor wan \
    --resolution_preset 512p \
    --caption_format sidecar

python -m tools.diffusion.preprocessing_multiprocess video \
    --video_dir /data/videos \
    --output_dir /cache \
    --processor hunyuan \
    --target_frames 121 \
    --caption_format meta_json

python -m tools.diffusion.preprocessing_multiprocess image \
    --image_dir /data/images \
    --output_dir /cache \
    --processor flux \
    --resolution_preset 512p

Формат вывода

Предобработка создаёт каталог кэша, организованный по бакетам разрешения. NeMo Automodel поддерживает обучение с несколькими разрешениями с помощью выборки по бакетам. Образцы группируются по пространственному разрешению, чтобы каждая партия содержала образцы одного размера, что исключает потери на заполнение (padding).

/cache/
├── 512x512/                      # Resolution bucket
│   ├── <hash1>.meta              # VAE latents + text embeddings
│   ├── <hash2>.meta
│   └── ...
├── 832x480/                      # Another resolution bucket
│   └── ...
├── metadata.json                 # Global config (processor, model, total items)
└── metadata_shard_0000.json      # Per-sample metadata (paths, resolutions, captions)

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

Конфигурация обучения

Настройка модели опирается на два компонента:

1. Скрипт рецепта (finetune.py) — это точка входа на Python, содержащая цикл обучения: загрузка модели, создание загрузчика данных, выполнение прямых и обратных проходов, вычисление потерь flow matching, сохранение контрольных точек и ведение журнала.
2. Файл конфигурации YAML определяет все параметры, используемые рецептом: какую модель дообучать, где находятся данные, гиперпараметры оптимизатора, стратегию параллелизма и многое другое. Вы настраиваете обучение, редактируя этот файл, а не изменяя код, что позволяет масштабироваться от 1 до сотен GPU.

Любое поле YAML также можно переопределить из CLI:

torchrun --nproc-per-node=8 examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml \
  --optim.learning_rate 1e-5 \
  --step_scheduler.num_epochs 50

Ниже приведена аннотированная конфигурация для дообучения Wan 2.1 T2V 1.3B, с пояснением каждого раздела.

seed: 42

# ── Experiment tracking (optional) ──────────────────────────────────────────
# Weights & Biases integration for logging metrics, losses, and learning rates.
# Set mode: "disabled" to turn off.
wandb:
  project: wan-t2v-flow-matching
  mode: online
  name: wan2_1_t2v_fm

# ── Model ───────────────────────────────────────────────────────────────────
# pretrained_model_name_or_path: any Hugging Face model ID or local path.
# mode: "finetune" loads pretrained weights; "pretrain" trains from scratch.
model:
  pretrained_model_name_or_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
  mode: finetune

# ── Training schedule ───────────────────────────────────────────────────────
# global_batch_size: effective batch across all GPUs.
# Gradient accumulation is computed automatically: global / (local × num_gpus).
step_scheduler:
  global_batch_size: 8
  local_batch_size: 1
  ckpt_every_steps: 1000     # Save a checkpoint every N steps
  num_epochs: 100
  log_every: 2               # Log metrics every N steps

# ── Data ────────────────────────────────────────────────────────────────────
# _target_: the dataloader factory function.
#   Use build_video_multiresolution_dataloader for video models (Wan, HunyuanVideo).
#   Use build_text_to_image_multiresolution_dataloader for image models (FLUX).
# model_type: "wan" or "hunyuan" (selects the correct latent format).
# base_resolution: target resolution for multiresolution bucketing.
data:
  dataloader:
    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
    cache_dir: PATH_TO_YOUR_DATA
    model_type: wan
    base_resolution: [512, 512]
    dynamic_batch_size: false  # When true, adjusts batch per bucket to maintain constant memory
    shuffle: true
    drop_last: false
    num_workers: 0

# ── Optimizer ───────────────────────────────────────────────────────────────
# learning_rate: 5e-6 is a good starting point for fine-tuning.
# Adjust weight_decay and betas for your dataset.
optim:
  learning_rate: 5e-6
  optimizer:
    weight_decay: 0.01
    betas: [0.9, 0.999]

# ── Learning rate scheduler ─────────────────────────────────────────────────
# Supports cosine, linear, and constant schedules.
lr_scheduler:
  lr_decay_style: cosine
  lr_warmup_steps: 0
  min_lr: 1e-6

# ── Flow matching ───────────────────────────────────────────────────────────
# adapter_type: model-specific adapter — must match the model:
#   "simple" for Wan 2.1, "flux" for FLUX.1-dev, "hunyuan" for HunyuanVideo.
# timestep_sampling: "uniform" for Wan, "logit_normal" for FLUX and HunyuanVideo.
# flow_shift: shifts the flow schedule (model-dependent).
# i2v_prob: probability of image-to-video conditioning during training (video models).
flow_matching:
  adapter_type: "simple"
  adapter_kwargs: {}
  timestep_sampling: "uniform"
  logit_mean: 0.0
  logit_std: 1.0
  flow_shift: 3.0
  num_train_timesteps: 1000
  i2v_prob: 0.3
  use_loss_weighting: true

# ── FSDP2 distributed training ──────────────────────────────────────────────
# dp_size: number of GPUs for data parallelism (typically = total GPUs on node).
# tp_size, cp_size, pp_size: tensor, context, and pipeline parallelism.
#   For most fine-tuning, dp_size is all you need; leave others at 1.
fsdp:
  tp_size: 1
  cp_size: 1
  pp_size: 1
  dp_replicate_size: 1
  dp_size: 8

# ── Checkpointing ──────────────────────────────────────────────────────────
# checkpoint_dir: where to save checkpoints (use a persistent path with Docker).
# restore_from: path to resume training from a previous checkpoint.
checkpoint:
  enabled: true
  checkpoint_dir: PATH_TO_YOUR_CKPT_DIR
  model_save_format: torch_save
  save_consolidated: false
  restore_from: null

Справочник по полям конфигурации

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

Запуск обучения

torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml

export MASTER_ADDR=node0.hostname
export MASTER_PORT=29500
export NODE_RANK=0  # 0 on master, 1 on second node, etc.

torchrun \
  --nnodes=2 \
  --nproc-per-node=8 \
  --node_rank=${NODE_RANK} \
  --rdzv_backend=c10d \
  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow_multinode.yaml

Генерация

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

python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_wan.yaml

python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_wan.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'

python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml

python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'

python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml

python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'

Интеграция с Diffusers

NeMo Automodel построен на базе Diffusers и использует его в качестве основы для загрузки моделей и вывода. Он загружает модели непосредственно из Hugging Face Hub с помощью классов моделей Diffusers, таких как WanTransformer3DModel, FluxTransformer2DModel и HunyuanVideoTransformer3DModel, и генерирует результаты через пайплайны Diffusers, например WanPipeline и FluxPipeline.

Эта интеграция предоставляет пользователям Diffusers несколько преимуществ:

• Конвертация контрольных точек не требуется: предобученные веса с Hub работают сразу. Укажите в pretrained_model_name_or_path любой идентификатор модели в формате Diffusers и начните обучение.
• Поддержка моделей с первого дня: когда новая диффузионная модель добавляется в Diffusers и загружается на Hub, её можно дообучить с помощью NeMo Automodel, не дожидаясь появления специального скрипта обучения.
• Совместимые с Pipeline выходные данные: дообученные контрольные точки сохраняются в формате, который можно напрямую загрузить обратно в пайплайны Diffusers для инференса, публикации на Hub или дальнейшей оптимизации с помощью инструментов, таких как квантование и компиляция.
• Масштабируемая тренировка моделей Diffusers: NeMo Automodel добавляет возможности распределённой тренировки (FSDP2, мультинодовый режим, мультирезольюционное бакетирование), выходящие за рамки того, что предлагают встроенные скрипты тренировки Diffusers, при этом сохраняя те же интерфейсы моделей и пайплайнов.
• Общая экосистема: любая модель, адаптер LoRA или компонент конвейера из экосистемы Diffusers сохраняет совместимость на всех этапах обучения и инференса.

Команда NVIDIA

• Pranav Prashant Thombre, pthombre@nvidia.com
• Linnan Wang, linnanw@nvidia.com
• Alexandros Koumparoulis, akoumparouli@nvidia.com

Ресурсы

• NeMo Automodel GitHub
• Руководство по дообучению диффузионных моделей
• Подготовка набора данных для диффузии
• Поддержка моделей диффузии
• NeMo Automodel для Transformers (дообучение LLM/VLM)