Zrychlení jemného ladění Transformerů s NVIDIA NeMo AutoModel

NVIDIA NeMo AutoModel představuje významný krok vpřed v oblasti jemného ladění modelů Transformer, zejména těch založených na architektuře Mixture-of-Experts (MoE). Tento nástroj, postavený na základech HuggingFace Transformers v5, přináší řadu optimalizací, které umožňují dosáhnout až 3.4-3.7x vyššího tréninkového průtoku a zároveň snižují paměťové nároky na GPU o 29-32 %.

Výhody NeMo AutoModel

Jednou z klíčových vlastností NeMo AutoModel je jeho kompatibilita s API HuggingFace Transformers. To znamená, že vývojáři mohou snadno integrovat NeMo AutoModel do svých stávajících projektů bez nutnosti zásadních změn v kódu. Například, pro model Nemotron 3 Ultra 550B A55B bylo při plném jemném ladění dosaženo průměrného tréninkového průtoku (TPS/GPU) 815.

Expert Parallelism a TransformerEngine

NeMo AutoModel využívá technologii Expert Parallelism (EP), která snižuje paměťový tlak tím, že distribuuje váhy expertů napříč GPU. Tímto způsobem se dosahuje efektivního využití paměti, což je zvláště důležité při trénování velkých modelů. Při trénování modelu Qwen3-30B-A3B byl průměrný TPS/GPU 11,340.

Další technologií, která přispívá k výkonu NeMo AutoModel, je TransformerEngine. Tato jádra zrychlují základní operace a poskytují konzistentní zrychlení oproti PyTorch/Flash Attention ekvivalentům. Kombinace EP a TransformerEngine umožňuje efektivní trénování modelů s využitím více GPU bez nutnosti dalších přepisů kódu.

Benchmarky a výkon

V rámci benchmarků byly provedeny testy na různých modelech, přičemž NeMo AutoModel prokázal své schopnosti v širokém spektru konfigurací. Například, při trénování modelu Qwen3-30B-A3B byla maximální paměťová náročnost 48.1 GiB.

Zajímavé na tom je, že NeMo AutoModel integruje také DeepEP, což fúzuje komunikaci s výpočtem a dále zrychluje proces trénování. Tato optimalizace je klíčová pro dosažení vysokého výkonu při trénování velkých modelů.