Support uneven tensor parallel sharding in distributed linear layers - #4465
Open
twallgren wants to merge 1 commit into
Open
Support uneven tensor parallel sharding in distributed linear layers#4465twallgren wants to merge 1 commit into
twallgren wants to merge 1 commit into
Conversation
twallgren
marked this pull request as ready for review
September 4, 2026 22:08
twallgren
added a commit
to twallgren/PersonalSite
that referenced
this pull request
Sep 8, 2026
Adds a third retraction, which is the substantial change. A shared function every pipeline-capable architecture uses computed its layer partition by multiplying a per-rank size instead of accumulating an offset, so whenever the layer count did not divide evenly by the node count exactly one layer was loaded and then executed by nobody. Five of the eight models here were affected; the two that were ever correct were correct by arithmetic accident. It belongs with the retractions because what it falsifies is a claim the page kept making — that a coherent completion plus a three-way memory split proved the thing worked. Re-measures every model on the current harness from a clean cluster, and replaces the pipeline-vs-tensor comparison with figures taken after the fix. Qwen2.5-72B now appears in both modes for the first time, so the "tensor never wins" line is restated: nothing wins meaningfully there. Also: links the sharding patch to ml-explore/mlx#4465, records the wired-memory leak that only a reboot clears, ports the sharding pattern from one model file to seven, and restructures the page for length — a section index, twelve one-line rules over a linked appendix, and the failure log collapsed behind its own titles.
twallgren
added a commit
to twallgren/mlx-lm
that referenced
this pull request
Sep 8, 2026
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).
twallgren
force-pushed
the
uneven-tensor-parallel-sharding
branch
from
September 8, 2026 19:02
24acf94 to
f69125a
Compare
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.
This change adds support for uneven sharding so the sharded dimension doesn't have to divide evenly across all ranks.
I needed this change so I could use a wider variety of models on my 3 node cluster. I've been running it successfully for a while and testing it with Qwen2.5-72B-Instruct-4bit, GLM-4.5-Air, Llama-3.1-8B, DeepSeek-V2-Lite, and Qwen3-30B-A3B with great results.