Feature Request: Support multiple devices on a single rpc-server

[nguha]

Prerequisites

• I am running the latest code. Mention the version if possible as well.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

Make a signle rpc-server able to handle multiple devices (GPUs). That way it could pass data directly from one device to the next when applicable.

Motivation

Goal: fully utilize the interconnect (PCIe) available on the rpc-server machine.

As of now (version b6084) you must launch one rpc-server instance per device (typically GPU) on your inference server.
If you happen to have 2 or more devices this results in a sub-optimal usage of the available interconnects (e.g. PCIe).
As far as I can see, even when devices on the rpc-server machine host contiguous model layers, all communication between them is still going through the network, to the client (llama-cli) and back.

I guess this was done for the sake of simplicity and ease of implementation.
There's therefore some performance gains to be made by making a single rpc-server able to handle all devices on a machine.

P.S. Thank you so much to everyone who put in the effort to build the rpc-server. It is such a great feature!

Possible Implementation

No response

[rgerganov]

One way to implement this is to introduce a new backend which is able to wrap existing backends, similar to what @slaren proposes for supporting tensor parallelism here. Implementation wise, this should be simpler than tensor parallelism support but it is still non trivial task. I will think about this more and try to come up with some PoC.

P.S. Thank you so much to everyone who put in the effort to build the rpc-server. It is such a great feature!

thank you for the kind words

[slaren]

It may be simpler to allow the server to expose multiple devices, and then just forward the ggml_backend_tensor_copy calls to the server. This would avoid most of the large transfers between the server and the client.

[jukofyork]

Goal: fully utilize the interconnect (PCIe) available on the rpc-server machine.

Yeah, I tried "weaving" RPC through 6 cards over 3 nodes, and even with 10gbit Ethernet I didn't get any real gains (it was a couple of months ago I tried this, so it's possible the RPC code has improved since then though).

[rgerganov]

It may be simpler to allow the server to expose multiple devices, and then just forward the ggml_backend_tensor_copy calls to the server. This would avoid most of the large transfers between the server and the client.

I can think of two ways to implement this:

1. Start a new server thread for each GPU device, each listening on a separate port. So if you have 3 GPUs, rpc-server will open ports 50052, 50053 and 50054 for each of them. On the client side ggml_backend_rpc_buffer_cpy_tensor will check if both src and dst tensors are on the same host (not on the same socket as it is now) and forward the call if they are. In this case the RPC protocol remains the same and RPC client usage remains unchanged with the assumption that you cannot start more than one server on the same host (can be enforced if needed). I made a quick and dirty prototype on a Linux host with two NVIDIA cards and it appears to work although I am not sure if CUDA backends are completely thread-safe.
2. Change the RPC protocol to include a device number in each command and multiplex on the server side to different backends/GPUs. This will require a lot more work but we can continue to use a single port for the entire server.

@slaren what do you think?

[slaren]

@rgerganov sounds about right. I think option 2 would be nicer to use (only need to write one address per server, instead of per device), but it also requires more effort to implement.

[wallentri88]

I played around with rpc a while ago and, in addition to option 2, made something like this: in function ggml_backend_sched_compute_splits when meeting RPC backend, I iterated until first non-rpc backend, collecting tensors to send to server, compute graphs, and tensors to get back from server (from first non-rpc). And after doing it, I send all this data in one batch. This eliminates:

1. All copy tensors requests when tensors are on the same servers
2. Multiple compute graphs requests (instead doing one)
3. Multiple get tensor requests from first non-rpc backend (they all returned as part of return data of batch request)

This is, of course, the case when I have one RPC server with multiple devices, didn't think about several RPCs at the moment.

While it sounds good, the draft code in ggml_backend_sched_compute_splits looked really awful, so I decided not to share it, haha. I think probably the proper solution would be creating a more abstract RPC backend without exposing devices to main node. So RPC server should do ggml_backend_sched_split_graph for his layers itself. And when we meet this abstract multidevice RPC backend in main node ggml_backend_sched_compute_splits, we just need to send input tensors and get final output tensors, but not send compute graphs: RPC server should create it by himself because it knows his devices best. Is it even possible, I wonder?