RPC: add -sm tensor - #26610
RPC: add -sm tensor#26610am17an wants to merge 5 commits into
-sm tensor#26610Conversation
|
confirmed working on metal(RDMA 2 x M3 Ultra with #26421), testing same model (ds4 MXFP4): it does however break with dspark applied because some ops it depends on appear to not be supported with TP (add across the split), so this is without any mtp. |
|
@ryan5rdx yeah I know it doesn't work with dspark, but what do you get when for just |
command for reference, let me know happy to test an alternate config: and yup just confirmed - with |
|
@ryan5rdx try using two rpc servers, one on each machine and connect via
|
|
It might be worth to add I don't know if this would be observable with RDNA, but sometimes manually moving tensors instead of standard |
|
Yeah I messed up, I think #26490 should be okay to merge though |
Don't we want to fix the DSpark support first? |
|
The support is broken over RPC I think(i.e. this PR), not in general. But I can check |
Sorry - to clarify - is this RPC servers on two RDMA linked nodes, and then a llama-server instance on one of them(if so I suppose this will just connect to the localhost RPC server)? Sorry I've never run a llama-cli/server instance where it's not also doing compute ha |
|
@ryan5rdx yes, the llama-server is just a client of these two. Basically we're trying to activate the RPC<>RPC all-reduce path rather than the one you probably got (Metal<>RPC) |
Currently on master |
Tested with Qwen327B, 0.6B and ds4. no obvious error on rpc servers or llama-server, I see the tensors copy on the RPC nodes then nothing. final client logs, llama server web UI never comes up(hangs here): command: |
for the all-reduce to work - don't RPC nodes need direct(RDMA) connections to each other in addition to (at least)TCP to the client? with an A - B(just client) - C topology where A<>B and B<>C are RDMA links, but there is no A <> C link, can this work? (in MLX they achieve this with a mesh + rdma, but here RPC servers don't know about peers yet) |
rgerganov
left a comment
There was a problem hiding this comment.
Please create a mermaid sequence diagram (similar to this one ) which describe how peers communicate when -sm tensor is being used, I am still trying to understand the new flows being added and that would be very helpful. In fact, I think this should be part of our dev documentation (feel free to create an .md file) so we can maintain this in the long term.
| return true; | ||
| } | ||
|
|
||
| // minor protocol version of each connected server, used to gate newer commands (comm collectives) |
There was a problem hiding this comment.
no need to do this, we don't care about backward compatibility and we prefer to keep the code simple; just bump the version to 6.0.0 and expect all peers to be running this version
|
@ryan5rdx I run this using the directly the RDMA interfaces, for example in my case it is |
in my case the RPC nodes cannot talk to each other via the Possible I'm misunderstanding here - or maybe a mac difference because it only supports p2p (TB)rdma without routing vs the spark? |
|
@ryan5rdx not sure, you can probably ask an LLM to debug? |
Ok yup confirmed - so if I use if I instead I use the LAN IPs - it still works becuase each worker can discover the other using the ip from the client, just super slow as expected (>>10s/token). However - if we want this to work on metal, with worker<>worker RDMA we need a mesh and to allow for passing the TB IP of other workers to each rpc server(because peer addresses are not derivable from the addr we get from the client in a TB p2p network). I'm not familiar with sparks but apparently it works because "all three nodes sit on one routable RDMA fabric" |
|
@rgerganov I added a mermaid diagram. Note that the Meta backend creates a lot of subgraphs so that's why graph caching is vital. |
rgerganov
left a comment
There was a problem hiding this comment.
you can improve the diagram by making a clear separation between messages sent to the standard RPC port and messages sent to the new "comm" port
| return nullptr; | ||
| } | ||
| ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *) backends[i]->context; | ||
| // one rank per endpoint: a server processes its socket sequentially, so a second |
There was a problem hiding this comment.
"one rank per endpoint" -- why having this limitation? i can have an endpoint with two devices which communicate very fast (because they are on the same physical host)
There was a problem hiding this comment.
In that case you can just create two rpc servers, one for each device?
There was a problem hiding this comment.
On local MoE models with -cmoe weights are copied to GPU during PP for faster compute. Utilizing separate processes per device would break this optimization.
The example is using local, but same optimization should be possible on rpc with multiple backends.
Utilizing same process can also be useful in general for AllReduce - 2 servers 2 GPU each could merge local results lowering the amount of network hops needed.
| if (!parse_endpoint(ranks[0].endpoint, host0, port0)) { | ||
| return nullptr; | ||
| } | ||
| const uint32_t comm_port = (uint32_t) port0 + 1000; |
There was a problem hiding this comment.
expect this port to be specified by the user via command line flag
| } | ||
|
|
||
| static void * ggml_backend_rpc_comm_init(ggml_backend_t * backends, size_t n_backends) { | ||
| if (n_backends != 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) { |
There was a problem hiding this comment.
is this going to work with more than 2 backends or it will require major redesign?
There was a problem hiding this comment.
it's going to fall-back to meta backend's all-reduce, which is slow but works
|
@am17an have you looked at implementing https://github.com/mk1-project/quickreduce/ QuickReduce seems pretty cool and it could probably be of huge benefit for RDMA performance. |
|
it's rebased on latest master. The async events seem to have added a lot of latency in sync due to worker threads wake up. When using |
a269ab9 to
28bd87b
Compare
Assisted-by: OpenAI Codex
|
Here are my MTP bench results on 2x sparks after #28387, #28390 This is without using NCCL and any extra dsv4 specific optimizations (which in my tests bring the wall time < 40s) |
This comment was marked as low quality.
This comment was marked as low quality.
f9b4233 to
04eec7a
Compare
| // Aux graph contents are rewritten on every compute but are identical across calls while the subgraphs are reused, | ||
| // so they can get stable uids on rebuild. Only safe without a comm backend, where the fallback usage is deterministic. | ||
| if (backend_ctx->comm_ctx == nullptr) { | ||
| for (ggml_cgraph * cgraph_aux : backend_ctx->cgraphs_aux) { | ||
| cgraph_aux->uid = ggml_graph_next_uid(); | ||
| } | ||
| } |
There was a problem hiding this comment.
This is not safe to do in general because with the API a user can pass arbitrary graphs. The correct logic would I think not be easy to implement and introduce non-negligible complexity. What is the specific motivation for adding this change?
There was a problem hiding this comment.
Sorry I just saw this. These changes are not required for this to work.
fc8b53b to
718ecbc
Compare
|
Hi, I have a bit different graph cache implementation made as separate patch. Could it be beneficial for this sm tensor feature branch? |
Partial cherry-pick of the graph caching and the RPC_CMD_SET_TENSOR_2D / RPC_CMD_GET_TENSOR_2D parts of PR ggml-org#26610. The -sm tensor comm support is not included. - server caches computed graphs per uid with bounded eviction - client tracks cached graph uids and reuses them - add RPC_CMD_SET_TENSOR_2D / RPC_CMD_GET_TENSOR_2D - bump the RPC protocol to 7 Assisted-by: DeepSeek V4 Flash
|
@rgerganov do you have time to review this PR? I'm pretty confident of the changes and I volunteer to maintain them |
cf00d3b to
d09d1c7
Compare
Partial cherry-pick of the graph caching and the RPC_CMD_SET_TENSOR_2D / RPC_CMD_GET_TENSOR_2D parts of PR ggml-org#26610. The -sm tensor comm support is not included. - server caches computed graphs per uid with bounded eviction - client tracks cached graph uids and reuses them - add RPC_CMD_SET_TENSOR_2D / RPC_CMD_GET_TENSOR_2D - bump the RPC protocol to 7 Assisted-by: DeepSeek V4 Flash
d09d1c7 to
9a70593
Compare

Overview
Add RPC
-sm tensor. This is on 2x Sparks connected via RDMA.For RPC following changes are required:
all_reduceset_tensor_2d,get_tensor_2dLooking for feedback @ggerganov @rgerganov
Additional information
sequenceDiagram participant C as Client (RPC backend) box Server A (rank 0) participant A as A: rpc port participant Ac as A: comm port end box Server B (rank 1) participant Bc as B: comm port participant B as B: rpc port end Note over C,B: initialization C->>A: COMM_INIT (rank 0) C->>B: COMM_INIT (rank 1) Note over Ac: listen on comm port Bc-->>Ac: connect + caps negotiation (transport upgrade, e.g. RDMA) A->>C: response (ok) B->>C: response (ok) Note over C,B: for each subgraph (split at reduction boundaries) C->>A: GRAPH_COMPUTE (uid) [GRAPH_RECOMPUTE on reuse] C->>B: GRAPH_COMPUTE (uid) Note over A: compute subgraph async Note over B: compute subgraph async C->>A: COMM_ALLREDUCE (partial tensor) [fire and forget, no response] C->>B: COMM_ALLREDUCE (partial tensor) Note over A: sync pending graph<br/>cast F32→BF16 if ne ≥ 32768<br/>copy partial to send buffer Note over B: sync pending graph<br/>cast F32→BF16 if ne ≥ 32768<br/>copy partial to send buffer Ac-->>Bc: partial A (rank 0 sends first) Bc-->>Ac: partial B (rank 1 receives first) Note over A: upload partial B<br/>dst = dst + partial B (async ADD) Note over B: upload partial A<br/>dst = dst + partial A (async ADD) Note over C,B: read back the output C->>A: GET_TENSOR (output) Note over A: sync all backends A->>C: dataRequirements
Stack created with GitHub Stacks CLI • Give Feedback 💬