Accelerated Out of Core Shuffling
2 days ago
- RapidsMPF是一个可重用的、外核的洗牌器，旨在在大规模数据洗牌过程中防止OOM错误，在DGX B200上实现高达1.8 TiB/s的全局吞吐量。
- 洗牌对于分布式连接、分组、合并和排序至关重要，但由于内存密集型、数据传输成本和同步障碍，代价高昂。
- 分布式连接需要将匹配的键路由到相同的等级，这可能会迫使一个等级持有多个数据副本：源表、暂存缓冲区、接收到的切片、哈希表和输出。
- RapidsMPF包括一个洗牌库（C++/Python）和一个用于流管线的参与者网络；它被cuDF Polars、NeMo-Curator使用，并在Ray Data中实验性使用。
- 在8个GPU和每个等级20 GiB的基准测试中，无约束洗牌以约1.8 TiB/s的全局吞吐量运行，每个等级的峰值设备内存使用量为60 GiB，是输入大小的3倍。
- 当设备内存限制为32 GiB时，吞吐量下降到约480 GiB/s，但RapidsMPF通过直接在主机上接收传入缓冲区来避免OOM，防止无意义的设备到主机驱逐。
- 将内存限制从32 GiB扫到12 GiB会产生平滑、单调的性能下降，不会出现OOM；只有当限制接近输入大小时，真正的溢出才开始。
- 溢出期间的主要瓶颈是通过PCIe Gen 5的主机到设备传输，平均约32 GiB/s；Grace-Blackwell C2C可以提供5-10倍的带宽。
- RapidsMPF通过可调溢出和主机侧缓冲区接收、通过UCXX（NVLink、InfiniBand、EFA、TCP）的传输以及通过异步/流执行的同步来解决内存压力。
- 该项目证明了外核洗牌器可以扩展、避免抖动，并处理比VRAM更大的数据，使其成为ETL引擎的实用构建块。