Sharded weight transfer pro model Kimi K2: 7.53 sekund na 48 uzlech

V oblasti strojového učení, zejména v kontextu reinforcement learning (RL), je efektivní synchronizace modelových vah klíčová pro zajištění aktuálnosti a výkonu modelů. Článek se zaměřuje na implementaci sharded weight transfer v rámci vLLM, která využívá Ray Direct Transport (RDT) a přináší zlepšení v rychlosti a efektivitě přenosu vah.

Pozadí

Tradiční metoda synchronizace vah pomocí NCCL broadcast má své limity, zejména při práci s velkými modely. Každý pracovník přijímá celý model, což vede k neefektivnímu využití paměti a zpomaluje přenos. To je obzvlášť problematické pro modely s velkým počtem parametrů, jako je Kimi K2.

Sharded weight-transfer engine s RDT

Implementace sharded weight transfer v vLLM s RDT umožňuje, aby každý inference worker vyžádal pouze ty sharded váhy, které potřebuje, podle svého záznamového plánu. To snižuje množství dat, která je třeba přenést, a zlepšuje celkovou efektivitu.

Podle testování trvá sharded weight transfer pro model Kimi K2 v BF16 pouze 7.53 sekund na 48 uzlech, což je významné zrychlení oproti předchozí implementaci, která trvala 64.72 sekund. Tato optimalizace přenosu vah přispěla k poklesu end-to-end latence z 25.02 sekund na 3.49 sekund.

Výkonové optimalizace

Celkové množství dat přenesených během synchronizace vah činí 7.9 TB, přičemž průměrná šířka pásma dosažená během synchronizace vah je 1,049 GB/s. Tyto výsledky ukazují, jak efektivně je možné přenášet velké objemy dat v rámci sharded weight transferu.

NIXL backend byl zvolen pro flexibilní P2P komunikaci, což zajišťuje odolnost proti chybám během rollouts. To znamená, že pokud dojde k selhání některého z inference engine, systém může pokračovat v provozu a obnovit se po opravě.

Omezení a budoucí plány

I když implementace RDT přináší mnoho výhod, existují i omezení. V současné době je podpora omezena na GPU -> GPU transfer, což může omezit flexibilitu nasazení v různých prostředích. Dále není implementace kompatibilní s EPLB v vLLM a váhy se aktuálně přenášejí sériově napříč skupinami PP, aby se předešlo OOM.

Budoucí plány zahrnují rozšíření podpory pro remote GPU -> CPU transfer, což by mohlo dále zjednodušit a zefektivnit proces synchronizace vah.