feat: multi-GPU tensor-parallel runtime with per-rank device placement and cross-GPU collectives
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #486.
Depends on #487.
## Problem
The in-process TP runtime shards weights correctly but materializes every rank on the single default GPU and reduces partial outputs locally, so there is no memory scaling and no real parallelism. A model larger than one GPU still cannot be loaded.
- `src/distributed/tensor_parallel/llama_runtime.rs:76-95` `from_full_weights`: all ranks built on the default device.
- `src/distributed/tensor_parallel/llama_runtime.rs:98-148`: forward runs ranks sequentially with `.iter().map()` and sums via `reduce_sum` (`llama_runtime.rs:2288`), all on one device.
Goal: place rank r's sharded weights and compute on GPU r, and replace the local sum with a real cross-GPU collective, so aggregate weights can exceed a single GPU and the forward runs across GPUs.
## Feasibility spike (do this first; it gates the design)
Confirm, within a single process on the pinned MLX (`ml-explore/mlx` @ `a6ec7123`, `src/lib/mlx-cpp/CMakeLists.txt:92-95`): (a) arrays can be allocated on a non-zero GPU index, (b) ops can run on a non-default device, (c) arrays can be moved between devices (peer copy). Record the outcome in this issue.
- If intra-process multi-GPU works: implement the intra-process design below.
- If it does not: fall back to a multi-process design (one worker process per GPU on localhost/NVLink, reusing the existing distributed transport and collective scaffolding in `src/distributed/`), or propose bumping the MLX pin. Document the decision before writing the runtime.
## Design (intra-process path)
Weight placement (memory scaling, the core win):
- Build rank r's model with its sharded weights resident on GPU r. Wrap the per-rank construction in `from_full_weights` so weight materialization runs on a device-bound (thread-local) stream for index r, using the sub-issue 1 API. Verify each GPU holds approximately 1/N of the weights.
Forward execution:
- Run each rank's per-layer attention and FFN on a worker pinned to GPU r (per-rank thread, each owning a device-bound thread-local stream; reuse the existing `MlxThreadLocalStream` pattern already used by the batch scheduler).
- Keep the residual stream `h` replicated on every device (standard Megatron-style TP). Broadcast `h` to all devices before the attention and FFN regions.
- Replace `reduce_sum` of partial attention (after o_proj) and partial FFN (after down_proj) with a cross-GPU all-reduce. The existing shape metadata already marks where reduction is required (`requires_allreduce_after_o_proj`, `CommPattern`).
Collective backend abstraction:
- Define a `CrossDeviceCollective` trait with two implementations:
1. Portable copy-and-reduce fallback: gather partials onto a root device using the sub-issue 1 copy-to-device primitive, sum, then broadcast back. Correct first, slower.
2. NCCL-backed all-reduce on CUDA for performance, by binding MLX's distributed/NCCL path or NCCL directly. Ship the fallback first; add NCCL as an optimization.
Reuse existing scaffolding where possible: `TPBarrier`, `synchronized.rs` step decisions, `split_rank_caches` (already per-rank), and the per-rank `TPCacheManager`. KV caches must be device-resident per rank.
## Acceptance criteria
- [ ] Spike outcome recorded and architecture chosen (intra-process or multi-process).
- [ ] `load_model_with_tensor_parallel` (`src/loading/mod.rs:431`) places shard r on GPU r; a memory probe confirms roughly 1/N weights per GPU.
- [ ] `mlxcel generate --tp-size N` and `mlxcel-server --tp-size N` execute across N GPUs and produce logits matching the single-GPU replicated reference within tolerance (extend the existing TP parity tests, for example `tests/tp_e2e.rs`, `tests/tensor_parallel_real_models.rs`).
- [ ] A model that does not fit on a single GPU loads and serves with `tp-size N` on real multi-GPU hardware (the DGX goal).
- [ ] `tp_size > gpu_device_count()` produces a clear, actionable error.
- [ ] Documented behavior on Metal/Apple (one GPU): `tp_size > 1` keeps the current in-process emulation or errors clearly; no silent wrong results.
- [ ] At least the currently supported TP model families (Llama/Qwen2, Qwen3, Qwen3.5, Gemma3, Gemma4, Ernie4.5, HunyuanV1Dense) work on the multi-GPU path, or the gap is explicitly documented.
- [ ] Docs updated; no clippy/fmt regressions; cross-platform checks pass.
## Risks and fallback
- MLX intra-process multi-GPU maturity is the main risk; the spike gates the design, with a multi-process fallback over the existing distributed transport.
- Cross-device synchronization correctness (barriers around each all-reduce) must be tested under concurrency.
## Out of scope
- Vocab-parallel embeddings and LM head for very large vocabularies (current default is replicated); track as a follow-up if needed.
- 2D pipeline-plus-TP composition (tracked separately).
- MoE expert-parallel multi-GPU placement beyond the listed dense families (follow-up).
Contributor guide
Research direction
Start with the MLX feasibility spike in src/lib/mlx-cpp/CMakeLists.txt:92-95, then inspect from_full_weights and forward handling in src/distributed/tensor_parallel/llama_runtime.rs. Review the existing distributed scaffolding and TP parity tests in tests/tp_e2e.rs and tests/tensor_parallel_real_models.rs. Done means the chosen architecture is documented and multi-GPU loading, execution, collectives, parity, and device-count errors meet the listed acceptance criteria.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100