Het artikel bespreekt RapidsMPF, een bibliotheek die is ontwikkeld om de problemen rondom 'shuffling' bij grootschalige data-analyse (zoals joins, groupby en merge) op te lossen. Shuffling is vaak de bottleneck in data-pipelines omdat het extreem geheugenintensief is en vaak leidt tot Out-Of-Memory (OOM) fouten.
De kernoplossing van RapidsMPF is een out-of-core benadering. Wanneer het VRAM van de GPU onvoldoende is, kan de bibliotheek data 'spillen' naar het hostgeheugen. Dankzij de integratie met UCXX kan RapidsMPF buffers direct op de host ontvangen, waardoor onnodige bewegingen van data tussen device en host (thrashing) worden voorkomen.
Uit benchmarks op een DGX B200 blijkt dat RapidsMPF een globale doorvoersnelheid van 1,8 TiB/s kan behalen onder optimale omstandigheden. Zelfs bij extreme geheugenbeperkingen blijft het systeem stabiel; de prestaties degraderen geleidelijk in plaats van dat het proces crasht, waarbij de PCIe-bus de uiteindelijke bottleneck vormt.
Out-Of-Core Shuffling met RapidsMPF
Shuffling is de kern van gestructureerde data-analyse, of deze nu gedistribueerd is of niet. Het is een cruciaal onderdeel van belangrijke data-operaties zoals: join, groupby, merge, sort, etc. Een volledige gedistribueerde shuffle kan alle data van elk proces naar elk ander proces verplaatsen, een zogenaamde 'all-to-all'. Dit is zeer kostbaar en er zijn veel geavanceerde technieken ontwikkeld om deze operatie zoveel mogelijk te vermijden. Shuffles zijn rekentechnisch niet bijzonder uitdagend: het berekenen van de hashes om data te routeren is relatief goedkoop. Ze zijn echter wel duur in een workflow om verschillende redenen:
- Geheugenintensief: Shuffles kunnen vereisen dat er een volledige kopie van alle data wordt vastgehouden of, in streaming-gevallen, kan de geheugendruk oplopen en OOM-fouten veroorzaken.
- Transport: De data moet fysiek worden verplaatst van Proces A naar Proces B of van Node A naar Node B, waardoor de snelheid beperkt wordt door de transportslaag.
- Synchronisatie: Output-data kan pas worden geconsumeerd nadat alle producenten zijn gestopt met het bijdragen van data. In een bulk-synchrone engine kan deze barrière het volledige uitvoeringsplan vertragen.
Omdat shuffling moeilijk, traag, geheugenintensief en kritiek is, is dit het punt waar RapidsMPF is gestart.
Waarom hebben Joins Shuffles nodig?
Een korte inleiding over het joinen van tabellen. Stel dat we twee tabellen hebben: partsupp en lineitem, en we willen deze joinen:
partsupp.join(
lineitem,
left_on=["ps_partkey", "ps_suppkey"],
right_on=["l_partkey", "l_suppkey"],
)
Of in SQL:
SELECT *
FROM partsupp
JOIN lineitem
ON partsupp.ps_partkey = lineitem.l_partkey
AND partsupp.ps_suppkey = lineitem.l_suppkey
In-Memory Joins
Inner joins bestaan uit twee fasen:
- Build-fase: De kleinere tabel (
partsupp) wordt gescand en er wordt een hashtabel geconstrueerd over de join-keys (pspartkey, pssuppkey), waarbij elke gehashte key wordt gekoppeld aan de rij waar deze vandaan komt. Deze hashtabel moet volledig zijn ingevuld voordat de probe-fase kan beginnen.
- Probe-fase: De grotere tabel (
lineitem) wordt gescand en de keys van elke rij (lpartkey, lsuppkey) worden gehasht. Matches tussen de hash van de build- en probe-tabel join-keys produceren een output-rij die kolommen van beide tabellen combineert. Rijen zonder matches worden weggegooid.
Opmerking: left, right en full outer joins gebruiken dezelfde build/probe-strategie, maar met verschillende regels voor rijden zonder match.
Minimaal houdt een in-memory join drie tabellen vast: de build-tabel, de probe-tabel en de output-tabel, plus de hashtabel die over de build-zijde is gebouwd.
Gedistribueerde Join
In het in-memory geval bevindt alle data zich al in dezelfde geheugenruimte. Dat is niet langer het geval zodra de tabellen verspreid zijn over vele processen, nodes of ranks, of wanneer tabellen in batches worden verwerkt voor "streaming" joins. Een rank/proces kan alleen rijen joinen die in het resident geheugen aanwezig zijn. Uiteindelijk zal er een in-memory join plaatsvinden, maar eerst moeten we alle matching keys voor de build- en probe-tabellen op dezelfde rank krijgen.
Om een gedistribueerde hash-join uit te voeren, moet men het volgende doen:
- Scan de build-tabel, hash de join-keys van elke rij om een bestemmingspartitie te kiezen:
hash(keys) % noutpartitions. Pak en verstuur elke rij naar de rank die deze zal beheren.
- Scan de probe-tabel en routeer de rijen op dezelfde manier, zodat probe-keys terechtkomen op de rank die al de build-keys met dezelfde hash bezit.
- Wacht tot elke rank klaar is met verzenden. Pas dan is een rank gegarandeerd in het bezit van elke rij uit beide tabellen voor de keys die deze beheert.
- Voer de in-memory join uit op de lokale slice van elke rank: bouw een hashtabel over de build-rijen, probe deze met de probe-rijen en produceer matches.
In het slechtste geval, als elke fase volledig gematerialiseerd is voordat de volgende begint, houdt een enkele rank al het volgende tegelijkertijd vast:
- build-tabel (bronscan)
- probe-tabel (bronscan)
- gestaged build-tabel (gepakt voor verzending)
- gestaged probe-tabel (gepakt voor verzending)
- geshuffelde build-slice (ontvangen)
- geshuffelde probe-slice (ontvangen)
- hashtabel over de build-slice
- output-tabel
Dit is waarom shuffling geheugenintensief is in plaats van rekenintensief. Het hashen zelf is goedkoop, en dat is waarom een out-of-core shuffle-implementatie een primair focuspunt moet zijn bij het bouwen van een ETL-engine. Daarnaast is het kunnen streamen van data, in plaats van het volledig materialiseren van tabellen vóór het shuffelen, kritiek voor het verminderen van de geheugendruk. Om deze redenen is RapidsMPF gestart met het oorspronkelijke doel om een streaming out-of-core shuffler te bouwen.
RapidsMPF
RapidsMPF is sinds het begin uitgebreid. Het is nu een bibliotheek die bestaat uit twee grote onderdelen:
- Een shuffle-bibliotheek ontworpen voor spilling / out-of-core geheugenbeheer, met versneld transport.
- Een actor-netwerk voor het construeren van streaming data-pipelines.
Gebruikers kunnen tegenwoordig alleen de shuffling-component van RapidsMPF adopteren (via C++ of Python interfaces). We hebben deze adoptie gezien in NeMo-Curator en, experimenteel, in Ray Data. Belangrijker is dat cuDF Polars RapidsMPF gebruikt voor zowel shuffles als het actor-netwerk.
Onze shuffle-implementatie moet aan de volgende eisen voldoen:
- Snel zijn
- Schaalbaar zijn
- Werken met data die groter is dan het VRAM (GPU) (out-of-core)
- Herbruikbaar zijn
In de rest van dit artikel ligt de focus op shuffling onder geheugendruk.
Benchmarking Setup
cuDF/RapidsMPF heeft een gebruiksvriendelijke C++ benchmark: bench_shuffle. Deze helpt ons te bestuderen hoe de RapidsMPF shuffling-implementatie werkt op diverse hardware (transports, aantal GPU's, etc.) en configuraties (input/output partitiegroottes, geheugenresources, etc.).
Hieronder volgt een overzicht van wat de huidige bench_shuffle test aanbiedt, samen met de waarden die in dit artikel worden gebruikt. In het algemeen bouwt deze benchmark instelbare hoeveelheden willekeurige 32-bit (4 byte) integers per rank (per GPU), shuffelt alle data en rondt dit af (er is hier geen join, alleen de shuffle).
| Vlag | Betekenis | Gebruikte waarden |
|---|
-C <naam> | Communicator | ucxx |
-c <n> | Aantal kolommen | 10 |
-r <n> | Aantal getimede runs | 10 |
-w <n> | Aantal warmup runs | 3 |
-n <n> | Aantal rijen per rank | 536870912 (2 GiB per kolom bij 4 bytes/rij) |
-p <n> | Aantal input-partities per rank | 1 |
-o <n> | Aantal output-partities per rank | 8 (één per rank) |
-m <naam> | RMM memory resource | pool |
-l <n> | Limiet apparaatgeheugen in MiB | weggelaten = onbeperkt; 32768 tot 12288 in de spill-sweep |
-s | Activeer output discard (simuleer streaming) | vlag, altijd ingesteld |
-x | Activeer geheugenprofiling | vlag, altijd ingesteld |
-g | Gebruik vooraf gepartitioneerde input-tabellen | vlag, altijd ingesteld |
Voor alle tests gebruiken we een enkele DGXB200 en rrun (een MPI-achtige multiprocess launch tool) om de shuffles te starten.
Eenvoudige Shuffling
Een DGXB200 heeft 8 Blackwell GPU's, elk met 180GB VRAM, en 2 Intel® Xeon® Platinum 8570 Processors. Om een baseline te krijgen, beginnen we met het shuffelen van data die comfortabel in alle GPU's past.
rrun -n 8 --bind-to cpu --bind-to memory -x UCX_MAX_RNDV_RAILS=1 -x UCX_PROTO_ENABLE=y -x UCX_WARN_UNUSED_ENV_VARS=n libcudf_streaming_bench_shuffle -C ucxx -w 3 -r 10 -m pool -g -s -x -p 1 -o 8 -c 10 -n 536870912
Berekening van de data:
536.870.912 rijen * 4 bytes (32-bit ints) = 2 GiB per kolom10 kolommen * 2 GiB = 20 GiB / rank8 ranks * 20 GiB = 160 GiB totaal
Voorbeeld output voor 20GiB/rank:
[0:PRINT:0:2026-09-16 02:09:54.328272061] elapsed: 86.89 ms | local throughput: 230.19 GiB/s | global throughput: 1.80 TiB/s
Elke rank rapporteert de tijd besteed aan shuffling, en de lokale en globale doorvoersnelheid. We kunnen al waarnemen dat de warmup-fase trager is dan de officiële run. Aan het einde geeft het programma de gemiddelde waarden per rank en statistieken over waar de tijd is besteed (geheugenallocatie, spilling, etc.).
Opmerking: De globale doorvoersnelheid varieert per rank, wat technisch gezien onjuist is en een rapportagefout is. De benchmark rapporteert de globale throughput als de lokale throughput van elke rank vermenigvuldigd met het aantal ranks (8x op een DGXB200).
Geheugenprofiling
De volgende gegevens zijn voor rank 0:
| ncalls | peak | g-peak | accum | max | filename:lineno(name) |
|---|
| 1 | 60 GiB | 60 GiB | 1.29 TiB | 2 GiB | main (all allocations using RmmResourceAdaptor) |
| 1 | 40 GiB | 40 GiB | 44.06 GiB | 2 GiB | /libcudfstreaming/src/partitionutils.cpp:183(partitionandpack) |
| 1 | 20 GiB | 20 GiB | 20 GiB | 320.39 MiB | /libcudfstreaming/src/partitionutils.cpp:242(splitandpack) |
| 8 | 2.50 GiB | 2.50 GiB | 20 GiB | 256.12 MiB | /libcudfstreaming/src/partitionutils.cpp:298(unpackandconcat) |
Met 1.440 GB VRAM over acht GPU's heeft RapidsMPF ruim voldoende plek om te shuffelen zonder te hoeven 'spillen'. In deze configuratie haalt RapidsMPF ongeveer 1,8 TiB/s globale doorvoersnelheid.
Het geheugenprofiel is interessant: elke rank heeft slechts 20 GiB input, maar het piekgebruik is 60 GiB (3x de input). Dit komt doordat er 20 GiB wordt gebruikt voor de input zelf en 40 GiB in partitionandpack, terwijl de input wordt gehasht en gekopieerd naar buffers per bestemming. In een shuffle kan het programma tijdelijk zowel het origineel als de kopie van de lokale data bezitten.
"Oops, you spilled a little..."
We gaan nu de geheugendruk verhogen door de apparaatlimiet te verlagen van onbeperkt naar 12 GiB. We zullen observeren dat RapidsMPF niet alleen geen OOM-fouten geeft, maar ook niet 'thrasht' (onnodig data heen en weer verplaatsen tussen device en host).
Laten we een kunstmatige geheugenlimiet van 32GB toepassen (-l 32768):
Statistieken (laatste run):
alloc-device: 17.50 GiB | 8.11 ms | 2.11 TiB/salloc-pinned_host: 5.62 GiB | 844 us | 6.51 TiB/scopy-pinned_host-to-device: 5.62 GiB | 171.80 ms | 32.74 GiB/s
Interessant! De globale doorvoersnelheid is aanzienlijk gedaald naar ~480 GiB/s. We zien nieuwe regels voor alloc-pinnedhost en copy-pinnedhost-to-device, maar we zien geen copy-device-to-pinned_host.
Dit is het cruciale punt: er is niets van het apparaat verwijderd (evicted). In plaats daarvan kan de ontvangende zijde de geheugendruk waarnemen en buffers direct op de host accepteren in plaats van op het apparaat, om thrashing te voorkomen. Dit is mogelijk dankzij UCXX. Spilling die nooit hoeft te spillen is de goedkoopste vorm van spilling.
Waarom de prestatiehit? Van de 335,15 ms besteedt RapidsMPF 171,8 ms aan het verplaatsen van data terug naar het apparaat. De DGXB200 heeft PCIe Gen 5 (64 GB/s voor 16 lanes). RapidsMPF verplaatst data over deze bus met een gemiddelde snelheid van 32 GiB/s per GPU. Dit is de bottleneck. Met een Grace-Blackwell configuratie met C2C (tot 900 GB/s) zou deze bandbreedte tussen host en apparaat 5x tot 10x kunnen toenemen.
Scaled Spilling DGXB200
Onderstaande tabel toont wat er gebeurt wanneer de geheugenlimiet stapsgewijs wordt verlaagd:
| Label | Spill limiet | Piek device | Lokale perf | Globale throughput | D2H (pinned) vol / tijd / rate | H2D (pinned) vol / tijd / rate |
|---|
| no-spill | $\infty$ GiB | 60.0 GiB | 229.7 GiB/s | 1.79 TiB/s | 0 | 0 |
| onset-spill | 32 GiB | 60.0 GiB | 60.0 GiB/s | 480.2 GiB/s | 0 | 5.6 GiB / 191.4 ms / 29.4 GiB/s |
| light-spill | 28 GiB | 60.0 GiB | 35.9 GiB/s | 287.0 GiB/s | 0 | 9.7 GiB / 379.0 ms / 25.6 GiB/s |
| moderate-spill | 24 GiB | 60.0 GiB | 26.5 GiB/s | 211.7 GiB/s | 0 | 13.8 GiB / 491.0 ms / 28.0 GiB/s |
| heavy-spill | 20 GiB | 60.0 GiB | 14.8 GiB/s | 118.6 GiB/s | 0.3 GiB / 9.6 ms / 32.6 GiB/s | 17.5 GiB / 622.3 ms / 28.1 GiB/s |
| very-heavy-spill | 16 GiB | 60.0 GiB | 10.1 GiB/s | 80.5 GiB/s | 4.1 GiB / 128.7 ms / 31.6 GiB/s | 18.1 GiB / 1340.0 ms / 13.5 GiB/s |
| extreme-spill | 12 GiB | 60.0 GiB | 6.9 GiB/s | 55.5 GiB/s | 8.1 GiB / 264.5 ms / 30.7 GiB/s | 18.8 GiB / 1170.0 ms / 16.1 GiB/s |
Observaties uit de tabel:
- Geleidelijke degradatie: De prestaties dalen monotoon naarmate de limiet lager wordt, maar het systeem loopt niet vast (geen OOM).
- Efficiënte ontvangst: De
copy-device-to-pinned_host is 0 voor de eerste vier spill-niveaus. RapidsMPF ontvangt buffers direct op de host om de limiet te respecteren. Pas bij een limiet van 20 GiB (gelijk aan de inputgrootte) treedt echte device-to-host spilling op.
- De kosten van spilling: De tijd om data terug te verplaatsen naar het apparaat (
copy-pinned_host-to-device) groeit van 191 ms naar 1170 ms.
Samenvatting
We hebben aangetoond hoe uitdagend en geheugenintensief shuffles (en de daarmee samenhangende joins en sorts) kunnen zijn. Zonder beperkingen haalt RapidsMPF ongeveer 1,8 TiB/s globale doorvoersnelheid op een DGX B200. Belangrijker is dat RapidsMPF succesvol kan shuffelen zonder onnodige thrashing of OOM-fouten.
Terugblikkend op de uitdagingen van shuffles:
- Geheugenintensief: RapidsMPF kan buffers van het apparaat spillen/evicten bij instelbare limieten én kan bytes op de host ontvangen om onnodige thrashing te voorkomen.
- Transport: RapidsMPF maakt gebruik van UCXX, wat transport via NVLink, InfiniBand, EFA, TCP, etc. mogelijk maakt.
- Synchronisatie: RapidsMPF lost dit op met een asynchroon/streaming uitvoeringsmodel dat de shuffle overlapt met ander klaarstaand werk.
We wilden een shuffle die snel is, schaalt en data groter dan het VRAM aankan, en dat hebben we met RapidsMPF gerealiseerd.
Out-Of-Core Shuffling met RapidsMPF
Shuffling is de kern van gestructureerde data-analyse, of deze nu gedistribueerd is of niet. Het is een cruciaal onderdeel van belangrijke data-operaties zoals: join, groupby, merge, sort, etc. Een volledige gedistribueerde shuffle kan alle data van elk proces naar elk ander proces verplaatsen, een zogenaamde 'all-to-all'. Dit is zeer kostbaar en er zijn veel geavanceerde technieken ontwikkeld om deze operatie zoveel mogelijk te vermijden. Shuffles zijn rekentechnisch niet bijzonder uitdagend: het berekenen van de hashes om data te routeren is relatief goedkoop. Ze zijn echter wel duur in een workflow om verschillende redenen:
- Geheugenintensief: Shuffles kunnen vereisen dat er een volledige kopie van alle data wordt vastgehouden of, in streaming-gevallen, kan de geheugendruk oplopen en OOM-fouten veroorzaken.
- Transport: De data moet fysiek worden verplaatst van Proces A naar Proces B of van Node A naar Node B, waardoor de snelheid beperkt wordt door de transportslaag.
- Synchronisatie: Output-data kan pas worden geconsumeerd nadat alle producenten zijn gestopt met het bijdragen van data. In een bulk-synchrone engine kan deze barrière het volledige uitvoeringsplan vertragen.
Omdat shuffling moeilijk, traag, geheugenintensief en kritiek is, is dit het punt waar RapidsMPF is gestart.
Waarom hebben Joins Shuffles nodig?
Een korte inleiding over het joinen van tabellen. Stel dat we twee tabellen hebben: partsupp en lineitem, en we willen deze joinen:
partsupp.join(
lineitem,
left_on=["ps_partkey", "ps_suppkey"],
right_on=["l_partkey", "l_suppkey"],
)
Of in SQL:
SELECT *
FROM partsupp
JOIN lineitem
ON partsupp.ps_partkey = lineitem.l_partkey
AND partsupp.ps_suppkey = lineitem.l_suppkey
In-Memory Joins
Inner joins bestaan uit twee fasen:
- Build-fase: De kleinere tabel (
partsupp) wordt gescand en er wordt een hashtabel geconstrueerd over de join-keys (pspartkey, pssuppkey), waarbij elke gehashte key wordt gekoppeld aan de rij waar deze vandaan komt. Deze hashtabel moet volledig zijn ingevuld voordat de probe-fase kan beginnen.
- Probe-fase: De grotere tabel (
lineitem) wordt gescand en de keys van elke rij (lpartkey, lsuppkey) worden gehasht. Matches tussen de hash van de build- en probe-tabel join-keys produceren een output-rij die kolommen van beide tabellen combineert. Rijen zonder matches worden weggegooid.
Opmerking: left, right en full outer joins gebruiken dezelfde build/probe-strategie, maar met verschillende regels voor rijden zonder match.
Minimaal houdt een in-memory join drie tabellen vast: de build-tabel, de probe-tabel en de output-tabel, plus de hashtabel die over de build-zijde is gebouwd.
Gedistribueerde Join
In het in-memory geval bevindt alle data zich al in dezelfde geheugenruimte. Dat is niet langer het geval zodra de tabellen verspreid zijn over vele processen, nodes of ranks, of wanneer tabellen in batches worden verwerkt voor "streaming" joins. Een rank/proces kan alleen rijen joinen die in het resident geheugen aanwezig zijn. Uiteindelijk zal er een in-memory join plaatsvinden, maar eerst moeten we alle matching keys voor de build- en probe-tabellen op dezelfde rank krijgen.
Om een gedistribueerde hash-join uit te voeren, moet men het volgende doen:
- Scan de build-tabel, hash de join-keys van elke rij om een bestemmingspartitie te kiezen:
hash(keys) % noutpartitions. Pak en verstuur elke rij naar de rank die deze zal beheren.
- Scan de probe-tabel en routeer de rijen op dezelfde manier, zodat probe-keys terechtkomen op de rank die al de build-keys met dezelfde hash bezit.
- Wacht tot elke rank klaar is met verzenden. Pas dan is een rank gegarandeerd in het bezit van elke rij uit beide tabellen voor de keys die deze beheert.
- Voer de in-memory join uit op de lokale slice van elke rank: bouw een hashtabel over de build-rijen, probe deze met de probe-rijen en produceer matches.
In het slechtste geval, als elke fase volledig gematerialiseerd is voordat de volgende begint, houdt een enkele rank al het volgende tegelijkertijd vast:
- build-tabel (bronscan)
- probe-tabel (bronscan)
- gestaged build-tabel (gepakt voor verzending)
- gestaged probe-tabel (gepakt voor verzending)
- geshuffelde build-slice (ontvangen)
- geshuffelde probe-slice (ontvangen)
- hashtabel over de build-slice
- output-tabel
Dit is waarom shuffling geheugenintensief is in plaats van rekenintensief. Het hashen zelf is goedkoop, en dat is waarom een out-of-core shuffle-implementatie een primair focuspunt moet zijn bij het bouwen van een ETL-engine. Daarnaast is het kunnen streamen van data, in plaats van het volledig materialiseren van tabellen vóór het shuffelen, kritiek voor het verminderen van de geheugendruk. Om deze redenen is RapidsMPF gestart met het oorspronkelijke doel om een streaming out-of-core shuffler te bouwen.
RapidsMPF
RapidsMPF is sinds het begin uitgebreid. Het is nu een bibliotheek die bestaat uit twee grote onderdelen:
- Een shuffle-bibliotheek ontworpen voor spilling / out-of-core geheugenbeheer, met versneld transport.
- Een actor-netwerk voor het construeren van streaming data-pipelines.
Gebruikers kunnen tegenwoordig alleen de shuffling-component van RapidsMPF adopteren (via C++ of Python interfaces). We hebben deze adoptie gezien in NeMo-Curator en, experimenteel, in Ray Data. Belangrijker is dat cuDF Polars RapidsMPF gebruikt voor zowel shuffles als het actor-netwerk.
Onze shuffle-implementatie moet aan de volgende eisen voldoen:
- Snel zijn
- Schaalbaar zijn
- Werken met data die groter is dan het VRAM (GPU) (out-of-core)
- Herbruikbaar zijn
In de rest van dit artikel ligt de focus op shuffling onder geheugendruk.
Benchmarking Setup
cuDF/RapidsMPF heeft een gebruiksvriendelijke C++ benchmark: bench_shuffle. Deze helpt ons te bestuderen hoe de RapidsMPF shuffling-implementatie werkt op diverse hardware (transports, aantal GPU's, etc.) en configuraties (input/output partitiegroottes, geheugenresources, etc.).
Hieronder volgt een overzicht van wat de huidige bench_shuffle test aanbiedt, samen met de waarden die in dit artikel worden gebruikt. In het algemeen bouwt deze benchmark instelbare hoeveelheden willekeurige 32-bit (4 byte) integers per rank (per GPU), shuffelt alle data en rondt dit af (er is hier geen join, alleen de shuffle).
| Vlag | Betekenis | Gebruikte waarden |
|---|
-C <naam> | Communicator | ucxx |
-c <n> | Aantal kolommen | 10 |
-r <n> | Aantal getimede runs | 10 |
-w <n> | Aantal warmup runs | 3 |
-n <n> | Aantal rijen per rank | 536870912 (2 GiB per kolom bij 4 bytes/rij) |
-p <n> | Aantal input-partities per rank | 1 |
-o <n> | Aantal output-partities per rank | 8 (één per rank) |
-m <naam> | RMM memory resource | pool |
-l <n> | Limiet apparaatgeheugen in MiB | weggelaten = onbeperkt; 32768 tot 12288 in de spill-sweep |
-s | Activeer output discard (simuleer streaming) | vlag, altijd ingesteld |
-x | Activeer geheugenprofiling | vlag, altijd ingesteld |
-g | Gebruik vooraf gepartitioneerde input-tabellen | vlag, altijd ingesteld |
Voor alle tests gebruiken we een enkele DGXB200 en rrun (een MPI-achtige multiprocess launch tool) om de shuffles te starten.
Eenvoudige Shuffling
Een DGXB200 heeft 8 Blackwell GPU's, elk met 180GB VRAM, en 2 Intel® Xeon® Platinum 8570 Processors. Om een baseline te krijgen, beginnen we met het shuffelen van data die comfortabel in alle GPU's past.
rrun -n 8 --bind-to cpu --bind-to memory -x UCX_MAX_RNDV_RAILS=1 -x UCX_PROTO_ENABLE=y -x UCX_WARN_UNUSED_ENV_VARS=n libcudf_streaming_bench_shuffle -C ucxx -w 3 -r 10 -m pool -g -s -x -p 1 -o 8 -c 10 -n 536870912
Berekening van de data:
536.870.912 rijen * 4 bytes (32-bit ints) = 2 GiB per kolom10 kolommen * 2 GiB = 20 GiB / rank8 ranks * 20 GiB = 160 GiB totaal
Voorbeeld output voor 20GiB/rank:
[0:PRINT:0:2026-09-16 02:09:54.328272061] elapsed: 86.89 ms | local throughput: 230.19 GiB/s | global throughput: 1.80 TiB/s
Elke rank rapporteert de tijd besteed aan shuffling, en de lokale en globale doorvoersnelheid. We kunnen al waarnemen dat de warmup-fase trager is dan de officiële run. Aan het einde geeft het programma de gemiddelde waarden per rank en statistieken over waar de tijd is besteed (geheugenallocatie, spilling, etc.).
Opmerking: De globale doorvoersnelheid varieert per rank, wat technisch gezien onjuist is en een rapportagefout is. De benchmark rapporteert de globale throughput als de lokale throughput van elke rank vermenigvuldigd met het aantal ranks (8x op een DGXB200).
Geheugenprofiling
De volgende gegevens zijn voor rank 0:
| ncalls | peak | g-peak | accum | max | filename:lineno(name) |
|---|
| 1 | 60 GiB | 60 GiB | 1.29 TiB | 2 GiB | main (all allocations using RmmResourceAdaptor) |
| 1 | 40 GiB | 40 GiB | 44.06 GiB | 2 GiB | /libcudfstreaming/src/partitionutils.cpp:183(partitionandpack) |
| 1 | 20 GiB | 20 GiB | 20 GiB | 320.39 MiB | /libcudfstreaming/src/partitionutils.cpp:242(splitandpack) |
| 8 | 2.50 GiB | 2.50 GiB | 20 GiB | 256.12 MiB | /libcudfstreaming/src/partitionutils.cpp:298(unpackandconcat) |
Met 1.440 GB VRAM over acht GPU's heeft RapidsMPF ruim voldoende plek om te shuffelen zonder te hoeven 'spillen'. In deze configuratie haalt RapidsMPF ongeveer 1,8 TiB/s globale doorvoersnelheid.
Het geheugenprofiel is interessant: elke rank heeft slechts 20 GiB input, maar het piekgebruik is 60 GiB (3x de input). Dit komt doordat er 20 GiB wordt gebruikt voor de input zelf en 40 GiB in partitionandpack, terwijl de input wordt gehasht en gekopieerd naar buffers per bestemming. In een shuffle kan het programma tijdelijk zowel het origineel als de kopie van de lokale data bezitten.
"Oops, you spilled a little..."
We gaan nu de geheugendruk verhogen door de apparaatlimiet te verlagen van onbeperkt naar 12 GiB. We zullen observeren dat RapidsMPF niet alleen geen OOM-fouten geeft, maar ook niet 'thrasht' (onnodig data heen en weer verplaatsen tussen device en host).
Laten we een kunstmatige geheugenlimiet van 32GB toepassen (-l 32768):
Statistieken (laatste run):
alloc-device: 17.50 GiB | 8.11 ms | 2.11 TiB/salloc-pinned_host: 5.62 GiB | 844 us | 6.51 TiB/scopy-pinned_host-to-device: 5.62 GiB | 171.80 ms | 32.74 GiB/s
Interessant! De globale doorvoersnelheid is aanzienlijk gedaald naar ~480 GiB/s. We zien nieuwe regels voor alloc-pinnedhost en copy-pinnedhost-to-device, maar we zien geen copy-device-to-pinned_host.
Dit is het cruciale punt: er is niets van het apparaat verwijderd (evicted). In plaats daarvan kan de ontvangende zijde de geheugendruk waarnemen en buffers direct op de host accepteren in plaats van op het apparaat, om thrashing te voorkomen. Dit is mogelijk dankzij UCXX. Spilling die nooit hoeft te spillen is de goedkoopste vorm van spilling.
Waarom de prestatiehit? Van de 335,15 ms besteedt RapidsMPF 171,8 ms aan het verplaatsen van data terug naar het apparaat. De DGXB200 heeft PCIe Gen 5 (64 GB/s voor 16 lanes). RapidsMPF verplaatst data over deze bus met een gemiddelde snelheid van 32 GiB/s per GPU. Dit is de bottleneck. Met een Grace-Blackwell configuratie met C2C (tot 900 GB/s) zou deze bandbreedte tussen host en apparaat 5x tot 10x kunnen toenemen.
Scaled Spilling DGXB200
Onderstaande tabel toont wat er gebeurt wanneer de geheugenlimiet stapsgewijs wordt verlaagd:
| Label | Spill limiet | Piek device | Lokale perf | Globale throughput | D2H (pinned) vol / tijd / rate | H2D (pinned) vol / tijd / rate |
|---|
| no-spill | $\infty$ GiB | 60.0 GiB | 229.7 GiB/s | 1.79 TiB/s | 0 | 0 |
| onset-spill | 32 GiB | 60.0 GiB | 60.0 GiB/s | 480.2 GiB/s | 0 | 5.6 GiB / 191.4 ms / 29.4 GiB/s |
| light-spill | 28 GiB | 60.0 GiB | 35.9 GiB/s | 287.0 GiB/s | 0 | 9.7 GiB / 379.0 ms / 25.6 GiB/s |
| moderate-spill | 24 GiB | 60.0 GiB | 26.5 GiB/s | 211.7 GiB/s | 0 | 13.8 GiB / 491.0 ms / 28.0 GiB/s |
| heavy-spill | 20 GiB | 60.0 GiB | 14.8 GiB/s | 118.6 GiB/s | 0.3 GiB / 9.6 ms / 32.6 GiB/s | 17.5 GiB / 622.3 ms / 28.1 GiB/s |
| very-heavy-spill | 16 GiB | 60.0 GiB | 10.1 GiB/s | 80.5 GiB/s | 4.1 GiB / 128.7 ms / 31.6 GiB/s | 18.1 GiB / 1340.0 ms / 13.5 GiB/s |
| extreme-spill | 12 GiB | 60.0 GiB | 6.9 GiB/s | 55.5 GiB/s | 8.1 GiB / 264.5 ms / 30.7 GiB/s | 18.8 GiB / 1170.0 ms / 16.1 GiB/s |
Observaties uit de tabel:
- Geleidelijke degradatie: De prestaties dalen monotoon naarmate de limiet lager wordt, maar het systeem loopt niet vast (geen OOM).
- Efficiënte ontvangst: De
copy-device-to-pinned_host is 0 voor de eerste vier spill-niveaus. RapidsMPF ontvangt buffers direct op de host om de limiet te respecteren. Pas bij een limiet van 20 GiB (gelijk aan de inputgrootte) treedt echte device-to-host spilling op.
- De kosten van spilling: De tijd om data terug te verplaatsen naar het apparaat (
copy-pinned_host-to-device) groeit van 191 ms naar 1170 ms.
Samenvatting
We hebben aangetoond hoe uitdagend en geheugenintensief shuffles (en de daarmee samenhangende joins en sorts) kunnen zijn. Zonder beperkingen haalt RapidsMPF ongeveer 1,8 TiB/s globale doorvoersnelheid op een DGX B200. Belangrijker is dat RapidsMPF succesvol kan shuffelen zonder onnodige thrashing of OOM-fouten.
Terugblikkend op de uitdagingen van shuffles:
- Geheugenintensief: RapidsMPF kan buffers van het apparaat spillen/evicten bij instelbare limieten én kan bytes op de host ontvangen om onnodige thrashing te voorkomen.
- Transport: RapidsMPF maakt gebruik van UCXX, wat transport via NVLink, InfiniBand, EFA, TCP, etc. mogelijk maakt.
- Synchronisatie: RapidsMPF lost dit op met een asynchroon/streaming uitvoeringsmodel dat de shuffle overlapt met ander klaarstaand werk.
We wilden een shuffle die snel is, schaalt en data groter dan het VRAM aankan, en dat hebben we met RapidsMPF gerealiseerd.
Groetjes,