Vědci dosáhli superlineárního zrychlení při trénování 100miliardového modelu

Nová éra heterogenního strojového učení

Výzkumníci ze Shanghai Artificial Intelligence Laboratory pod vedením Ding Tang, Jiecheng Zhou, Jiakai Hu, Shengwei Li, Huihuang Zheng, Zhilin Pei, Hui Wang a Xingcheng Zhang představili průlomový framework nazvaný H2 (HyperHetero), který umožňuje efektivní trénování velkých jazykových modelů na clusterech obsahujících více než 1000 heterogenních čipů od různých výrobců. Tato technologie řeší jeden z nejvýznamnějších problémů současného vývoje umělé inteligence - jak efektivně využít výpočetní prostředky různých typů čipů, které mají odlišné vlastnosti v oblasti výpočetního výkonu, paměti a komunikačních schopností.

Běžné přístupy k distribuovanému trénování velkých modelů byly navrženy pro homogenní prostředí, kde všechny čipy pocházejí od jednoho výrobce a mají podobné charakteristiky. Současná situace v průmyslu však vyžaduje využívání diverzifikovaných hardwarových zdrojů, protože spoléhání pouze na jeden typ čipu může být nepostačující kvůli omezením výrobní kapacity nebo specifické vhodnosti pro určité výpočetní úlohy. Hyper-heterogenní výpočetní prostředí se vyznačuje třemi klíčovými charakteristikami: výrazně odlišnými hardwarovými specifikacemi bez jasných vzorců, softwarovou izolací mezi čipy různých výrobců a vysoce nevyváženým zastoupením různých typů čipů v rámci jednoho clusteru.

Překonávání technické izolace

Framework H2 řeší problém technické izolace mezi různými typy čipů pomocí tří hlavních komponent. První komponentou je DiTorch, což je jednotné PyTorch-kompatibilní rozhraní, které standardizuje knihovny operátorů a runtime prostředí napříč heterogenními čipy. DiTorch využívá PyTorch jako jednotnou programovací vrstvu a implementuje bottom-up pipeline pro zarovnání přesnosti, který nejdříve vynucuje numerickou konzistenci na úrovni operátorů napříč různými čipy a poté dosahuje end-to-end zarovnání přesnosti na úrovni modelu.

DiTorch integruje dvě hlavní strategie pro zvýšení interoperability napříč různými AI hardwarovými platformami. Torch Adapter může být dodán výrobcem čipu, nebo lze přímo použít Device-Independent Process Unit (DIPU), který harmonizuje základní runtime systémy od různých výrobců do jednotných API. Kromě toho Device-Independent Operator Interface (DIOPI) propojuje jednotné Torch operátorové API s knihovnami operátorů přizpůsobenými konkrétním výrobcům, čímž zajišťuje konzistentní sémantickou interpretaci napříč operátory. V současnosti je v DIOPI implementováno více než 300 standardizovaných rozhraní operátorů.

Druhou komponentou je DiComm, jednotná komunikační knihovna, která překonává izolaci a neefektivitu v komplexních heterogenních prostředích. DiComm využívá knihovnu libibverbs a umožňuje RDMA komunikaci napříč různými architekturami čipů, podporuje jak homogenní, tak heterogenní chip-to-chip interakce. Knihovna implementuje dva odlišné komunikační modely: CPU-mediated a device-direct. V device-direct metodě každý čip registruje své lokální paměťové oblasti s RDMA ovladačem, což umožňuje RDMA-enabled NIC mapovat fyzické adresy pro vzdálený přístup. Experimentální výsledky ukazují, že device-direct RDMA snižuje průměrnou latenci o 9,94× oproti konvenčnímu TCP/IP schématu.

Heterogenní paralelismus a automatické vyhledávání strategií

Třetí klíčovou komponentou je HeteroPP framework spolu s HeteroAuto, adaptivní strategií pipeline paralelismu, která dynamicky balancuje výpočetní zátěž, paměťová omezení a komunikační overhead. HeteroPP rozšiřuje pipeline paralelismus na heterogenní AI Chip-Clustery pro distribuované LLM trénování. Každá pipeline stage se skládá výhradně z homogenního typu čipů, přičemž heterogenní uzly jsou strategicky distribuovány napříč různými stages. Čipy s větší paměťovou kapacitou jsou mapovány na dřívější stages pipeline, zatímco čipy s menší pamětí jsou alokovány na pozdější stages.

HeteroAuto představuje automatizovanou vyhledávací metodu pro optimalizaci konfigurací paralelních strategií v heterogenních čipových prostředích. Systém využívá depth-first search (DFS) metodologii a heuristické vyhledávání pro odvození nejlepší distribuce vrstev. Algoritmus nejdříve identifikuje vhodnou velikost data paralelismu a poté rozděluje velké skupiny homogenních čipů na menší podskupiny, které jsou zpracovávány jako odlišné heterogenní entity. Pro zvýšení efektivity trénování v heterogenních nastaveních byl vyvinut topology-aware activation resharding strategy, který konfiguruje každý čipový server tak, aby přiřadil dedikovaný komunikační NIC každému čipu na základě jeho NIC afinity.

Experimentální validace a dosažené výsledky

Výzkumníci validovali svůj framework na 100miliardovém parametrickém modelu s architekturou založenou na strukturních principech LLaMA. Model obsahuje 96 vrstev, 64 attention heads, 8192 hidden size a 36864 intermediate size s maximální délkou sekvence 4096 tokenů. Testování probíhalo se čtyřmi různými typy AI čipů označených jako Chip-A, Chip-B, Chip-C a Chip-D, které se lišily výpočetním výkonem, paměťovou kapacitou a počtem čipů na uzel.

Experimentální výsledky demonstrují výjimečnou efektivitu H2 frameworku. V konfiguraci s 768 čipy tří typů dosáhl systém HeteroSpeedupRatio 109,03%, zatímco s 1024 čipy čtyř typů dosáhl 104,29%. I při konstantní globální batch size dosáhl systém 89,56% a 77,45% respektive. Tyto výsledky ukazují, že heterogenní trénování může dosáhnout konkurenceschopné, a v některých případech i nadřazené efektivity ve srovnání s baseline throughput homogenního trénování.

Superlineární zlepšení výkonu je vysvětlitelné tím, že konvenční 3D paralelní trénování má tendenci přehlížet nevyvážené požadavky na zdroje mezi různými výpočetními úlohami, zatímco HeteroPP framework s HeteroAuto kapitalizuje na těchto nerovnováhách inteligentní alokací čipových úloh a jemným laděním trénovacích hyperparametrů na základě specifických požadavků na zdroje. V praktických produkčních prostředích obvykle čipy s nižšími specifikacemi vykazují výrazně nižší ceny a sníženou spotřebu energie ve srovnání s high-spec čipy, což činí framework ekonomicky výhodným.

Budoucí perspektivy a význam

Framework H2 představuje významný pokrok v oblasti distribuovaného strojového učení tím, že umožňuje efektivní integraci různorodých hardwarových zdrojů do jednotného výpočetního systému. Tato technologie má potenciál revolucionalizovat způsob, jakým organizace přistupují k tréninku velkých jazykových modelů, protože jim umožňuje využívat existující heterogenní infrastrukturu namísto nutnosti investovat do homogenních clusterů. Výsledky experimentů na 100miliardovém modelu potvrzují proveditelnost a efektivitu hyper-heterogenního trénování v dosud nevídaných měřítkách, což otevírá nové možnosti pro vývoj ještě větších a komplexnějších AI systémů využívajících diverzifikované výpočetní zdroje.