huggingface / huggingface/candle
`candle-nn`: column- and row-parallel linear layers for multi-GPU inference
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Problem
There is no tensor parallelism in this crate. A model too large for one device
can be loaded layer-by-layer, or quantized until it fits, but it cannot be
*sharded* — and #1813 ("Quantized models on multi-GPU") has been open since
March 2024 without an answer.
The missing piece is small and old. Megatron-style tensor parallelism needs
exactly two layers: a column-parallel linear that splits its output features
across devices and concatenates, and a row-parallel linear that splits its input
features and all-reduces. Everything else — which layers to shard, how devices
are discovered, how ranks talk to each other — is the caller's, and should stay
the caller's.
The MLP seam merged for #3809 makes this the natural next step: a caller can now
substitute a block's feed-forward computation, but has to write the sharded
linear layers itself to have anything worth substituting. Tachyon-Mesh wrote
them, and they are not mesh-specific: they are two `Tensor` operations and a
device list.
## What the layers need to get right
1. **Refuse an uneven split.** `out_features` (column) or `in_features` (row)
not divisible by the device count must be an error at shard time. Silently
dropping the remainder produces a model that runs and is wrong, which is the
worst available outcome.
2. **Contiguity before the matmul.** `narrow`, `to_device` and `t()` all yield
non-contiguous views; the CPU matmul tolerates them and the CUDA one does
not, so the shards must be made contiguous inside the layer rather than by
every caller who happens to hit the GPU path.
3. **A single all-reduce per row-parallel layer**, not per shard — that
reduction is the whole communication cost of the scheme, and doing it
per-shard is what turns tensor parallelism into a slowdown.
4. **Degenerate cleanly to one device.** A single-element device list must
produce exactly the dense result, so the same code path serves both and a
caller does not branch.
5. **Leave the communicator to the caller.** Whether the all-reduce goes through
NCCL, through host memory, or through something else is a deployment
decision; the layer should accept a reduction strategy rather than assume
one.
## Proposed API
```rust
pub struct ColumnParallelLinear { /* … */ }
impl ColumnParallelLinear {
/// Split `weight` (`[out_features, in_features]`) into `devices.len()`
/// equal column shards.
pub fn shard(weight: &Tensor, devices: &[Device]) -> Result;
/// Run every shard and gather onto `gather_device`.
pub fn forward(&self, x: &Tensor, gather_device: &Device) -> Result;
}
pub struct RowParallelLinear { /* … */ }
impl RowParallelLinear {
/// Split `weight` (`[out_features, in_features]`) into `devices.len()`
/// equal row shards.
pub fn shard(weight: &Tensor, devices: &[Device]) -> Result;
/// Run every shard on its own slice of `x` and reduce onto
/// `reduce_device`.
pub fn forward(&self, x_shards: &[Tensor], reduce_device: &Device) -> Result;
}
/// Split a `[rows, in_features]` activation the way `RowParallelLinear`
/// expects it.
pub fn split_for_row_parallel(x: &Tensor, devices: &[Device]) -> Result>;
```
With these plus #3809's `BlockMlp`, a tensor-parallel Llama is a factory closure
and nothing else — which is how we run it today.
## Note
Deliberately *not* proposed here: cluster topology discovery, NCCL bootstrap,
NUMA binding, pipeline-stage transports. Those are deployment concerns and they
belong in whatever is orchestrating the devices, not in `candle-nn`. Keeping the
boundary there is what makes this two types rather than a subsystem.
I am happy to submit the PR, with the equality test we use — sharded output must
match the dense reference bit for bit on a real checkpoint, which is also the
only way to catch an uneven-split bug that a shape check misses.
Related: #3809, #1813.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in the candle-nn crate and read the MLP seam from #3809, then trace the Tensor operations used for narrow, transpose, contiguous views, matmul, device transfers, concatenation, and reduction. Implement the two proposed layers and split helper around the stated API, keeping communication supplied by the caller. Done means uneven splits error, one device matches dense output, and the sharded result matches the dense reference in the described equality test.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100