Vysoce výkonný přenos paměti GPU na AWS Sagemaker Hyperpod

Moderní infrastruktura pro hluboké učení často vyžaduje přenos velkých objemů dat mezi GPU napříč stroji. V Perplexity jsme narazili na jedinečnou technickou výzvu: efektivně přenášet nesouvislé oblasti paměti GPU mezi stroji maximální možnou rychlostí. Naše cílová platforma, instance AWS p5, nabízí působivou síťovou šířku pásma 3200 Gbps prostřednictvím 32 síťových karet. Tento článek sdílí naši cestu budování vlastního vysoce výkonného síťového řešení, které dosahuje 97,1 % této teoretické šířky pásma.

Technická výzva

Náš případ použití představoval několik klíčových technických požadavků:

• Vysoce propustný přenos mezi vzdálenými GPU s nesouvislými paměťovými bloky
• Schopnost dynamicky přidávat nebo odebírat uzly z nasazení Kubernetes bez narušení probíhajících operací
• Podpora pro vzorce peer-to-peer komunikace

Zatímco knihovna NCCL od NVIDIA je de facto standardem pro distribuované hluboké učení, nebyla ideální pro náš případ použití:

• NCCL vyniká v kolektivní komunikaci, ale vyžaduje vytvoření statického "světa", což vyžaduje restartování celého clusteru při úpravě zúčastněných uzlů.
• Synchronní komunikační model NCCL přidává složitost pro naši asynchronní pracovní zátěž
• Chtěli jsme přímou kontrolu nad vzorci přenosu paměti pro optimalizaci
• Budování vlastního řešení poskytlo cenné příležitosti k učení

Moderní vysoce výkonné sítě

Abychom pochopili naše řešení, prozkoumejme nejprve, jak se moderní vysoce výkonné sítě liší od tradičních sítí.

Většina sítí, které denně používáme, spoléhá na protokoly TCP/IP, kde aplikace komunikují se síťovou kartou přes jádro operačního systému pomocí soketů. Vysoce výkonné sítě však využívají RDMA (Remote Direct Memory Access) – zcela odlišný hardwarový a softwarový zásobník, který umožňuje přímý přístup do paměti mezi stroji bez zapojení CPU.

AWS poskytuje Elastic Fabric Adapter (EFA), vlastní síťové rozhraní, které implementuje vlastní protokol Amazonu nazvaný Scalable Reliable Datagram (SRD). Na rozdíl od tradičních sítí TCP/IP, kde musejí být data kopírována víckrát mezi uživatelským prostorem, prostorem jádra a síťovými vyrovnávacími pamětmi, EFA s RDMA umožňuje přímý přenos dat mezi pamětí GPU a síťovou kartou, čímž zcela obchází CPU.

Filozofie návrhu vysoce výkonných sítí

Budování vysoce výkonných síťových systémů vyžaduje přehodnocení několika základních předpokladů:

1. Vlastnictví vyrovnávací paměti: Na rozdíl od tradičních soketů, kde jádro spravuje síťové vyrovnávací paměti a vyžaduje kopírování mezi uživatelským prostorem a prostorem jádra, RDMA vyžaduje, aby si aplikace spravovaly vlastní vyrovnávací paměti. Když aplikace zahájí síťovou operaci, převede vlastnictví vyrovnávací paměti na síťovou kartu, dokud operace neskončí, což eliminuje potřebu kopírování dat.
2. Registrace paměti: Aplikace musí zaregistrovat oblasti paměti u jádra operačního systému. Jádro nastaví mapování virtuálních adres, které umožňuje CPU, GPU a síťovým kartám porozumět stejným virtuálním adresám. Tato registrace je jednorázová operace, která umožňuje následné přenosy dat s nulovým kopírováním.
3. Řídicí rovina vs. datová rovina: Vysoce výkonné sítě rozdělují operace do dvou kategorií:
   1. Operace řídicí roviny (jako je nastavení připojení a registrace paměti) procházejí jádrem pro zajištění bezpečnosti
   2. Operace datové roviny (skutečný přenos dat) obcházejí jádro pro dosažení maximálního výkonu
4. Příjem před přenosem: Bez vyrovnávacích pamětí spravovaných jádrem musejí aplikace předem odeslat operace příjmu a určit, kam mají být příchozí data umístěna. Jedná se o zásadní odklon od modelu soketů, kde aplikace mohou přijímat data kdykoli.
5. Dokončení založené na dotazování: Namísto čekání na oznámení jádra prostřednictvím mechanismů, jako je epoll, aplikace přímo dotazují hardwarové fronty dokončení. To eliminuje režii systémových volání a umožňuje okamžitou reakci na dokončené operace.
6. Povědomí o hardwarové topologii: Porozumění hardwarové topologii a její optimalizace je zásadní pro dosažení maximálního výkonu.

Pochopení hardwarové topologie

Instance AWS p5 mají sofistikovanou vnitřní architekturu. Jak je znázorněno níže, každá instance obsahuje dva sokety CPU tvořící dva uzly NUMA, přičemž každý uzel NUMA se připojuje ke čtyřem přepínačům PCIe:

Zvětšit

Pod každým přepínačem PCIe najdeme:

• Čtyři 100 Gbps síťové karty EFA
• Jednu NVIDIA H100 GPU
• Jeden NVMe SSD

Datové cesty pro přenosy TCP/IP a RDMA demonstrují základní rozdíly v jejich přístupech:

Zvětšit

U TCP/IP (levá strana) musejí být data kopírována víckrát:

1. Z paměti GPU do vyrovnávací paměti aplikace v uživatelském prostoru
2. Z vyrovnávací paměti uživatelského prostoru do vyrovnávací paměti jádra
3. Z vyrovnávací paměti jádra do síťové karty
4. Obrácený proces probíhá na přijímací straně

Každá operace kopírování spotřebovává cykly CPU a šířku pásma paměti. Aplikace musí také provádět přepínání kontextu mezi uživatelským prostorem a prostorem jádra pro každou síťovou operaci.

Naproti tomu RDMA (pravá strana) poskytuje skutečný přenos dat s nulovým kopírováním (zero-copy):

1. Síťová karta čte přímo z paměti GPU
2. Data cestují přímo na vzdálenou síťovou kartu
3. Vzdálená síťová karta zapisuje přímo do cílové paměti GPU
4. Aplikace potřebuje pouze zkontrolovat frontu dokončení (CQ) v uživatelském prostoru, aby věděla, kdy je přenos hotov

S RDMA a správným párováním hardwaru vyžaduje přenos dat mezi dvěma GPU pouze průchod místním přepínačem PCIe a sítí:

Zvětšit

Naproti tomu přenosy TCP/IP musejí kopírovat data víckrát přes hlavní paměť, což způsobuje značné přetížení sběrnice PCIe:

Zvětšit

Budování pomocí libfabric

Použili jsme libfabric – framework, který poskytuje obecné rozhraní pro služby fabric. Naše implementace využívá dva typy operací RDMA:

• Oousměrné RDMA (SEND/RECV) pro řídicí zprávy, které nesou metadata o oblastech paměti (např. posuny a velikosti).
• Jednosměrné RDMA WRITE pro vlastní přenos dat, kde každá operace WRITE zpracovává jeden souvislý paměťový blok.

Cesta k špičkovému výkonu

Náš vývoj postupoval v několika fázích:

1. Implementován základní jednosměrný přenos zpráv pomocí SEND/RECV
2. Rozšířeno na obousměrnou komunikaci
3. Přidáno GPUDirect RDMA WRITE pro přímý přenos mezi GPU
4. Rozšířeno pro zpracování více souběžných přenosů
5. Zavedeno řazení operací do fronty kvůli robustnosti
6. Dosaženo 97,4% využití šířky pásma na jedné síťové kartě

Při škálování na 32 síťových karet jsme implementovali několik klíčových optimalizací:

1. Řazení operací do fronty: Místo přímého odesílání operací na síťové karty udržujeme frontu na úrovni aplikace. To zajišťuje robustnost proti přetížení sítě a zjednodušuje programovací model.
2. Zahřátí sítě: Předběžné navázání spojení zlepšilo výkon při spuštění.
3. Vícevláknovost: Vyhrazená vlákna pro síťové operace každé GPU.
4. Připnutí jader CPU: Navázání vláken na konkrétní jádra CPU, aby se předešlo efektům NUMA a minulým mezipamětem.
5. Sharding stavu: Snížení konkistence mezi vlákny.
6. Dávkování operací: Odesílání několika operací najednou.
7. Zpětné odesílání operací (Lazy Operation Posting): Operace jsou nejprve zařazeny do fronty v aplikaci. Po dotázání fronty dokončení se pokusíme odeslat čekající operace na síťovou kartu, což zajišťuje efektivní využití síťových zdrojů.
8. Alokace prostředků s ohledem na NUMA: Alokace prostředků libfabric, jako jsou fronty dokončení, na správném uzlu NUMA za účelem minimalizace latence přístupu do paměti.

Prostřednictvím těchto optimalizací jsme dosáhli konečného výkonu 3 108 Gbps – 97,1 % teoretického maxima šířky pásma.

Níže uvedené video ukazuje náš program příkazové řádky v akci. Přenáší nesouvislé bloky paměti GPU do vzdáleného uzlu a dosahuje přenosové rychlosti 3108,283 Gbps – což demonstruje téměř teoretické využití šířky pásma síťové infrastruktury:

Závěr

Budování vysoce výkonného síťového systému vyžaduje pochopení jak hardwarové architektury, tak principů návrhu systému. Zatímco knihovny jako NCCL poskytují vynikající řešení pro běžné vzorce, někdy jsou pro specifické požadavky nutná vlastní řešení.

Naše cesta ukazuje, že dosažení téměř teoretického síťového výkonu je možné s pečlivou pozorností věnovanou architektuře systému, hardwarové topologii a různým optimalizačním technikám. Klíčem není jen porozumění jednotlivým komponentám, ale tomu, jak spolu interagují a tvoří kompletní systém.

Úplný technický rozbor této cesty včetně podrobností implementace a příkladů kódu je k dispozici v našem open-source repozitáři.

Citace

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}