mlx tensor-parallel: replicate vector-quantized codebooks - #2268
Open
noahzelezny wants to merge 1 commit into
Open
mlx tensor-parallel: replicate vector-quantized codebooks#2268noahzelezny wants to merge 1 commit into
noahzelezny wants to merge 1 commit into
Conversation
VQ-quantized MoE experts store each expert tensor as codes + scales + a [K, d] codebook. The codes and scales shard correctly on the default axes, but the codebook is a shared lookup table indexed by the codes -- slicing it across ranks makes every rank decode against a fraction of the table and produce garbage. Guard both sharding closures so the codebook is replicated instead. Also adds model cards for three VQ builds of Qwen3.5-397B-A17B, which are what this was found on. They load on stock mlx-lm via the bundled model.py (config.json model_file); no engine changes beyond the above.
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
1 participant
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.
What
Vector-quantized (VQ) models store each expert weight as
codes+scales+ a[K, d]codebook. Under tensor-parallel sharding,codesandvq_scalesshard correctly on the default axes, but the codebook is a shared lookup table indexed by the codes — slicing it across ranks leaves every rank decoding against a fraction of the table, producing garbage output. This adds a guard to both sharding closures inauto_parallel.pyso any parameter path ending incodebookis replicated instead of sliced.The guard can't misfire on existing models: no parameter anywhere in current mlx-lm is named
codebook.Also adds model cards for the three VQ builds this was found on — Qwen3.5-397B-A17B-VQ-2.2bpw / 2.4bpw / 3.1bpw — which load on stock mlx-lm via the bundled
model.py(config.jsonmodel_file), so no engine changes are needed beyond this one.Why
I built these quants to make this model more accessible to the community. The 2.2bpw and 2.4bpw builds each run on a single 128 GB Mac (~101 / ~112 GiB on disk); the 3.1bpw quality build is what tensor-parallel across two Macs is for — which is where this guard matters. Verified serving across a 2-node ring: output identical to the single-box run.
The artifacts also ship the model's full vision tower at source precision as a top-level indexed shard, and the cards declare the
visioncapability —mlx-lmis text-only for this architecture, but exo's vision path loads the tower from the folder directly. Verified: an image query through/v1/chat/completionsagainst the 2.4bpw build on a live 2-node tensor instance returns a correct description.Testing
ModelCard.load_from_path(existing card tests pass).