lablup / lablup/mlxcel

feat: single-node multi-GPU tensor parallelism to load and serve models larger than one GPU (CUDA/DGX)

Open
#486 1 comment 0 reactions 1 assignee Claimed by @inureyes View on GitHub
area:architecture area:core area:inference priority:high status:in-progress type:enhancement
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Goal

Make it possible to load and serve a model whose weights exceed a single GPU's memory by sharding those weights across multiple physical GPUs on one node (for example a DGX-class box with 8 GPUs over NVLink), with each tensor-parallel rank's weights and compute resident on its own GPU and real cross-GPU collectives between ranks.

The headline deliverable: `mlxcel-server --tp-size N` (and `mlxcel generate --tp-size N`) loads a model that does NOT fit on a single GPU and serves it correctly across N GPUs on one node.

## Current state

Tensor parallelism is extensively implemented (shard planning, ring collectives, parallel attention/FFN/MoE shape math, sharded loading, TP scheduler/executor/barrier, per-rank KV cache) and an in-process TP runtime is wired into the load path:

- `src/loading/mod.rs:431` `load_model_with_tensor_parallel` constructs `TensorParallel*Model` when `tp_size > 1` (supported: Llama/Qwen2, Qwen3, Qwen3.5, Gemma3, Gemma4, Ernie4.5, HunyuanV1Dense).
- Reachable from CLI generate (`src/commands/generate.rs:130`) and the server (`src/server/model_worker.rs:234,753`).

However the runtime is a single-device numerical emulation, not real multi-GPU execution:

- `src/distributed/tensor_parallel/llama_runtime.rs:98-148`: the forward iterates ranks with a sequential `.iter().map()` on the default device and sums partials locally via `reduce_sum` (`llama_runtime.rs:2288`). No threads, no per-rank device, no network/NVLink collective.
- `src/distributed/tensor_parallel/llama_runtime.rs:76-95` `from_full_weights`: every rank's sharded weights are materialized on the single default device, so aggregate memory is not reduced and a model larger than one GPU still cannot be loaded.

Root cause in the bindings: the device API is binary CPU/GPU only, with no device index.

- `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:3780` `set_default_device(bool gpu)` maps to `Device::gpu` (index 0).
- `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:18` `new_stream_on_device(bool gpu)` likewise.
- `src/lib/mlxcel-core/src/lib.rs:45,54,1708` expose only the boolean forms.

So even on a CUDA box with many GPUs, all work runs on the one default GPU.

MLX itself is fetched from upstream `ml-explore/mlx` at GIT_TAG `a6ec7123dac814417147e21d4aeed694924ddd4d` with a CUDA backend and custom CUDA patches (`src/lib/mlx-cpp/CMakeLists.txt:92-95`). `mlx::core::Device` carries a `DeviceType` plus an integer index, and MLX provides a `distributed` module (MPI/ring, and NCCL on CUDA), but it is process-based and not currently bound in mlxcel-core.

## Decomposition

This epic is delivered by two dependent sub-issues:

1. Sub-issue: device-index-aware device and stream bindings in mlxcel-core (the foundation that lets us target GPU 1..N and query device count).
2. Sub-issue: multi-GPU tensor-parallel runtime that places each rank's weights and compute on its own GPU and reduces partials with real cross-GPU collectives. Depends on (1).

## Sub-issues

### Phase 1
- [x] #487

### Phase 2
- [ ] #488 (depends on #487)

## Cross-cutting feasibility spike (gating)

Before committing to the intra-process design in sub-issue 2, verify that MLX's CUDA backend at the pinned commit can, within a single process: (a) allocate arrays on a non-zero GPU index, (b) run ops on a non-default device, and (c) move arrays between devices (peer copy). If intra-process multi-GPU is not supported by MLX as pinned, fall back to a multi-process design (one process per GPU on localhost/NVLink, reusing the existing distributed transport path), or evaluate bumping the MLX pin. The spike outcome decides sub-issue 2's architecture and must be recorded in that issue.

## Integration and validation (epic acceptance)

The epic is complete only when the capability is integrated end to end, not just present as modules:

- [ ] `gpu_device_count()` (from sub-issue 1) is surfaced at server/CLI startup (a diagnostic log line such as "detected N GPUs").
- [ ] `load_model_with_tensor_parallel` places shard r on GPU r; a memory probe shows each GPU holds approximately 1/N of the weights.
- [ ] `mlxcel generate --tp-size N` and `mlxcel-server --tp-size N` run on N GPUs and produce logits matching the single-GPU replicated reference within tolerance.
- [ ] A model that does not fit on one GPU loads and serves with `tp-size N` (the DGX goal), validated on real multi-GPU hardware.
- [ ] Clear, actionable errors when `tp_size > device_count()`, and a documented behavior on Metal/Apple (single GPU): `tp_size > 1` keeps the current in-process emulation or errors clearly.
- [ ] Docs updated (multi-GPU usage, flags, hardware requirements, limitations).

## Non-goals

- Multi-node scale-out already exists via pipeline parallelism over TCP/Thunderbolt/RDMA and disaggregated prefill/decode; this epic is about a single node.
- Metal/Apple Silicon multi-GPU is not applicable (one GPU per machine, unified memory).
- 2D pipeline-plus-TP composition is tracked separately.

## References

- Bindings device API: `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:18,3780`; `src/lib/mlxcel-core/src/lib.rs:45,54,1708`
- TP runtime: `src/distributed/tensor_parallel/llama_runtime.rs:76-148,2288`
- TP load wiring: `src/loading/mod.rs:431-476`
- MLX pin: `src/lib/mlx-cpp/CMakeLists.txt:92-95`
- CUDA feature: `Cargo.toml:48`

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.