NVIDIA i Hugging Face integrują NeMo Automodel z Diffusers dla skalowania treningu modeli AI

NVIDIA i Hugging Face ogłosiły integrację biblioteki NeMo Automodel z ekosystemem Diffusers, która umożliwia skalowalne trenowanie i dostrajanie modeli dyfuzji bezpośrednio z Hugging Face Hub. Współpraca ta eliminuje potrzebę konwersji checkpointów i przepisywania modeli, oferując gotowe przepisy dla najpopularniejszych modeli generatywnych, jak podaje oficjalny wpis na blogu Hugging Face.

Integracja została udokumentowana w oficjalnym przewodniku treningowym Diffusers i jest dostępna jako oprogramowanie open source na licencji Apache 2.0. Rozwiązanie adresuje rosnące zapotrzebowanie na narzędzia do trenowania modeli dyfuzji, które oferują efektywne zarządzanie pamięcią, cachowanie latentów oraz konfiguracje skalujące się od jednego GPU do setek.

Kluczowe wnioski

• NeMo Automodel obsługuje natywnie modele z Hugging Face Hub bez konieczności konwersji checkpointów — wystarczy wskazać ID modelu i rozpocząć trening.
• Biblioteka oferuje gotowe przepisy dla sześciu głównych modeli: Wan 2.1/2.2, FLUX.1-dev/FLUX.2-dev, HunyuanVideo 1.5 i Qwen-Image.
• Wspiera zarówno pełne dostrajanie, jak i efektywne metody LoRA, z automatyczną konfiguracją różnych schematów równoległości (FSDP2, tensor parallel, pipeline parallel).
• Nowe modele dyfuzji w Diffusers wymagają jedynie małych dodatków kodu zamiast kompletnych skryptów treningowych.
• Integracja umożliwia skalowanie od pojedynczego GPU (Wan 2.1 1.3B na A100 40GB) do wielowęzłowych klastrów dla największych modeli.

Architektura i możliwości NeMo Automodel

NeMo Automodel to biblioteka treningowa PyTorch DTensor oparta na dwóch kluczowych zasadach projektowych. Pierwsza to natywna kompatybilność z Hugging Face — użytkownicy mogą wskazać dowolny identyfikator modelu z Hub jako pretrained_model_name_or_path i od razu rozpocząć trening. Biblioteka wykorzystuje klasy modeli Diffusers (np. WanTransformer3DModel) do ładowania oraz pipeline'y Diffusers (WanPipeline) do generowania.

Druga zasada to ”jeden program, dowolna skala” — przepisy i skrypty treningowe można łatwo modyfikować do treningu w każdej skali. Równoległość to wybór konfiguracyjny, nie przepisywanie kodu. Przełączanie między FSDP2, tensor parallel, expert parallel, context parallel i pipeline parallel odbywa się przez deklaracje konfiguracji.

Pod maską biblioteka wykorzystuje flow matching jako cel treningowy, z treningiem w przestrzeni latentnej (via pre-encoded VAE outputs) oraz multiresolution bucketed dataloading dla przyspieszenia przepustowości.

Obsługiwane modele i przepisy treningowe

Biblioteka dostarcza gotowe przepisy dostrajania dla sześciu rodzin modeli. Wan 2.1 T2V w wersji 1.3B mieści się na pojedynczym A100 40GB, podczas gdy wersja 14B wymaga większych zasobów. FLUX.1-dev (12B parametrów) i jego następca FLUX.2-dev (32B) obsługują generowanie obrazów z tekstu z opcją LoRA.

HunyuanVideo 1.5 oferuje 13 miliardów parametrów dla zadań text-to-video, a Qwen-Image wykorzystuje architekturę MMDiT z 20 miliardami parametrów. Wan 2.2 T2V A14B wyróżnia się jako model Mixture of Experts z 27 miliardami parametrów łącznie, ale tylko 14 miliardami aktywnymi na krok.

Każdy model ma dedykowany przepis w katalogu examples/diffusion/finetune, a większość obsługuje zarówno pełne dostrajanie, jak i efektywne metody LoRA. Wyjątkiem jest Wan 2.2 MoE, dla którego przepis LoRA nie jest jeszcze dostępny.

Workflow i praktyczne zastosowanie

Przepływ pracy składa się z trzech głównych kroków. Pierwszy to pre-enkodowanie datasetu, które przyspiesza trening przez wcześniejsze przetworzenie danych przez VAE. Drugi krok to uruchomienie treningu z wykorzystaniem istniejącego pliku YAML — na przykład dla FLUX wystarczy komenda wskazująca na odpowiedni przepis konfiguracyjny.

Ostatni krok to generowanie z dostrojonego checkpointu, które działa bezpośrednio z DiffusionPipeline bez dodatkowych konwersji. Wydajność skaluje się efektywnie — przykładowe testy pokazują, że można skutecznie wykorzystać zasoby od pojedynczych GPU po wielowęzłowe klastry.

Autorzy zapowiadają również nadchodzące Pythonic recipe APIs, które uproszą konfigurację i uruchamianie treningu bezpośrednio z kodu Python, eliminując potrzebę zarządzania plikami YAML.

Ta integracja stanowi znaczący krok w demokratyzacji dostępu do zaawansowanych technik trenowania modeli dyfuzji, łącząc łatwość użytkowania Diffusers z mocą obliczeniową rozwiązań NVIDIA.