GPU-Native Operators in Ray Data

Ray Data is a data engine built on Ray focused on AI workloads such as multimodal data processing and training data loading. Due to its focus on AI workloads, GPUs are a core feature of Ray Data - Ray Data is able to treat GPUs as a scheduling resource and thus is able to easily scale data processing tasks across GPUs.

Despite this, Ray Data has had limited native functionality that executes directly on the GPU, delegating that responsibility to user defined functions that call GPU-enabled libraries like PyTorch or JAX. 

Recently, the Ray Data team at Anyscale has collaborated with the NVIDIA cuDF team to better integrate GPU-based operators into Ray Data. In particular, we’re excited to announce two new features:

• cuDF as a backend for Ray Data map_batches

• GPU-based shuffle based on RapidsMPF 

On select data curation workloads, these enhancements have enabled up to 3x better TCO compared to CPU-based solutions.

You can try this out today on Ray by pip install -U ray==2.58.

LinkBackground

Ray Data offers a set of canonical data processing operations, from IO to shuffle to UDFs. One of Ray Data’s differentiating features is that it has first-class support for GPUs for its UDFs. For example, one can schedule custom code leveraging a GPU as follows:

import ray

ds = ray.data.read_parquet("s3://my-bucket/documents/")

def compute_gpu(batch: pd.DataFrame):
    ...

ds = ds.map_batches(compute_gpu, num_gpus=1, batch_size=4096, batch_format="pandas")

The above code will automatically launch multiple compute_gpu calls, all running on different processes attached to different GPUs across your Ray cluster.

In using Ray Data for data deduplication tasks, we noticed two opportunities for improvement:

1. There were compute-intensive dataframe operations in the pipeline that could be accelerated by using cuDF, but this would require the user to convert data formats manually

2. The specific tasks had heavy hash-aggregation steps that were compute intensive and could benefit from improved GPU acceleration.

In order to address these opportunities, we added two key features to Ray Data: cuDF batch format support and a GPU-native shuffle backend based on RapidsMPF.

LinkcuDF as a first-class batch format

Ray Data has 3 native batch formats - Pandas Dataframes, Numpy Arrays, and PyArrow Tables. Batch formats serve two purposes -- to be the data exchange format between stages, and to provide users with a more familiar interface to write custom processing functions.

As of Ray 2.58, map_batches now accepts batch_format="cudf", which hands your UDF a cudf.DataFrame instead of a pandas DataFrame or Arrow table.

cuDF is NVIDIA's open-source GPU-accelerated DataFrame library, one of the NVIDIA CUDA-X^TM libraries for data processing. It provides a Pandas-like API with GPU-accelerated operations. For certain computationally-heavy operations, cuDF can be up to order of magnitude faster and cheaper than alternative CPU equivalent operations.

import ray

ds = ray.data.read_parquet("s3://my-bucket/documents/")

def add_minhash(batch):  # batch is a cudf.DataFrame
    batch["minhash"] = batch["text"].str.minhash(seed=42, ...)
    return batch

ds = ds.map_batches(add_minhash, batch_format="cudf", num_gpus=1, batch_size=4096)

LinkGPU-native shuffle via RapidsMPF

Users can now leverage GPU-accelerated repartition and hash aggregations through a new GPU execution path built on RapidsMPF. Ranks talk to each other directly over UCXX. Between ranks, shuffled data take advantage of high performance transports like RDMA or NVLink, skipping the CPU and Ray object store.

Similar to cuDF, for certain computationally-heavy operations, RapidsMPF-based shuffle can be up to order of magnitude faster and cheaper than alternative CPU equivalent operations.

from ray.data.context import DataContext, ShuffleStrategy

ctx = DataContext.get_current()
ctx.shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE

ds.repartition(num_blocks=256, keys=["doc_id"])
ds.groupby("band_id").count()

LinkBenchmarks: fuzzy deduplication

The motivating workload for our collaboration has been fuzzy document deduplication, a standard workload for data curation for foundation model training. For our benchmarking, we used the standard MinHash-LSH approach , outlined below:

Fig 1: Stages in the deduplication workload

1. Generate MinHash signatures for every document.

2. Generate LSH bands from those signatures.

3. Group by band to find candidate duplicate pairs.

4. Deduplicate the candidate edges.

5. Compute connected components by traversing the graph.

6. Filter duplicates

7. Persist the survivors.

Step 1 involves generating hash values for each n-gram, and hashing is a compute-intensive operation.

Step 3 and Step 5 are the most shuffle-intensive -- a groupby operation is used to form duplicate pairs, and the connected components algorithm is a series of groupby operations.

In the below experiments, we were able to see improved performance for these steps due to GPU integration and acceleration.

In the below experiments, we used the FineWeb 10BT dataset on cluster sizes between 2 to 8 nodes. For CPU, we used the m8id.4xlarge instance type, which has 16 vCPUs, 64 GiB memory, and costs 1.04416 dollars per hour.

For GPU evaluation, we used the g6.4xlarge instance type, which has 16 vCPUs, 64 GiB memory, and costs 1.3232 dollars per hour.

Experiment code can be found here using Ray 2.58

LinkGenerating MinHash signatures

In this phase, we break each document into a set of n-grams (where n is set to 5), and we generate 128 hashes per each n-gram. 

The hashing step is compute-intensive, and we compare a baseline Numpy-based solution with the cuDF built-in hashing implementation.

From the above results, we see that GPU acceleration improves this stage of the workload up to 15x faster than a similarly sized CPU cluster.

LinkThe grouping stages

In steps 3 and 5, we perform locality-sensitive hashing, which creates groups documents if they have hash collisions; and we implement an iterative connected components detection algorithm in order to find documents that have transitive collisions (for example, if document A is found to be similar to document B, and document C is found to be similar to document B, then document A should be similar to document C).

In this comparison, while increasing the number of GPUs improves runtime, the marginal improvement diminishes at higher scale, thereby decreasing cost savings.

Since these stages are shuffle intensive, adding more GPUs adds more communication overhead.  For this particular dataset and workload, 2 GPUs was optimal, so the dataset was likely too small to take advantage of all 8 GPUs.

LinkEnd-to-end Pipeline

From the full pipeline on a fixed hardware configuration, we see that the GPU-accelerated solution is:

• 4x faster and 3.1x better TCO than the lowest cost CPU comparison point (at 2 nodes)

• Provides TCO and runtime improvements across the board on all node scales

LinkWhat's next

While the results above are promising, we are still far from the upper limits of performance. In particular, operator fusion between sequential operators that use cuDF and RapidsMPF will provide a significant step change in performance. This will allow Ray Data to avoid unnecessary host-to-device round trips between stages, allowing the data to stay within GPU memory.

In working with end users, another step is to enable GPU shuffle for data preprocessors so that feature-engineering pipelines built on Ray Data's preprocessor API pick up the GPU path.

We’re excited to continue investing in this area to bring the best data processing performance for your AI data workloads. Ray Data GPU shuffle is experimental today and available as part of Ray 2.58. If you try it, we'd like to hear your feedback!