Conversation
twallgren
force-pushed
the
pr3-uneven-tensor-parallel-sharding
branch
from
September 8, 2026 18:30
c586b30 to
ecb1014
Compare
Contributor
|
Read through the sharding approach: distributing whole KV-head groups so n_heads_local stays an exact multiple of n_kv_heads_local on every rank is the right invariant for GQA under uneven splits, and the zero-head rank check failing loudly matches how the pipeline side now behaves. I have not run this branch since it's gated on ml-explore/mlx#4465. An offer for when 4465 lands: I run a 4-Mac Thunderbolt/JACCL ring with genuinely uneven hardware (an M4 Max and an M5 Max with different memory headroom), which is the exact deployment this PR exists for. Happy to validate real uneven tensor-parallel inference on it, including memory-per-rank and throughput against the even-split baseline and against uneven pipeline splits from #1816. That's data a single-machine test can't produce. Ping me here when it's testable. |
computed from this rank's own size rather than a real prefix sum, silently omitting a layer whenever ranks have different sizes) with a cleaner API than what this branch originally proposed for that part -- dropped here in favor of it. deepseek_v32.py wasn't touched by that fix because it never used PipelineMixin at all -- it carried its own fully independent copy of the same buggy logic, unrelated to and unfixed by ml-explore#1816. Rather than patching that copy a third time, this converts DeepseekV32Model to actually inherit PipelineMixin (matching qwen3_5.py/ministral3.py), so it gets ml-explore#1816's tested fix directly instead of maintaining a parallel implementation of the same bug-prone logic. Same behavior-preserving transformation applied elsewhere: __call__ walks pipeline_layers instead of manually indexing layers[start_idx + i] over a hand-tracked num_layers, and the outer Model.layers property returns model.pipeline_layers instead of re-deriving the same slice by hand. Also added the same empty-cache guards used elsewhere in this codebase (cache and cache[-1] is not None, etc.) for a rank that ends up with zero local layers. qwen3_moe.py's sanitize() has an unrelated bug in the same problem space, not touched by ml-explore#1816: it gates MoE expert-weight-stacking on a fixed model.layers.0... key. Under pipeline sharding a rank only has its own local layers' weight files downloaded, so a rank not owning layer 0 sees that key absent and skips stacking entirely for its own local layers -- silently wrong for any raw (unconverted) checkpoint under pipeline parallelism. Fixed to detect and iterate by which layers are actually present in the local weights dict instead.
Depends on ml-explore/mlx#4465 ("Support uneven tensor parallel sharding in distributed linear layers"), which adds the sizes= parameter shard_linear() needs -- open, not yet merged. Every shard() in this diff currently requires num_heads/num_kv_heads to divide evenly by the number of ranks (floor-divides and discards the remainder, `n_heads //= N`), which fails outright for most real models at odd rank counts -- grouped-query attention has pushed KV-head counts down to 8 or 4, so e.g. N=3 breaks nearly everything off the shelf. Adds a small _rank_sizes(dim, N, block=1) helper to pipeline.py (used by every file below to compute per-rank sizes; unrelated to and doesn't touch pipeline.py's own pipeline() method) and ports a GQA-aware uneven-sharding fix (originally landed and verified on Qwen2) to llama.py, qwen3.py, qwen3_moe.py, glm4_moe.py, deepseek_v2.py, and deepseek_v3.py: distribute whole KV-head *groups* (each covering num_heads/num_kv_heads query heads) across ranks, rather than splitting q_proj's and k_proj's/v_proj's output features independently -- this keeps the local head-to-KV-head ratio exactly correct on every rank even when the counts don't divide evenly among each other. deepseek_v2.py/ deepseek_v3.py use MLA (multi-head latent attention) instead of GQA, which has no separate KV-head count at all, so heads there are distributed individually instead of in fixed-size groups; deepseek_v3.py's batched embed_q/unembed_out (MultiLinear modules, sliced directly on their head axis rather than through shard_linear) use the same cumulative, possibly-uneven boundaries. Verified live on a real 3-node cluster with genuinely non-divisible head counts -- the exact case this fixes -- across four of the six architectures with a model small enough to test (GLM-4.5-Air 96/8 heads, Llama-3.1-8B 32/8, DeepSeek-V2-Lite 16 heads MLA, Qwen3-30B-A3B 32/4): all produced coherent completions with evenly-balanced per-node memory. Not live-tested: plain (non-MoE) qwen3.py, no dense Qwen3 checkpoint small enough to download, and deepseek_v3.py, where real DeepSeek-V3 is too large for any cluster this size regardless of sharding. Went through 4 rounds of review (Claude Opus, fresh context each round).
This branch's old duplicate PipelineMixin-conversion commit reintroduced the cache and cache[0]/cache[-1] guards that were removed from deepseek_v32.py during ml-explore#1861's review (inconsistent with qwen3_moe.py's established PipelineMixin pattern). Match upstream/main exactly.
twallgren
force-pushed
the
pr3-uneven-tensor-parallel-sharding
branch
from
September 14, 2026 16:56
ecb1014 to
7ce3051
Compare
twallgren
marked this pull request as ready for review
September 14, 2026 17:15
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The issue this is addressing:
Tensor parallelism required the model's attention head count to divide evenly across the ranks.
The solution:
Adds support to more model files for uneven sharding so ranks can get different amounts of work and adds a small shared helper in pipeline.py used by all of them to compute the per-rank split sizes.
This depends on ml-explore/mlx#4465 so it should not be merged until that PR is merged.
I needed this change so I could use tensor parallelism with a wider variety of models on my 3 node cluster. I've been running it successfully for a while and testing it with a variety of models with great results.