Augstas veiktspējas GPU atmiņas pārsūtīšana AWS Sagemaker Hyperpod

Modernā dziļās mācīšanās infrastruktūra bieži prasa lielu datu apjomu pārsūtīšanu starp GPU dažādās mašīnās. Uzņēmumā Perplexity mēs saskārāmies ar unikālu tehnisku izaicinājumu: efektīvi pārsūtīt nesadalītas GPU atmiņas apgabalus starp mašīnām ar maksimālo iespējamo ātrumu. Mūsu mērķa platforma, AWS p5 instances, piedāvā iespaidīgu tīkla joslas platumu 3200 Gbps, izmantojot 32 tīkla kartes. Šajā rakstā ir aprakstīts mūsu ceļš, izveidojot pielāgotu augstas veiktspējas tīkla risinājumu, kas sasniedz 97,1% no šī teorētiskā joslas platuma.

Tehniskais izaicinājums

Mūsu izmantošanas gadījums izvirzīja vairākas galvenās tehniskās prasības:

• Liela joslas platuma pārsūtīšana starp attāliem GPU no nesadalītiem atmiņas blokiem
• Iespēja dinamiiski pievienot vai noņemt mezglus no Kubernetes izvietojumiem, nepārtraucot notiekošās darbības
• Atbalsts vienādu (peer-to-peer) sakaru modeļiem

Lai gan NVIDIA NCCL bibliotēka ir de facto standarts izkliedētai dziļajai mācīšanai, tā nebija ideāli piemērota mūsu vajadzībām:

• NCCL izceļas kolektīvajā saziņā, taču prasa izveidot statisku "pasauli", kas savukārt prasa restartēt visu klasteri, kad tiek pielāgoti iesaistītie mezgli.
• NCCL sinhronās komunikācijas modelis rada sarežģītību mūsu asinhronajai darba slodzei
• Mēs vēlējāmies tiešu kontroli pār mūsu atmiņas pārsūtīšanas modeļiem optimizācijas nolūkos
• Pašu risinājuma izveide sniedza vērtīgas mācīšanās iespējas

Mūsdienu augstas veiktspējas tīkli

Lai izprastu mūsu risinājumu, vispirms aplūkosim, ar ko mūsdienu augstas veiktspējas tīkli atšķiras no tradicionālajiem tīkliem.

Lielākā daļa tīklu, ko izmantojam ikdienā, balstās uz TCP/IP protokoliem, kur lietojumprogrammas sazinās ar tīkla karti, izmantojot operētājsistēmas kodolu un soketus. Tomēr augstas veiktspējas tīklos tiek izmantots RDMA (Remote Direct Memory Access) — pilnīgi atšķirīga aparatūras un programmatūras kaudze, kas nodrošina tiešu atmiņas piekļuvi starp mašīnām bez procesora iesaistes.

AWS piedāvā Elastic Fabric Adapter (EFA) — pielāgotu tīkla interfeisu, kas ievieš Amazon pielāgoto protokolu ar nosaukumu Scalable Reliable Datagram (SRD). Atšķirībā no tradicionālajiem TCP/IP tīkliem, kur dati ir jākopē vairākas reizes starp lietotāja telpu, kodola telpu un tīkla buferiem, EFA ar RDMA nodrošina tiešu datu pārsūtīšanu starp GPU atmiņu un tīkla karti, pilnībā apejot procesoru.

Augstas veiktspējas tīkla projektēšanas filozofija

Augstas veiktspējas tīkla sistēmu izveide prasa pārskatīt vairākus pamata pieņēmumus:

1. Bufera īpašumtiesības: Atšķirībā no tradicionālajiem soketiem, kur kodols pārvalda tīkla buferus un prasa kopēšanu starp lietotāja telpu un kodola telpu, RDMA prasa lietojumprogrammām pašām pārvaldīt savus buferus. Kad lietojumprogramma uzsāk tīkla darbību, tā nodod bufera īpašumtiesības tīkla kartei, līdz darbība ir pabeigta, novēršot datu kopēšanas nepieciešamību.
2. Atmiņas reģistrācija: Lietojumprogrammām ir jāreģistrē atmiņas apgabali ar operētājsistēmas kodolu. Kodols iestatī virtualās adreses kartējumus, kas ļauj CPU, GPU un tīkla kartēm saprast vienas un tās pašas virtuālās adreses. Šī reģistrācija ir vienreizēja operācija, kas nodrošina turpmāku datu pārsūtīšanu bez kopēšanas.
3. Vadības plakne pret datu plakni: Augstas veiktspējas tīkli sadala operācijas divās kategorijās:
   1. Vadības plaknes operācijas (piemēram, savienojuma iestatīšana un atmiņas reģistrācija) iet caur kodolu, lai nodrošinātu drošību
   2. Datu plaknes operācijas (faktiskā datu pārsūtīšana) apved kodolu, lai nodrošinātu maksimālu veiktspēju
4. Saņemšana pirms pārsūtīšanas: Bez kodola pārvaldītiem buferiem lietojumprogrammām ir iepriekš jāpublicē saņemšanas operācijas, norādot, kur ienākošajiem datiem vajadzētu tikt novietotiem. Šī ir fundamentāla maiņa salīdzinājumā ar soketu modeli, kur lietojumprogrammas var saņemt datus jebkurā laikā.
5. Uz aptauju balstīta izpilde: Tā vietā, lai gaidītu kodola paziņojumus, izmantojot tādus mehānismus kā epoll, lietojumprogrammas tieši aptaujā aparatūras izpildes rindas. Tas novērš sistēmas izsaukuma pieskaitāmās izmaksas un ļauj nekavējoties reaģēt uz pabeigtajām operācijām.
6. Aparatūras topoloģijas apzināšanās: Izpratne un optimizācija aparatūras topoloģijai ir būtiska, lai sasniegtu maksimālu veiktspēju.

Aparatūras topoloģijas izpratne

AWS p5 instancēm ir sarežģīta iekšējā arhitektūra. Kā parādīts zemāk, katra instance satur divus CPU soketus, veidojot divus NUMA mezglus, un katrs NUMA mezgls savienojas ar četriem PCIe slēdziem:

Izvērst

Zem katra PCIe slēdža atrodas:

• Četras 100 Gbps EFA tīkla kartes
• Viens NVIDIA H100 GPU
• Viens NVMe SSD

Datu ceļi TCP/IP un RDMA pārsūtīšanai parāda fundamentālas atšķirības to pieejās:

Izvērst

Izmantojot TCP/IP (kreisajā pusē), dati ir jākopē vairākas reizes:

1. No GPU atmiņas uz lietotāja telpas lietojumprogrammas buferi
2. No lietotāja telpas bufera uz kodola buferi
3. No kodola bufera uz tīkla karti
4. Pretējs process notiek saņemšanas pusē

Katra kopēšanas operācija patērē procesora ciklus un atmiņas joslas platumu. Lietojumprogrammai ir arī jāpārslēdz konteksts starp lietotāja telpu un kodola telpu katrai tīkla darbībai.

Turpretim RDMA (labajā pusē) nodrošina patiesu datu pārsūtīšanu bez kopēšanas (zero-copy):

1. Tīkla karte nolasa datus tieši no GPU atmiņas
2. Dati ceļo tieši uz attālo tīkla karti
3. Attālā tīkla karte ieraksta datus tieši mērķa GPU atmiņā
4. Lietojumprogrammai ir jāpārbauda tikai izpildes rinda (CQ) lietotāja telpā, lai uzzinātu, kad pārsūtīšana ir pabeigta

Izmantojot RDMA un pareizu aparatūras savienošanu pārī, datu pārsūtīšanai starp diviem GPU ir jāšķērso tikai vietējais PCIe slēdzis un tīkls:

Izvērst

Turpretim TCP/IP pārsūtīšanai dati ir jākopē vairākas reizes caur galveno atmiņu, izraisot ievērojamu PCIe kopnes noslodzi:

Izvērst

Veidošana, izmantojot libfabric

Mēs izmantojām libfabric — ietvaru, kas nodrošina vispārēju interfeisu auduma pakalpojumiem. Mūsu ieviešanā tiek izmantoti divu veidu RDMA operācijas:

• Divpusējs RDMA (SEND/RECV) vadības ziņojumiem, kas satur metadatus par atmiņas apgabaliem (piemēram, nobīdes un izmērus).
• Vienpusējs RDMA WRITE faktiskajai datu pārsūtīšanai, kur katra WRITE operācija apstrādā vienu nepārtrauktu atmiņas bloku.

Ceļš uz maksimālo veiktspēju

Mūsu izstrāde noritēja vairākos posmos:

1. Ievviesta pamata virziena ziņojumu pārsūtīšana, izmantojot SEND/RECV
2. Paplašināts līdz divvirzienu komunikācijai
3. Pievienots GPUDirect RDMA WRITE tiešai GPU-GPU pārsūtīšanai
4. Paplašināts, lai apstrādātu vairākas vienlaicīgas pārsūtīšanas
5. Ieviesta operāciju rindas veidošana robustumam
6. Sasniegta 97,4% joslas platuma izmantošana vienā tīkla kartē

Mērogojot līdz 32 tīkla kartēm, mēs ieviesām vairākas būtiskas optimizācijas:

1. Operāciju rindu veidošana: Tā vietā, lai tieši iesniegtu operācijas tīkla kartēm, mēs uzturam lietojumprogrammas līmeņa rindu. Tas nodrošina noturību pret tīkla noslodzi un vienkāršo programmēšanas modeli.
2. Tīkla iesildīšana: Savienojumu iepriekšēja izveide uzlaboja palaišanas veiktspēju.
3. Daudzpavedienu apstrāde: Īpaši pavedieni katra GPU tīkla darbībām.
4. Procesora kodolu piesaiste: Pavedienu piesaistīšana konkrētiem procesora kodoliem, lai izvairītos no NUMA efektiem un kešatmiņas kļūmēm.
5. Stāvokļa sadalīšana: Sāncensības samazināšana starp pavedieniem.
6. Operāciju pakešu apstrāde: Vairāku operāciju iesniegšana kopā.
7. Tīša operāciju publicēšana: Operācijas vispirms tiek ievietotas rindā lietojumprogrammā. Pēc izpildes rindas aptaujāšanas mēs mēģinām iesniegt gaidošās operācijas tīkla karttei, nodrošinot efektīvu tīkla resursu izmantošanu.
8. NUMA apzināta resursu sadale: Allocate libfabric resources like completion queues on the correct NUMA node to minimize memory access latency.

Pateicoties šīm optimizācijām, mēs sasniedzām galējo veiktspēju 3108 Gbps — 97,1% no teorētiskā maksimālā joslas platuma.

Zemāk redzamajā videoklipā ir parādīta mūsu komandrindas programma darbībā. Tā pārsūta nesadalītas GPU atmiņas daļas uz attālu mezglu, sasniedzot pārraides ātrumu 3108,283 Gbps, kas demonstrē gandrīz teorētisku tīkla infrastruktūras joslas platuma izmantošanu:

Secinājums

Augstas veiktspējas tīkla sistēmas izveide prasa izpratni gan par aparatūras arhitektūru, gan sistēmas projektēšanas principiem. Lai gan bibliotēkas, piemēram, NCCL, piedāvā lieliskus risinājumus izplatītiem modeļiem, dažkārt konkrētām prasībām ir nepieciešami pielāgoti risinājumi.

Mūsu ceļojums parāda, ka, rūpīgi pievēršot uzmanību sistēmas arhitektūrai, aparatūras topoloģijai un dažādām optimizācijas tehnikām, ir iespējams sasniegt gandrīz teorētisku tīkla veiktspēju. Galvenais ir izprast ne tikai atsevišķus komponentus, bet arī to, kā tie mijiedarbojas, veidojot pilnīgu sistēmu.

Pilnīgu tehnisko apskatu par šo ceļojumu, ieskaitot ieviešanas detaļas un koda piemērus, var atrast mūsu atvērtā koda krātuvē.

Citāts

@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}
}