Moving GPU Tensors Between Ray Workers Without the Serialization Tax

The pattern is common. A pipeline is split into stages, each stage runs as a Ray actor on its own GPU, and Ray passes the output of one stage to the next. It works, it scales out on paper, and then nvidia-smi shows the GPUs sitting well below full utilisation while requests queue up.

Adding GPUs rarely helps, because the GPUs are not the bottleneck. The time is going into moving data between them. Every handoff between stages has a cost that is easy to miss in the code and very visible on the bill.

Ray's object store lives in host memory, not on the GPU. So when one actor returns a CUDA tensor and another actor receives it, the data takes a long way round:

A copy off the GPU and another back onto it over PCIe, a serialisation pass, a deserialisation pass and possibly a network hop, all to move data that could have gone straight from one GPU to the other. Two GPUs in the same machine may even share an NVLink connection that this path never touches.

1@ray.remote(num_gpus=1)2class Encoder:3    def forward(self, batch):4        return self.model(batch.cuda())  # returns a CUDA tensor5 6@ray.remote(num_gpus=1)7class Decoder:8    def forward(self, features):9        return self.head(features)10 11# Ray stores the encoder's output and hands it to the decoder.12features = encoder.forward.remote(batch)13result = ray.get(decoder.forward.remote(features))

Do not start rewriting the data path on a hunch. Three tools, used in this order, will tell you whether handoffs are really the problem and which one hurts most.

Write down three numbers before you touch the code: time per request, GPU utilisation per stage, and how many bytes cross each handoff. Every fix below should move at least one of them, and you want to know which.

Many pipelines start with a CPU stage: decoding audio or images, tokenising, resizing. Ray can read NumPy arrays out of its shared-memory object store without copying them, as long as the reader is on the same node. The array you get back is a read-only view into shared memory. Large data handed to Ray as other Python objects is usually rebuilt with a copy instead, so check the serialisation notes for the Ray version you run.

1@ray.remote2def preprocess(raw) -> np.ndarray:3    x = decode(raw)  # your CPU work4    return np.ascontiguousarray(x, dtype=np.float32)5 6@ray.remote(num_gpus=1)7class Model:8    def predict(self, x: np.ndarray):9        # x is a read-only view into the node's object store10        t = torch.from_numpy(x)  # no copy (warns that x is read-only)11        t = t.pin_memory().to("cuda", non_blocking=True)12        return self.model(t).cpu().numpy()

Two habits make this pay off. Batch before you hand off, because each Ray object has a fixed overhead and thousands of tiny objects cost more than one large one. And keep arrays contiguous with a fixed dtype, so nothing has to be converted on the other side.

If both stages run on GPUs, the cleanest fix is to stop sending the tensor through Ray at all. Let Ray do what it is good at, scheduling, placement and restarts, and move the data with NCCL, the library PyTorch already uses for distributed training. NCCL uses NVLink between GPUs in the same machine when it is available, and the network between machines.

1# Once, when each actor starts: join the same NCCL group.2dist.init_process_group(3    "nccl", init_method=f"tcp://{head_ip}:29500",4    rank=rank, world_size=2,5)6 7# Stage A (rank 0): features is a CUDA tensor.8dist.send(features, dst=1)9 10# Stage B (rank 1): receive straight into GPU memory.11features = torch.empty(shape, dtype=torch.float16, device="cuda")12dist.recv(features, src=0)

The trade-off is coordination. Both sides must agree on shape and dtype before the transfer, and a missing send leaves the receiver waiting, so send a small header with the shape first or keep shapes fixed per batch size. Recent Ray releases also include Compiled Graphs, which can move tensors between actors over NCCL for you. It is worth evaluating before you hand-roll the plumbing above, since it handles the coordination.

Some copies are unavoidable, such as the first transfer from host to device. You can still stop the GPU from waiting on them. Copy from pinned host memory on a separate CUDA stream while the model works on the previous batch, so the transfer and the compute run at the same time.

1copy_stream = torch.cuda.Stream()2 3with torch.cuda.stream(copy_stream):4    next_gpu = next_cpu.to("cuda", non_blocking=True)  # next_cpu is pinned5 6out = model(current_gpu)  # runs on the default stream meanwhile7 8torch.cuda.current_stream().wait_stream(copy_stream)9next_gpu.record_stream(torch.cuda.current_stream())10current_gpu = next_gpu

The record_stream call matters. Without it, PyTorch's caching allocator only knows the tensor was used on the copy stream and may hand its memory to someone else while the default stream is still reading it. That bug produces wrong results, not crashes, which makes it painful to find.

Teams sometimes reach for a hand-built transport: a pinned shared-memory ring buffer, CUDA IPC handles passed between processes, a small native daemon. These can be very fast. They also bring memory management, crash recovery, back-pressure and synchronisation bugs that only show up under load.

“Build your own transport only after the standard ones are measured, tuned and still not fast enough. Most pipelines never get that far.”

If you do go there, keep control messages and data apart. Ray passes a small descriptor, such as a buffer index or an IPC handle, while the bytes move through your transport. Then load-test at full saturation for hours, not minutes. Lock and memory-pressure bugs in this kind of code tend to appear long after a short benchmark has passed.