Modele dyfuzyjne napędzały jedne z najbardziej ekscytujących wydań open-source ostatnich dwóch lat — pomyśl o FLUX.1-dev do generowania obrazu z tekstu, czy Wan 2.1 i HunyuanVideo do generowania wideo z tekstu. Biblioteka 🤗 Diffusers stała się domem dla tych modeli, oferując badaczom i twórcom spójny interfejs do wnioskowania, adaptacji i komponowania potoków.

Trenowanie i dostrajanie tych modeli również zyskuje na popularności, co wymaga narzędzi obsługujących pamięciowo wydajne shardowanie, cachowanie latentów, wielorozdzielcze bucketing oraz konfiguracje, które płynnie skalują się od jednego GPU do setek.

Wkracza biblioteka open-source NVIDIA NeMo Automodel. Dziś przyglądamy się współpracy NVIDIA i Hugging Face, która wprowadza produkcyjne, rozproszone trenowanie dyfuzyjne dla dowolnego modelu w formacie Diffusers z Hugging Face Hub — bez konwersji punktów kontrolnych, bez przepisywania modeli dla nowych modeli. Integracja jest udokumentowana w przewodniku trenowania Diffusers i jest w pełni open-source na licencji Apache 2.0.

Czym jest NeMo Automodel?

NeMo Automodel to open-source'owa biblioteka treningowa natywna dla PyTorch DTensor, będąca częścią frameworka NVIDIA NeMo. Została zbudowana wokół dwóch zasad projektowych, które mają znaczenie dla ekosystemu Diffusers:

  • Natywność Hugging Face. Wskaż pretrained_model_name_or_path na dowolny identyfikator modelu Diffusers na Hubie i rozpocznij trenowanie. NeMo Automodel używa klas modeli Diffusers (np. WanTransformer3DModel) do ładowania oraz potoków Diffusers (WanPipeline) do generowania. Punkty kontrolne wracają czysto do ekosystemu Diffusers.

  • Jeden program, dowolna skala. Przepisy i skrypty treningowe można łatwo modyfikować, aby dostosować je do trenowania w dowolnej skali. Równoległość jest wyborem konfiguracyjnym, a nie przepisywaniem kodu — przełączaj się między FSDP2, równoległością tensorową, ekspercką, kontekstową i potokową, deklarując konfiguracje, a nie przepisując modele.

AutoModel obecnie obsługuje tylko modele z dopasowaniem przepływu (flow-matching). Pod maską używa dopasowania przepływu jako celu treningowego, z trenowaniem w przestrzeni latentnej (poprzez wstępnie zakodowane wyjścia VAE) i wielorozdzielczym bucketingiem danych w celu przyspieszenia przepustowości.

Obsługiwane modele dyfuzyjne

Integracja NeMo Automodel dostarcza gotowe przepisy do dostrajania dla poniższych otwartych modeli dyfuzyjnych. Lista odzwierciedla przepisy obecnie znajdujące się w examples/diffusion/finetune.

ModelIdentyfikator Hugging FaceZadanieParametryPrzepis LoRAWan 2.1 T2V 1.3B / 14BWan-AI/Wan2.1-T2V-1.3B-Diffusers / Wan-AI/Wan2.1-T2V-14B-DiffusersTekst na wideo1.3B (mieści się na pojedynczym 40GB A100) / 14BTakWan 2.2 T2V A14BWan-AI/Wan2.2-T2V-A14B-DiffusersTekst na wideo27B łącznie (MoE), 14B aktywne na krokNieFLUX.1-devblack-forest-labs/FLUX.1-devTekst na obraz12BTakFLUX.2-devblack-forest-labs/FLUX.2-devTekst na obraz32BTakHunyuanVideo 1.5hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2vTekst na wideo13BTakQwen-ImageQwen/Qwen-ImageTekst na obraz20B (MMDiT)Tak

Co ta współpraca odblokowuje

Dla użytkowników Diffusers praktyczne korzyści sprowadzają się do kilku konkretnych możliwości.

Brak konwersji punktów kontrolnych. Wstępnie wytrenowane wagi z Huba działają od razu. Nie ma osobnego „formatu treningowego” do konwersji, a potem z powrotem. Twój dostrojony punkt kontrolny ładuje się bezpośrednio do DiffusionPipeline do wnioskowania lub z powrotem na Huba do udostępniania. Narzędzia downstream — kwantyzacja, kompilacja, adaptery LoRA, niestandardowe samplery — nadal działają.

Szybka ścieżka do obsługi nowych modeli. Gdy nowy model dyfuzyjny pojawi się w Diffusers, włączenie go w NeMo Automodel wymaga niewielkiego, ograniczonego dodatku kodu — handlera przetwarzania danych i adaptera modelu — a nie pełnego niestandardowego skryptu treningowego. Reszta stosu przepisów (FSDP2, bucketing danych, punktowanie kontrolne, generowanie) pozostaje bez zmian, a ten sam przepływ pracy oparty na YAML ma zastosowanie.

Pełne i parametrycznie efektywne dostrajanie. Obsługiwane jest zarówno pełne dostrajanie, jak i PEFT w stylu LoRA, więc możesz wybrać między maksymalną jakością (pełne FT na dużym klastrze) a maksymalną wydajnością (LoRA na pojedynczym węźle). Ta sama struktura przepisu obsługuje oba.

Skalowalne trenowanie wykraczające poza to, co oferują wbudowane skrypty. NeMo Automodel dodaje schematy shardowania, takie jak FSDP2, równoległość tensorowa, kontekstowa i potokowa, orkiestrację wielowęzłową (obecnie SLURM, Kubernetes w przygotowaniu) oraz wielorozdzielcze bucketing. Te możliwości umożliwiają trenowanie większych modeli, takich jak FLUX.1-dev (12B) i HunyuanVideo (13B).

Przegląd przepływu pracy dostrajania

W tej sekcji przeprowadzimy Cię przez typowy przepływ pracy dostrajania dowolnego z obsługiwanych modeli. Zalecanym sposobem instalacji Automodel jest kontener Docker NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06), który zawiera wstępnie zbudowane PyTorch, TransformerEngine i inne zależności skompilowane z CUDA. Alternatywnie, zainstaluj za pomocą pip3 install nemo-automodel lub ze źródła (pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git); zobacz przewodnik instalacji, aby poznać wszystkie opcje.

Ten przewodnik przeprowadzi przez pełne dostrajanie transformera FLUX.1-dev na 78-kartowym zbiorze danych tarota Rider–Waite, a następnie generowanie z wynikowego punktu kontrolnego. Wykorzystuje on sprawdzone konfiguracje YAML i stosuje ustawienia specyficzne dla uruchomienia jako nadpisania wiersza poleceń, więc nie są wymagane żadne nowe pliki konfiguracyjne.

1. Wstępne kodowanie zbioru danych

Przepis dyfuzyjny zużywa cachowane latenty VAE i osadzenia tekstu zamiast kodowania obrazów źródłowych podczas każdego kroku treningowego. Przesyłaj strumieniowo 78 obrazów Rider–Waite bezpośrednio z Hugging Face i rozdziel przetwarzanie wstępne na wszystkie widoczne GPU:

uv run --locked --no-default-groups \
  --extra diffusion \
  --extra diffusion-media \
python -m tools.diffusion.preprocessing_multiprocess image \
  --dataset_name multimodalart/1920-raider-waite-tarot-public-domain \
  --dataset_media_column image \
  --dataset_caption_column caption \
  --dataset_streaming \
  --max_images 78 \
  --output_dir /cache/flux_tarot \
  --processor flux \
  --model_name black-forest-labs/FLUX.1-dev \
  --max_pixels 245760

Podpisy zawierają już token wyzwalający trtcrd. Przy tym budżecie pikseli i proporcjach portretowych zbioru danych, przetwarzanie wstępne przypisuje próbki do zasobnika 384×640 używanego w demonstracyjnym uruchomieniu.

W przypadku trenowania obrazów, przetwarzanie wstępne tworzy pliki pamięci podręcznej .pt i shardowane metadane:

/cache/flux_tarot/
├── 384x640/
│   ├── <hash1>.pt
│   └── ...
├── metadata_shard_0000.json
├── metadata.json
└── _hf_dataset/
    └── images/

2. Uruchom trenowanie z istniejącym YAML FLUX

Użyj bezpośrednio examples/diffusion/finetune/flux_t2i_flow.yaml. YAML już wybiera FLUX.1-dev, pełne dostrajanie transformera, adapter dopasowania przepływu FLUX, efektywny rozmiar partii 32 i ośmiokierunkowy FSDP2.

Podaj ścieżki i ustawienia specyficzne dla tarota jako nadpisania wiersza poleceń:

uv run --locked --no-default-groups --extra diffusion \
  torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/flux_t2i_flow.yaml \
  --model.transformer_engine_fp8 false \
  --data.dataloader.cache_dir /cache/flux_tarot \
  --data.dataloader.base_resolution '[384,640]' \
  --lr_scheduler.lr_decay_style constant \
  --lr_scheduler.lr_warmup_steps 20 \
  --step_scheduler.max_steps 200 \
  --step_scheduler.ckpt_every_steps 50 \
  --checkpoint.checkpoint_dir /tmp/flux_tarot/checkpoints/full \
  --checkpoint.save_consolidated true \
  --seed 2026

Uruchomienie tworzy punkty kontrolne w krokach 50, 100, 150 i 200. Ostateczny punkt kontrolny jest oznaczony jako epoch_66_step_199; etykieta jest oparta na zero, mimo że reprezentuje ukończony 200. krok optymalizatora.

3. Generuj z dostrojonego punktu kontrolnego

Użyj istniejącego YAML generowania FLUX i wskaż model.checkpoint na kompletny punkt kontrolny treningu:

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["a trtcrd of an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/full/step_200 \
  --seed 2026

Dołącz trtcrd, aby wywołać wyuczony styl tarota. Dla porównania kontrolnego, zachowaj ziarno i scenę, ale pomiń wyzwalacz:

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/control \
  --seed 2026

Wyniki

W kroku 200, wyzwolone podpowiedzi astronauty zachowują żądaną treść, jednocześnie nabierając kremowej, czerwonej i czarnej vintage palety, grubych konturów atramentowych, płaskich pól koloru, tonów starego papieru i alegorycznej kompozycji kart. Niewyzwolony astronauta pozostaje fotograficzny, co dowodzi, że wyuczony efekt jest w znacznym stopniu związany z trtcrd, a nie zastępuje globalnie modelu bazowego.

Podpowiedź (ziarno)Linia bazowaDostrojony (krok 200)Astronauta pielęgnujący ogród róż na Marsie (ziarno 2026)

4. Wydajność

Wszystkie pomiary zostały zebrane na jednym węźle z 8× NVIDIA H100 80GB GPU. Wyniki to średnie ± odchylenie standardowe próbki z trzech okien 10 kroków w stanie ustalonym.

Tekst na obraz — 512×512

ModelTrenowanieRównoległośćGBS / LBSCzas krokuObrazy/sObrazy/s/GPUSzczytowe przydzielone/GPUFLUX.1-devPełneFSDP232 / 40.902 ± 0.039 s35.51 ± 1.554.44 ± 0.1963.88 GiBFLUX.1-devLoRA r64DDP48 / 60.894 ± 0.008 s53.73 ± 0.486.72 ± 0.0667.43 GiBQwen-ImagePełneFSDP240 / 50.974 ± 0.075 s41.21 ± 3.065.15 ± 0.3853.55 GiBQwen-ImageLoRA r64DDP24 / 30.515 ± 0.006 s46.63 ± 0.545.83 ± 0.0766.33 GiB

Tekst na wideo — 512×512×49 klatek

Każda próbka to jeden 49-klatkowy klip wideo.

ModelTrenowanieGBS / LBSKontrola punktów aktywacjiCzas krokuKlipy/sKlipy/s/GPUSzczytowe przydzielone/GPUWan 2.1 1.3BPełne8 / 1Wył.0.942 ± 0.038 s8.50 ± 0.351.06 ± 0.046.09 GiBWan 2.1 14BPełne8 / 1Wł.3.798 ± 0.017 s2.107 ± 0.0060.263 ± 0.00633.35 GiBWan 2.1 14BLoRA r6416 / 2Wł.7.585 ± 0.014 s2.110 ± 0.0000.263 ± 0.00024.07 GiBWan 2.2 A14B, wysoki szumPełne8 / 1Wł.4.628 ± 0.031 s1.730 ± 0.0100.217 ± 0.00623.57 GiBHunyuanVideo 1.5Pełne8 / 1Wł.5.926 ± 0.046 s1.350 ± 0.0100.170 ± 0.00015.90 GiBHunyuanVideo 1.5LoRA r648 / 1Wł.5.575 ± 0.006 s1.433 ± 0.0060.180 ± 0.00010.58 GiB

Szczegóły pomiarów

  • Sprzęt: 8× H100 80GB HBM3, w pełni połączone NVLink.

  • Zbiór danych obrazów: lambda/naruto-blip-captions, 256 buforowanych próbek.

  • Zbiór danych wideo: svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned, 112 buforowanych próbek.

  • Pełne partie były wymuszane za pomocą drop_last=true.

  • Zapisy punktów kontrolnych były wyłączone.

  • Czas kroku obejmuje ładowanie danych, forward, backward, przycinanie gradientów, optymalizator i pracę harmonogramu.

  • Pamięć to szczytowe użycie alokatora CUDA PyTorch, a nie całkowita pamięć urządzenia NVML.

Inne przykłady dostrajania/LoRA

Wyniki dostrajania i LoRA pokazują moc NeMo Automodel w specjalizacji domenowej. Na przykład, dostrojenie modelu Wan 2.1 na zbiorze danych wideo Ghibli z powodzeniem dostosowało styl wyjściowy, co zademonstrowano zauważalną zmianą wyglądu kwiatu w porównaniu do linii bazowej.

Linia bazowa:

Dostrojony na filmach Ghibli:

Zaobserwowaliśmy również wyraźny wpływ użycia LoRA, gdzie zastosowanie adaptera do Wan 2.1 spowodowało, że wideo przyjęło charakterystyczny styl Ghibli, szczególnie widoczny w podkreśleniu oczu postaci.

Bez LoRA:

Z LoRA:

Te przykłady, w tym dla FLUX.2, potwierdzają, że użytkownicy mogą osiągnąć zarówno maksymalną jakość poprzez pełne dostrajanie, jak i maksymalną wydajność poprzez PEFT w stylu LoRA, dostosowując wyjście do określonych domen stylistycznych.

Wypróbuj dziś

Dowiedz się więcej o integracji i znajdź więcej przykładów dostrajania w dokumentacji NeMo Automodel

Co dalej: Pythoniczne API przepisów

YAML dobrze sprawdza się w przypadku powtarzalnej konfiguracji, szczególnie dla zespołów, które chcą plików, które mogą zatwierdzić, przejrzeć i ponownie wykorzystać, ale wiele zespołów potrzebuje również interfejsu programistycznego.

W nadchodzącej wersji NeMo Automodel planujemy udostępnić przepisy dyfuzyjne również poprzez w pełni typowane Pythoniczne API. Użytkownicy będą mogli skomponować te same elementy modelu, danych, optymalizatora, PEFT/LoRA, równoległości, punktowania kontrolnego i generowania bezpośrednio z Pythona.

Pythoniczna ścieżka ma na celu ułatwienie korzystania z przepisów w istniejącym kodzie treningowym, notebookach i przepływach pracy eksperymentów oraz zaoferowanie pierwszorzędnego interfejsu Pythonicznego obok szybkiej ścieżki YAML.