Skalowalne fine-tuning modeli Diffusers z NVIDIA NeMo Automodel

Integracja NVIDIA NeMo Automodel z biblioteką Hugging Face Diffusers wprowadza spójny, produkcyjny workflow do trenowania i fine-tuningu dużych modeli obrazów i wideo w formacie Diffusers – bez konwersji checkpointów i bez przepisywania modeli. Dzięki temu zespoły mogą przejść od pojedynczej karty GPU do klastrów z dziesiątkami GPU, używając tych samych przepisów (YAML) i tego samego kodu.

Czym jest NVIDIA NeMo Automodel?

NeMo Automodel to otwartoźródłowa biblioteka treningowa oparta na PyTorch, natywnie wykorzystująca DTensor i strategie równoległe (FSDP2, tensor parallel, pipeline parallel, expert i context parallel), zaprojektowana jako część frameworka NVIDIA NeMo. Kluczowe założenie: „one program, any scale” – ten sam kod treningowy działa od jednej GPU po setki, a sposób równoleglenia definiuje się w konfiguracji, zamiast przepisywać model.

Drugie istotne założenie to „Hugging Face native”: użytkownik wskazuje parametr pretrained_model_name_or_path na ID modelu w Hubie (np. black-forest-labs/FLUX.1-dev), a Automodel ładuje klasę modelu z Diffusers (np. WanTransformer3DModel) i korzysta z pipeline’ów Diffusers (np. WanPipeline) do generacji. Dzięki temu fine-tunowane checkpointy można bezpośrednio załadować do DiffusionPipeline albo opublikować ponownie w Hugging Face Hub, zachowując kompatybilność z istniejącym toolingiem (quantization, LoRA, custom samplery).

Jakie modele Diffusers są wspierane?

Aktualna integracja NeMo Automodel zawiera gotowe przepisy fine-tuningu dla kilku kluczowych modeli wideo i obrazów, zapisanych w katalogu examples/diffusion/finetune. Są to zarówno modele text-to-image, jak i text-to-video o liczbie parametrów od ~1,3B do 32B.

Pod maską NeMo Automodel używa „flow matching” jako celu treningowego i trenuje w przestrzeni latentów, bazując na wcześniej zakodowanych wyjściach VAE oraz wielorozdzielczościowym bucketingu danych, co zmniejsza koszty obliczeniowe.

Co odblokowuje integracja z Diffusers?

W praktyce integracja NeMo Automodel + Diffusers daje kilka istotnych korzyści:

• Brak konwersji checkpointów – nie ma osobnego „formatu treningowego”, do którego trzeba konwertować wagi, a następnie wracać do Diffusers; checkpoint po fine-tuningu działa od razu w DiffusionPipeline.
• Szybka droga do wsparcia nowych modeli – dla nowego modelu Diffusers wystarczy przygotować adapter modelu i handler preprocessingowy, zamiast pisać kompletny, dedykowany skrypt treningowy.
• Wsparcie pełnego fine-tuningu i PEFT/LoRA – można wybrać pełne strojenie przy dużej liczbie GPU (maksymalna jakość) albo adaptery LoRA dla efektywnego fine-tuningu na pojedynczym nodzie.
• Skalowalny trening ponad możliwości typowych skryptów – NeMo Automodel dodaje FSDP2, tensor/context/pipeline parallel, wielowęzłową orkiestrację (SLURM, zapowiedziane Kubernetes) oraz multiresolution bucketing.

W efekcie można trenować modele, które wcześniej były praktycznie nieosiągalne w standardowym workflow Diffusers, takie jak FLUX.1-dev (12B parametrów) czy HunyuanVideo (13B parametrów).

Przykładowy workflow: fine-tuning FLUX.1-dev na talii tarota

Blogowy przykład pokazuje pełny fine-tuning transformera FLUX.1-dev na publicznej talii Rider–Waite (78 kart), aby nauczyć model specyficznego stylu kart tarota z wyzwalaczem trtcrd w promptach. Proces jest podzielony na trzy główne kroki: pre-encoding danych, uruchomienie treningu z użyciem istniejącej konfiguracji YAML oraz generacja z końcowego checkpointu.

1. Pre-encoding danych i bucketing

Zamiast kodować obrazy w każdej iteracji, pipeline najpierw przetwarza dataset do postaci latentów VAE i embeddingów tekstowych – to znacząco redukuje czas kroku treningowego. Preprocessing odbywa się narzędziem tools.diffusion.preprocessing_multiprocess, które strumieniuje obrazy bezpośrednio z Hugging Face (dataset multimodalart/1920-raider-waite-tarot-public-domain) i rozdziela pracę na wszystkie dostępne GPU.

Rezultatem jest struktura katalogu zawierająca podkatalog z bucketem rozdzielczości (np. 384x640), pliki cache .pt oraz pliki metadanych (metadata.json, metadata_shard_0000.json) powiązane z oryginalnymi obrazami w podkatalogu _hf_dataset/images. Bucket rozdzielczości dobierany jest tak, aby równoważyć budżet pikseli (tu ok. 245 760 pikseli) i proporcje obrazu.

2. Trening z użyciem istniejącego YAML (flux_t2i_flow.yaml)

Kod treningowy korzysta bezpośrednio z dostarczonego pliku examples/diffusion/finetune/flux_t2i_flow.yaml, który z góry definiuje m.in. wybór modelu, typ trenowania (full transformer fine-tuning), adapter flow-matching dla FLUX oraz efektywny batch size (GBS/LBS) i tryb FSDP2.

Specyficzne parametry związane z datasetem (np. cache_dir, bazowa rozdzielczość), schedulerem, liczbą kroków i ścieżkami checkpointów nadpisuje się przez parametry CLI, co pozwala uniknąć tworzenia nowych plików konfiguracji dla każdego eksperymentu. Trening generuje checkpointy co określoną liczbę kroków; finalny checkpoint nosi etykietę epoch_66_step_199, odpowiadając zakończonemu 200. krokowi optymalizatora z zerową indeksacją.

3. Generacja z fine-tunowanego checkpointu

Do generacji używa się gotowego pliku YAML dla inferencji (examples/diffusion/generate/configs/generate_flux.yaml), wskazując pole model.checkpoint na ścieżkę do końcowego checkpointu treningowego. Parametry takie jak wysokość i szerokość obrazu, liczba kroków samplingowych i prompt definiuje się jako parametry CLI.

W przykładzie porównia się dwa uruchomienia: z wyzwalaczem trtcrd i bez niego, przy identycznym seedzie, scenie i kompozycji. Z triggerem obraz zachowuje treść (astronauta na Marsie), ale przyjmuje styl kart tarota (kremowo-czerwona paleta, grube kontury, płaskie pola barw, patyna papieru), natomiast bez triggera pozostaje fotograficzny – pokazuje to, że efekt styli jest skojarzony z tokenem, a nie globalnie narzuca się na model.

Wydajność: obrazy i wideo na 8× H100 80GB

W artykule podano szczegółowe benchmarki dla text-to-image (512×512) oraz text-to-video (512×512×49 klatek) na pojedynczym nodzie z 8 GPU NVIDIA H100 80GB, z pełnymi batchami, wyłączonym zapisem checkpointów i mierzeniem pełnego czasu kroku (dataloader + forward + backward + optymalizer + scheduler).

Dla FLUX.1-dev pełny fine-tuning z FSDP2 i efektywnym GBS 32 (LBS 4) osiąga czas kroku ok. 0,9 s i ~35,5 obrazów/s, co daje ~4,4 obrazów/s na GPU przy szczytowym przydziale ok. 63,9 GiB z alokatora PyTorch. Z kolei wariant LoRA (r64) z DDP przy GBS 48 i LBS 6 osiąga ~0,89 s na krok i ~53,7 obrazów/s (~6,7 obrazów/s/GPU), kosztem nieco większego szczytowego zużycia pamięci (~67,4 GiB).

Dla Qwen-Image pełny fine-tuning z FSDP2 i GBS 40 (LBS 5) daje ~0,97 s na krok i ~41,2 obrazów/s (~5,15/GPU), natomiast LoRA r64 z DDP przy GBS 24 (LBS 3) skraca krok do ~0,51 s i daje ~46,6 obrazów/s (~5,8/GPU). Benchmarki dla modeli wideo (Wan 2.1, Wan 2.2, HunyuanVideo) pokazują, że pełny fine-tuning dużych modeli MoE wymaga kilku sekund na krok nawet przy 8 GPU, ale LoRA wyraźnie zmniejsza peak memory i pozwala zwiększyć GBS.

Przykłady specjalizacji domenowej i LoRA

Oprócz tarota, autorzy pokazują, że fine-tuning Wan 2.1 na zestaw wideo w stylu Ghibli skutecznie adaptuje styl wizualny – np. wygląd kwiatów istotnie się zmienia względem baseline. W wariancie z LoRA adapter zastosowany do Wan 2.1 sprawia, że wideo przyjmuje charakterystyczny, „ghibliowski” styl, widoczny m.in. w sposobie rysowania oczu postaci.

Podobne efekty obserwuje się dla FLUX.2, co pokazuje, że NeMo Automodel pozwala zarówno na pełne „przestrojenie” modeli pod konkretną domenę, jak i na lekką, efektywną adaptację przy użyciu LoRA. Ważne jest, że w obu przypadkach pozostaje kompatybilność z Diffusers – można korzystać z tych samych pipeline’ów, narzędzi do downstreamu i publikować modele w Hubie.

Praktyczne zastosowania dla inżyniera ML / MLOps

Integracja NeMo Automodel i Diffusers otwiera kilka praktycznych ścieżek wykorzystania w projektach produkcyjnych:

• Budowa wyspecjalizowanych generatorów obrazów (np. grafika techniczna, styl korporacyjny, layouty produktów) z wykorzystaniem pełnego fine-tuningu na klastrze GPU.
• Tworzenie generatorów wideo w określonych stylach (np. anime, kreskówka edukacyjna, wizualizacje przemysłowe) poprzez trening dużych modeli text-to-video (Wan, HunyuanVideo) na domenowych datasetach.
• Wykorzystanie LoRA do szybkiej personalizacji modelu dla konkretnego klienta lub kampanii (nowy styl, branding) bez ruszania pełnych wag modelu i przy znacznie mniejszym zapotrzebowaniu na zasoby.
• Integracja z istniejącym MLOps: YAML jako źródło prawdy do wersjonowania konfiguracji, a zapowiadane API Pythonowe jako wygodna warstwa do eksperymentów w notebookach i internal frameworkach.

Dla zespołów pracujących już z Diffusers ważne jest, że nie trzeba budować równoległego stacku treningowego – NeMo Automodel staje się „silnikiem treningowym pod spodem”, a wszystkie istniejące pipeline’y, adaptery i narzędzia pozostają w użyciu.

Następny krok: Pythonowe API przepisów

Obecnie główną ścieżką jest konfiguracja YAML, która dobrze nadaje się do reprodykowalnych eksperymentów, przeglądu przez zespoły i wersjonowania w repozytoriach. Autorzy zapowiadają jednak pełne, typowane API Pythonowe, które będzie odzwierciedlać te same elementy receptury: model, dane, optymalizator, PEFT/LoRA, równoległość, checkpointing, generację.

Takie API ułatwi integrację z istniejącym kodem treningowym, notebookami i workflowami eksperymentalnymi, zapewniając jednocześnie możliwość użycia YAML jako „konfiguracji produkcyjnej”. W praktyce oznacza to, że zespoły będą mogły mieć wspólne, utrzymywalne przepisy, a jednocześnie elastycznie składać komponenty w kodzie.

Źródła

• Fine-tune video and image models at scale with NVIDIA NeMo Automodel and 🤗 Diffusers (Hugging Face Blog)
• NeMo Automodel – Diffusers training guide
• NVIDIA-NeMo/Automodel – GitHub repository
• NeMo Automodel – Transformers community integration