Accelerated Out of Core Shuffling
2 days ago
- RapidsMPF is a reusable, out-of-core shuffler designed to prevent OOM errors during large-scale data shuffling, achieving up to 1.8 TiB/s global throughput on a DGX B200.
- Shuffling is essential for distributed joins, groupbys, merges, and sorts, but is expensive due to memory intensity, data transport costs, and synchronization barriers.
- Distributed joins require routing matching keys to the same rank, which can force a single rank to hold many copies of data: source tables, staged buffers, received slices, hash tables, and output.
- RapidsMPF includes both a shuffle library (C++/Python) and an actor network for streaming pipelines; it is used by cuDF Polars, NeMo-Curator, and experimentally in Ray Data.
- In benchmark tests with 8 GPUs and 20 GiB per rank, the unconstrained shuffle ran at ~1.8 TiB/s global throughput with a peak device memory usage of 60 GiB per rank, 3x the input size.
- When device memory is limited to 32 GiB, throughput drops to ~480 GiB/s, but RapidsMPF avoids OOM by receiving incoming buffers directly on the host, preventing pointless device-to-host eviction.
- Sweeping memory limits from 32 GiB down to 12 GiB produces smooth, monotonic performance degradation without OOM; real spilling only starts once the limit approaches the input size.
- The main bottleneck during spills is host-to-device transfer over PCIe Gen 5, averaging ~32 GiB/s; Grace-Blackwell C2C could provide 5-10x higher bandwidth.
- RapidsMPF addresses memory pressure through tunable spilling and host-side buffer reception, transport via UCXX (NVLink, InfiniBand, EFA, TCP), and synchronization through asynchronous/streaming execution.
- The project demonstrates that an out-of-core shuffler can scale, avoid thrashing, and handle data larger than VRAM, making it a practical building block for ETL engines.