Shard quantized expert companions under EP and TP
A stacked experts projection is a Parameter, not a Module, so its scale
grid, bias and NVFP4 globals are SIBLINGS rather than children: the plan
matcher's one module-level fallback never reaches them, and the module
entry shards nothing. Each companion has to be named in the plan or the
weight becomes this rank's expert slice while its scales stay whole.
EP had a hand-written rule for exactly one of them (`grouped_gemm` ->
`_scale_inv`) that this branch had dropped; restore it and generalise to
`_bias`, both NVFP4 globals and a static activation scale. The gate_up
input global stays replicated: one per-tensor value for the pre-routing
hidden states.
TP never had the equivalent rule at all, since it uses packed_colwise /
rowwise and that loop only matched `grouped_gemm`. Add it, plus a second
fix: packed_colwise splits the output axis as two packed halves, which is
the [gate; up] STACK, while a quantized module holds the kernels'
INTERLEAVED rows. Register param-only styles for the axis actually wanted
-- the built-in styles count from the END of the shape, landing inside a
swizzled scale grid's inner tile and on a bias's expert axis.
Verified on 2 GPUs against a pre-quantized DeepSeek-V3 fixture: EP and TP
logits match an unsharded reference, and removing the TP half makes that
test fail with the scale left whole against a sharded weight. 8-rank smoke
runs pass on DeepSeek-V4-Flash (EP) and DeepSeek-V3 (TP); device_map smoke
passes on GLM-5.2-NVFP4, DeepSeek-V3, MiniMax-M3-MXFP8 and OLMoE.
Also in this pass, from the second review:
- `to_local` deprecated rather than deleted; every call site is gone now
that `_hf_quantized_needs_local_tp` hands modules local tensors.
- `validate_environment` gains its first tests -- the only ones that
existed built the frozen FP8 quantizer, which no `from_pretrained` can
reach since AUTO_QUANTIZER_MAPPING["fp8"] became FineGrainedHfQuantizer.
- `_swizzles_scales` no longer claims block-FP8 never reaches a scaled
MMA; it does, via UE8M0, but broadcasts one 128-block scalar in register
so there is no per-group grid to lay out.
- modelopt remapped to nvfp4 at construction; glob skip-lists moved to
the shared `subtree_pattern_to_regex`; GPT-OSS carries swiglu_alpha /
swiglu_limit on its config.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
