compressed-tensors MoE decompression: weight_shape sharded away under TP/EP, fallback formula unrecoverable for 3/5/6/7-bit

[verma8076]

Context

PR #47315 fixed DecompressExperts.convert (src/transformers/integrations/compressed_tensors.py) to prefer the checkpoint's stored per-expert weight_shape over reconstructing the unpacked in-dim as packed_cols * (32 // num_bits), since that formula only recovers the true in-dim when num_bits divides 32 (4-bit, 8-bit). For 3/5/6/7-bit it's wrong.

That PR's own review thread already surfaced a follow-up gap that never got its own issue, so filing it here to make it trackable. Credit for the analysis below goes to @ErenAta16 in #47315 (comment) (their comment on that PR) — reproducing it here since it was never filed separately.

The gap

Under TP/EP sharding, weight_shape is empty on most ranks (only whichever rank owns full checkpoint metadata has it), so the fallback formula still runs there. Per @ErenAta16's analysis, this isn't just a formula bug that a better formula can fix — the packed shape is genuinely ambiguous once sharded:

• compressed_tensors.compressors.pack_quantized.helpers.pack_to_int32 packs 32 // num_bits values per int32 with padding (the "dense/native" layout).
• AutoAWQ-style checkpoints instead pack bit-continuously: ceil(cols * num_bits / 32) int32s, no padding.
• For in_dim=512, 3-bit: native packing gives packed_cols=52 (formula recovers 520, over-counts by 8), AWQ/dense packing gives packed_cols=48 (formula recovers 480, under-counts by 32).
• Given only packed_cols, the true in-dim can't be distinguished between these two layouts, and the native-layout padding means its exact value isn't recoverable even in principle (only bounded within pack_factor - 1).

So no formula fix is possible here — weight_shape is the only exact source, and it needs to actually reach every rank. It's a 2-element tensor, so per @ErenAta16 the fix is likely to stop sharding it (or broadcast it) rather than reconstruct it. In quantizer_compressed_tensors.py it's currently pulled in as a plain source pattern alongside weight_packed/weight_scale, so it inherits their sharding even though it carries no per-rank data.

Why this matters

This is precisely the deployment scenario that motivates low-bit MoE quantization in the first place — large MoE models needing multi-GPU TP/EP sharding. Today, non-4/8-bit MoE checkpoints under TP/EP still hit wrong decompressed shapes (and the resulting torch._grouped_mm contraction-dimension crash that #47315 was meant to eliminate), just not for the single-GPU case #47315 already fixed.

Suggested direction

Stop letting the small weight_shape metadata (2 ints per expert) get sharded away by the TP/EP loading path — broadcast it to all ranks or otherwise keep it outside the sharding split applied to weight_packed/weight_scale. Validating this needs a real multi-GPU box, which is why it wasn't included in #47315.

cc @SunMarc @zucchini-nlp — flagged in the original thread as having context on why weight_shape ends up empty per-rank.

[peterlodri-sec]

@verma8076 @SunMarc @MekkCyber — I'd like to pick this up. I dug through the packing layout and the TP/EP path and I agree with @ErenAta16's analysis: the fallback is genuinely unrecoverable for non-divisor bit widths, and I have a concrete plan.

Findings (verified against compressed-tensors source)

• weight_shape is a state-dict tensor written by the compressor at save time (torch.tensor(weight.shape)), not part of quantization_config. Upstream (PackedQuantizationCompressor.decompress / vLLM's WNA16 MoE) treats it as required — unpack_from_int32 cannot run without it.
• In the transformers loader, weight_shape is routed as a plain *_shape$ source pattern in the same bucket as weight_packed$/weight_scale$, so it inherits the same TP/EP sharded routing — even though it carries no per-rank data. update_tp_plan has no entry for it.
• There is no packing-layout flag in the checkpoint: the current pack_to_int32 is dense bit-continuous (ceil(cols*num_bits/32)); the padded 32//num_bits layout is legacy. Both share the same format name, so the ambiguity is real and unresolvable from metadata alone.

Proposed fix

1. Broadcast weight_shape (and the other per-rank-invariant metadata: weight_zero_point, weight_g_idx) — keep it out of the TP/EP split so every rank gets the (num_experts, 2) shape info. It's a tiny tensor; it should not ride the sharded split. This matches how vLLM/compressed-tensors consume it (required, never reconstructed).
2. Hard-error safety net in DecompressExperts.convert: when weight_shape is missing/empty and 32 % num_bits != 0 (3/5/6/7-bit), raise a clear error instead of fabricating a wrong in-dim that silently crashes later in torch._grouped_mm or produces garbage padding. For dividing bit widths (1/2/4/8) the fallback is exact and stays.

I have a local fix ready to implement. If this sounds good, I'll open the PR — happy to split either part into its own PR if you'd prefer. (Note: #47430 covers weight_zero_point collection for asymmetric schemes — no overlap with the broadcast/safety-net work here.)

[verma8076]

@peterlodri-sec sounds good to me, please go ahead. The broadcast fix and hard-error safety net split makes sense, happy to see it as two PRs if that's easier to review. Thanks for digging in on this, nice catch on the update_tp_plan gap and the packing-format ambiguity confirmation.