rpc: cut the protocol cost of uploading weights to a remote server - #199
rpc: cut the protocol cost of uploading weights to a remote server#199danielhanchen wants to merge 5 commits into
Conversation
Loading a 27B layer split spends most of its time in the weight upload, and most of that time is protocol, not wire. Measured on one node with an rpc-server on the RoCE interface (Qwen3.8-27B UD-Q4_K_XL, 15.7 GiB pushed to the remote backend, RDMA active), the upload phase cost 41.5 s and broke down as: hashing 12.6 s every tensor over 10 MiB hashed with FNV-1a at about 1.2 GiB/s staging 4.3 s a fresh zero filled buffer per tensor, then a copy into it wire 4.3 s other 20.4 s client side stalls outside the RPC calls Four changes, all inside ggml/src/ggml-rpc: - The server now says at HELLO whether it keeps a tensor cache, in the byte that used to be padding in the response. Without a cache the answer to SET_TENSOR_HASH is always "not cached", so the hash pass over every large tensor was pure cost. The client only hashes when the server can use it. The message keeps its size, an older server sends a zero byte and an older client ignores it, so both directions interoperate unchanged. - SET_TENSOR is written from the header and the caller's payload directly instead of being copied into one contiguous buffer first. That buffer cost a zero fill and a full copy of every tensor. The bytes on the wire are identical. - The server reads a SET_TENSOR message off the connection instead of into a vector sized to the whole message, and receives the payload straight into the destination when the backend buffer is host memory. Non host backends reuse one staging allocation that is never zero filled. - The RDMA transport keeps up to eight 256 KiB chunks in flight instead of posting one and polling it to completion before posting the next, and drains them at the message boundary that flush() already marks. Receives are now consumed byte by byte from the completed buffer, so a peer that frames a message differently keeps working instead of losing the remainder of a frame. After the change the same upload phase is 23.5 s, with hashing and staging at zero and 4.9 s on the wire. GGML_RPC_LOAD_OPT=0 restores the previous behaviour for A/B. GGML_RPC_LOADPROF=1 turns on a load profiler on both sides: per command counts and times on the client, and the split of the upload into hashing, staging, wire and the gaps between calls. It is off by default and costs one relaxed atomic load per call.
|
@codex review |
|
@codex security review |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4b7f23fa19
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| const void * data = input.data() + sizeof(rpc_tensor) + sizeof(offset); | ||
| // only a host buffer can take the payload directly; hashing for the cache also needs it contiguous | ||
| const bool direct = ggml_backend_buffer_is_host(tensor->buffer) && cache_dir == nullptr; |
There was a problem hiding this comment.
When an RPC server uses a host-visible buffer whose set_tensor hook performs more than a memcpy, this shortcut silently bypasses required processing. For example, Hexagon enables host buffers by default, reports them as host-visible, and its hook repacks Q4/Q8/MXFP4 tensors into the tiled layout consumed by its kernels; with no cache directory, these tensors are now received directly in their original layout and the hook is never called, producing incorrect inference results. Direct streaming should be limited to buffers known to have memcpy semantics, or the backend hook must still be invoked.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in bace306. The substance is right and this was a real silent-wrong-results path, so thank you.
One correction to the example, which does not change the verdict. The Hexagon repack buffer type reports is_host = false, so it was never eligible for the shortcut. It is the plain Hexagon buffer type that qualifies: its is_host returns opt_hostbuf, and its set_tensor is the one that repacks Q4_0, Q4_1, Q8_0, IQ4_NL and MXFP4 into the tiled layout.
The shortcut is now limited to the CPU buffer type, whose set_tensor is a memcpy by construction. That keeps the benefit where it exists, since a device-backed rpc-server is not host-visible and was already taking the staged path.
| rdma->rx_buf = aligned_alloc(4096, static_cast<size_t>(RDMA_RX_DEPTH) * RDMA_CHUNK); | ||
| if (!rdma->tx_buf || !rdma->rx_buf) return false; | ||
|
|
||
| rdma->tx_mr = ibv_reg_mr(rdma->pd, rdma->tx_buf, RDMA_CHUNK, IBV_ACCESS_LOCAL_WRITE); | ||
| rdma->rx_mr = ibv_reg_mr(rdma->pd, rdma->rx_buf, static_cast<size_t>(RDMA_RX_DEPTH) * RDMA_CHUNK, | ||
| IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); |
There was a problem hiding this comment.
On systems whose memlock allowance is sufficient for the previous 256 KiB TX MR plus the 6 MiB RX MR but not the new 2 MiB TX MR plus RX, the largest TX registration succeeds by itself and the RX registration then fails, causing RDMA probing to fall back to TCP. The depth-reduction loop therefore does not provide the claimed graceful degradation for the relevant combined-memory limit; RX registration must be included in the retry so smaller TX depths are attempted.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in f06bbb2. Confirmed, and the code contradicted its own comment about degrading instead of failing, since the limit that binds is TX and RX together and only TX was being retried.
RX registration moved inside the loop. The RX buffer is a fixed size so only its registration is retried, and a TX depth that leaves no room for RX now deregisters and tries the next depth down.
ggml_backend_buffer_is_host says the memory is addressable, not that set_tensor is a plain memcpy. ggml-hexagon reports host buffers when opt_hostbuf is set, and its set_tensor repacks Q4_0, Q4_1, Q8_0, IQ4_NL and MXFP4 into the tiled layout its kernels read, so receiving into tensor->data and skipping the hook left those weights in the wrong layout and produced wrong results with nothing to show for it. The shortcut is now limited to the CPU buffer type, whose set_tensor is a memcpy by construction; every other backend keeps the staged path and its hook. That costs nothing in practice, since a device-backed rpc-server is not host-visible and was already staging.
The loop halved the TX ring until its registration succeeded, but RX was registered once afterwards, so the comment's promise that a tight memlock degrades instead of failing did not hold for the limit that actually binds. On a system whose allowance covers the largest TX on its own but not TX plus RX, the loop stopped at the largest TX, RX then failed, and RDMA fell back to TCP without any smaller depth being tried. RX is registered inside the loop now. Its buffer is a fixed size, so only the registration is retried.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f06bbb22d9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| printf("Accepted client connection\n"); | ||
| fflush(stdout); | ||
| rpc_serve_client(backends, cache_dir, client_socket); | ||
| g_rpc_loadprof_server.print(); |
There was a problem hiding this comment.
When the persistent server accepts a second client, this prints cumulative state from the global g_rpc_loadprof_server rather than statistics for that connection. None of its counters or timestamps are cleared, so first_ns still refers to the first client's upload and ns_gap includes the idle time between clients, making the reported span, gaps, and effective rates misleading for subsequent experiments. Reset the profiler after printing or keep one profiler per client connection.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Accepted. g_rpc_loadprof_server is global and nothing clears it between clients, so a second connection prints a span that starts at the first client upload and gaps that include the idle time between clients. The numbers are not merely stale, they are wrong in a way that reads as plausible, which is the worst kind for a profiling aid whose whole purpose is to be quoted.
Of your two options I prefer a reset after printing rather than one profiler per connection: the server serves one client connection at a time, so a per-connection instance would add lifetime management for no extra information, whereas a reset at the point the summary is emitted keeps the reporting boundary and the data boundary the same thing.
Not fixed yet, so no thumbs-up.
There was a problem hiding this comment.
Confirmed and fixed in d2ff662.
g_rpc_loadprof_server is a single global and print() is called once per client in the accept loop, with nothing clearing it in between, so every connection after the first reports the earlier ones folded in. All three consequences you list follow: the counters are cumulative, first_ns still points at the first client's first upload, and ns_gap has the idle time between clients added to it, which makes the span, the gap total and the effective rates wrong for precisely the repeat experiments this profiler exists to support.
Took the "reset after printing" option rather than one profiler per connection, since the accept loop serves one client at a time and a per-connection object would be a larger change for the same result.
One detail worth calling out beyond the counters: last_end_ns has to be cleared as well. Left set, the first SET_TENSOR of the next client takes the prev != 0 branch and is charged as a gap from the previous client's last upload, instead of the else branch that records it as a new first_ns. Clearing only the counters would have left that one wrong in a way that is easy to miss, because the report would look plausible.
The constructor now calls reset() too, so the field list is in one place rather than partly initialised in two.
The server profiler is a single global printed once per client connection, and nothing cleared it, so every connection after the first reported the earlier ones as well. The counters were cumulative, first_ns still pointed at the first client's first upload, and ns_gap had the idle time between clients folded into it, which makes the span, the gap total and the effective rates wrong for exactly the repeat experiments this is meant to support. Adds reset() and calls it after print(). It clears last_end_ns too, which matters as much as the counters: left set, the first upload of the next client is charged as a gap from the previous client's last one rather than treated as a new start. The constructor now uses it as well, so the field list lives in one place instead of being partly initialised in two.
1 participant
Loading a model that is split across two machines spends most of its time pushing the
remote half of the weights over RPC, and most of that time is protocol overhead rather
than wire time. On a pair of DGX Sparks with 111 Gb/s RoCE rails a 27B Q4_K_XL layer
split takes about fourteen minutes to become ready while both CPUs sit at 99 percent and
both GPUs read zero. This branch attacks the protocol side of that.
Where the time goes
Measured on a single node, with an
ggml-rpc-serverbound to the RoCE interface so thewhole transport is exercised (RDMA active,
mtu=4096), pushing all 15.7 GiB ofQwen3.8-27B-UD-Q4_K_XL to the remote backend. Numbers from the load profiler added here
(
GGML_RPC_LOADPROF=1), for the span from the firstSET_TENSORto the last:876
SET_TENSORcalls carry the 15.7 GiB; 369 of them are over 10 MiB and were the onesbeing hashed. The client side stalls are four events, one of them 10.7 s, and are not per
tensor overhead: they are outside the RPC calls and are left for a separate change.
What changed
Everything is inside
ggml/src/ggml-rpc/.GGML_RPC_LOAD_OPT=0restores the previousbehaviour on both sides for A/B.
The hash pass is skipped when it cannot help.
SET_TENSOR_HASHlets a server thatkeeps a tensor cache answer "I already have this" instead of taking the upload. A server
started without
-chas no cache, so the answer is always no and the FNV-1a pass overevery tensor above 10 MiB is pure cost, about 1.2 GiB/s of pure serial multiply. The
server now advertises whether it has a cache in the byte of the HELLO response that used
to be padding, and the client only hashes when that bit is set. The message keeps its
size and its version, an older server sends a zero byte and an older client ignores it,
so old and new interoperate in both directions.
No staging copy for the upload. The client used to allocate a zero filled buffer the
size of header plus tensor, copy the header and then the tensor into it, and send that.
For a 27B split that is several gigabytes of zero fill and several more of copy. The
header and the caller's payload are now written straight to the socket. The bytes on the
wire are byte for byte what they were.
The server streams the message instead of buffering it.
RPC_CMD_SET_TENSORis readoff the connection: header first, then the payload directly into the destination when the
backend buffer is host memory, or into one reused staging allocation that is never zero
filled when it is not (CUDA, or when a cache dir means the payload has to be hashed). It
used to
resize()a fresh vector per message, which zero fills, and then copy out of it.The RDMA transport stops waiting for every chunk.
rdma_sendposted one 256 KiBchunk and polled it to completion before posting the next, so the link was idle for a
full round trip on every chunk and only one chunk was ever in flight. It now keeps up to
eight in flight against a ring of registered buffers and drains them at the message
boundary that
flush()already marks. The receive ring is 24 deep, so the sender cannotoutrun it, and the number of send slots degrades down to one if a tight memlock limit
refuses the registration.
Receives are consumed byte by byte.
rdma_recvused to copy a whole completedreceive into the caller's buffer and decrement by the full
byte_len, so a peer whosemessage framing differed would have the remainder of a frame dropped, or would overrun
the caller's buffer. A completion is now consumed across as many
recv_datacalls as ittakes. That is what makes the split write above safe against an older peer.
Interoperability
All four combinations were run on the 27B, each loading the whole model over RPC and
generating the same text:
GGML_RPC_LOAD_OPT=0New client against old server keeps the hash saving and the send side pipelining but not
the server side streaming receive, which is why its wire time is higher than new against
new. Old client against new server keeps the streaming receive only.
Non-RPC workloads
The whole diff is three files in
ggml/src/ggml-rpc/. No file outside that directory istouched, so a build without the RPC backend compiles identical code.
-DGGML_RPC=OFF -DGGML_CUDA=ONconfigures and builds clean, all 81 targets.test-backend-ops -b CUDA0on the new build:2/2 backends passed,OK.Single GPU greedy, no
--rpc, three prompts, 64 tokens, temperature 0, seed 1:md5
111e44cd31da973afbe5aa190b2fbed7for base and for this branch. The raw outputdiffers only in the tokens per second footer llama-cli prints, which is not the same
from run to run on either build.
Single GPU
llama-batched-bench, no--rpc, 27B,-npp 512 -ntg 128 -npl 1,4,base / new / base in one window with the new build inside the bracket:
Pair
Two DGX Sparks, 27B UD-Q4_K_XL split across both,
--device CUDA0,RPC0 -sm layer -c 16384 --parallel 32 --cache-ram 0 -t 6, 8538 MiB ofweights going to the peer, RDMA active on both ends. Load is the
model loadedstamp inthe llama-server log; the arms alternate base, new, base, new in one window, both nodes
uncapped:
Load 19.1 s to 8.6 s on the mean of two arms each. Prefill and decode are unchanged, as
expected: nothing on the inference path changed.
Greedy output over the split is byte identical in all four arms,
md5
371c236256c5c537b2e72ed4add8571b(400 token prompt, 64 tokens, temperature 0).The client side breakdown of a base load on the pair, from the profiler: of the 19.4 s,
6.85 s is hashing, 2.39 s is host staging and 3.06 s is wire, so 12.3 s of the load is
inside the RPC calls and a third of the whole load is a hash whose answer is thrown away.
Profiler
GGML_RPC_LOADPROF=1prints, on both client and server, the number ofSET_TENSORcalls, the bytes, and the split between hashing, staging, wire and the gaps between
calls, plus a per command count and time on the client. It is off by default and costs
one relaxed atomic load per call when off.