lablup / lablup/mlxcel

feat: multi-GPU tensor-parallel runtime with per-rank device placement and cross-GPU collectives

Open
#488 0 comments 0 reactions 0 assignees View on GitHub
area:architecture area:inference priority:high status:ready type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.