Finegrained kernels: unified quantization kernels api (separate quantizers) - #49015
Draft
IlyasMoutawwakil wants to merge 62 commits into
Draft
IlyasMoutawwakil wants to merge 62 commits into
IlyasMoutawwakil wants to merge 62 commits into
Conversation
The kernels now store gate|up row-interleaved (gate j at row 2j, up j at 2j+1)
rather than as two N-apart halves, so the integration follows:
- GPT-OSS already ships this order on disk, so its de-interleave converters go
away entirely: _deinterleave_gate_up_rows, its use in FineGrainedMxfp4Deserialize,
and the whole FineGrainedGateUpBiasDeinterleave op plus its registration.
- Checkpoints that ship gate_up stacked are interleaved once post-load by
interleave_gate_up_after_loading, ordered BEFORE swizzle_scales_after_loading
(the swizzle packs whatever row order it finds) and skipped for mxfp4. This
cannot live in conversion_mapping's MergeModulelist+Concatenate chain — that is
shared with non-finegrained MoE models.
- _apply_gate de-interleaves with a stride-2 split, mirroring the fused epilogue's
split_gate_up; the two must agree or fused and unfused stop being comparable.
- swizzle_scales_after_loading drops gate= and gates on the doubled extent
(n_rows % 128, i.e. N % 64), so GPT-OSS N=2880 now pre-swizzles instead of
falling back to an affine scale read.
Adds tests/kernels/test_finegrained.py (24 tests), covering the post-load interleave,
that mxfp4 is left untouched, dispatch/swizzle gating, and the MergeModulelist path.
tests/kernels/test_finegrained.py: 24 passed.
The kernels store gate|up row-interleaved, but ``transform_weights_for_mega_moe`` does its OWN gate/up interleave and so takes the stacked form. Interleaving those modules at load and undoing it at the boundary is a pointless round trip — and interleaving twice is not the identity, it is a different permutation, i.e. silent garbage rather than a crash. So modules bound for megamoe are simply never interleaved: the post-load pass skips them and records what the module actually holds in ``_gate_up_interleaved`` (the bytes cannot be inspected to tell, and ``set_experts_implementation`` can switch backends after load). ``setup_megamoe_weights`` reads that flag and raises if it is ever handed an interleaved module, rather than re-permuting on an assumption. The other DeepGEMM expert paths need no change: their grouped GEMM is layout-agnostic and the gate split happens in ``_apply_gate``, which already de-interleaves stride-2. Verified: triton backend interleaves and flags; megamoe backend is left untouched with no round trip. The raise itself is unverified here — DeepGEMM's JIT needs a CUDA toolkit >= 12.9 and CUDA_HOME is unset on this box, so no megamoe path executes.
# Conflicts: # src/transformers/integrations/hub_kernels.py # src/transformers/quantizers/auto.py
The memoized `(epilogue, gate_up_quantization, down_quantization)` triple bought nothing where it matters. Constructing the three dataclasses measured ~1.4us per layer (~85us per step at 61 layers), and only in EAGER decode — under cudagraphs the host path does not run on replay, which is the mode this integration deploys in. Against that it held hidden mutable state on the module, and cached a gate_up quantization carrying an `output_recipe` that silently mismatches any caller asking for the unfused form. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The kernels now take a per-expert bias as a `bias=` operand, so a biased model no longer falls back to the unfused two-GEMM form just for having a bias. `_kernel_epilogue` stops bailing on `has_bias`, and both host-side adds go away: `_apply_unfused_gate_up` keeps only the activation, `_finish_down` only the routing-weighted reduce. That also drops the `torch.sort` the grouped path paid to index the gate_up bias by expert-sorted row -- the kernel indexes by the tile's own expert id instead. An unsupported act_fn no longer costs the bias its fusion either: the bias is an operand beside the epilogue rather than part of it, so it fuses whether or not the GLU does, and `epilogue is None` goes back to meaning exactly one thing. The bias rides the same output-row axis as the weight and scale grid, so `interleave_gate_up_after_loading` already delivers it in the kernels' interleaved order -- no reordering here. Tests assert the operand reaches both GEMMs, plus a new case pinning that an unfusable activation still fuses its bias. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Conflict resolutions: - quantizers/auto.py: keep the finegrained claims on fp8/nvfp4/modelopt (pre-quantized + MoE experts; NVFP4HfQuantizer only quantizes on the fly), take main's new gguf entry. - integrations/__init__.py: union — the finegrained module plus main's FP8Embedding exports. - integrations/deepgemm.py: ours (to_local unwrap + the stacked-gate guard); re-add the to_local import main's shim conversion dropped. - integrations/tensor_parallel.py: main's shim, with to_local restored on it — the kernel integrations still import it from this path. - distributed/sharding_utils.py: re-apply the 0-dim scalar guard (ModelOpt per-projection weight_scale_2/input_scale) inside the rewritten DtensorShardOperation.shard_tensor: replicated on the dense path, expert-ownership-filtered on the MoE path. - tests/tensor_parallel: main's rewrite, with the 0-dim sharding test re-added against the new engine (passes). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e kernels' MoE chain
The experts forwards are adapters over the kernels' `moe_fused_batched/grouped`: `_moe_operands`
hands them the module's tensors and the activation — a `get_supported_act_fns()` name the gate_up
epilogue fuses, else the module's own GLU as a callable the kernels run on the host, so any
activation works without a kernel release. The kernel bundle is loaded by symbol name
(`matmul_2d/batched/grouped`, the two forwards, swizzle/unswizzle, `get_supported_act_fns`).
Every layout difference between a checkpoint and what the modules hold is a `ConversionOps` with a
reverse, attached by the quantizer, so `save_pretrained` restores the checkpoint bitwise:
- `FineGrainedInterleaveGateUp`: stacked [gate; up] rows -> the kernels' [g0, u0, ...] order,
skipped when the model declares `is_concatenated=False` (GPT-OSS; the flag travels through
`use_experts_implementation`, which stamps its defaults after `__init__`) or the backend packs
gate|up itself (DeepGEMM Mega MoE);
- `FineGrainedScaleContainer`: a scale into the dtype its module holds — the same bytes for a
uint8 UE8M0 container (MiniMax), an exact cast for float32 values (dsv4-flash-base) — and the
module records the shipped container so the reverse restores it;
- `FineGrainedSwizzleScales`: the SWIZZLE_32_4_4 artifact for modules that hold 5-D scales
(`FineGrainedExperts.__init__` allocates that shape for every non-DeepGEMM backend on SM100
with a group-scaled format and activation quant; whole 128-row/4-col blocks only);
- `FineGrainedPackedBlocks` regroups GPT-OSS `{proj}_blocks` into the packed weight, so the
blocks/scales format is two one-to-one converters and round-trips too.
Catch-all converters cover keys that already arrive under the fused names and dense scale keys.
The post-load hooks, `_save_to_state_dict` override and post-load dtype cast are gone.
Also: the eager per-expert loop reads a swizzled stack's `[e:e+1]` slice and passes bias, NVFP4
global and activation format (it dropped all three); expert biases take the model dtype; the
`mxfp4` format pins UE8M0 scales; `FineGrainedEmbedding` (FP8 table, per-tensor scale) for
`modules_to_convert`; DeepGEMM backends refuse modules holding swizzled scales or interleaved rows
before loading the kernel; `disable_deepgemm_on_multi_device` is public.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…heckpoint precision for quantized params `FineGrainedQuantize` reads the module's weight format: block-FP8 in torch (per-block E4M3, fp32 or power-of-two UE8M0 inverse scales, from the module's block and held scale dtype), MXFP8 / MXFP4 / NVFP4 through the kernels' row-wise quantizers in one launch per tensor (NVFP4 normalized by the per-tensor / per-expert global `amax / (6 * 448)`), and emits the scale (and global) in the layout the module holds — the container dtype, the swizzled artifact — since the loader runs the quantization op after the converter ops. Core loader: a parameter the quantizer's op is about to quantize is materialized in the checkpoint's precision instead of the empty parameter's storage dtype (int8 storage zeroed FP4 weights; float8 double-rounded the existing FP8 path). Also: `auto.py` groups the finegrained keys under one comment and drops unused NVFP4 imports; the quantizer's config type hint names `FineGrainedConfig`; DeepGEMM guards (`_assert_affine_scales`, `_assert_stacked_gate_up`, group-32 divisibility) run before the kernel load; expert biases take the model dtype; comment trims for the noisy-comments check; tests for every path. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ion/Epilogue bundle fields The kernels' dispatchers now take input_recipe/output_recipe strings, so the bundle no longer needs the Quantization and Epilogue config classes; the linears hand the module-level activation_format straight through as input_recipe and the _kernel_quantization adapter is gone. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…uantization/Epilogue bundle fields" This reverts commit 07dae50. The kernels keep the Quantization/Epilogue op-boundary configs as the dispatcher API, so the bundle carries them and the linears build kernel.Quantization again. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…kip the kernels arch check The kernels now speak the module's vocabulary: activation_format (None = the weights' format, "bf16" = weight-only) goes straight to matmul_2d / matmul_grouped / the MoE forwards, so the Quantization/Epilogue bundle fields and both adapter helpers are gone. hub_kernels: kernels >= 0.16 checks a build's declared archs against the device; the deep-gemm build declares only 9.0a although it JIT-compiles per device, so the mapping opts out via a new check_arch passthrough (only forwarded when False, so older kernels releases without the keyword keep working). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…(pr-1018) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…s per-expert output norm modelopt NVFP4 checkpoints ship a second-level global per projection and a calibrated `input_scale`; both now reach the kernels, in either of the two layouts such a checkpoint uses — one tensor per expert per projection (GLM-5.2) or one stacked tensor per layer (the fused vLLM layout). A gate|up stack arrives with two weight globals per expert, calibrated separately. One conversion op owns every global of a layer and merges them to the one-per-expert the kernels take: the stack keeps the gate's, and the up half's folds onto the down projection, whose weight global scales the expert output back up and whose calibrated input global moves the other way, keeping the requantized intermediate on the range the checkpoint calibrated. That is exact arithmetic on fp32 scalars, unlike rescaling the up half's e4m3 block scales. Many-to-many converters are now opt-in per op (`ConversionOps.supports_many_to_many`) rather than a private whitelist. A model states its per-expert output norm the way it states its activation: `use_experts_implementation(post_expert_norm="<name>")` plus an `_apply_post_norm` the class must define. A name the kernels implement rides as that name and is folded into their reduce; anything else is the module, called on the routed rows. The DeepGEMM experts paths, which have no slot for such a norm, refuse instead of dropping it. Weight-only modules hold no activation global, since nothing there quantizes activations: the parameters are not allocated, the converters for them are not emitted, and a calibrated checkpoint's `input_scale` keys are ignored for that run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
One guard decorator on the DeepGEMM experts forwards instead of four statements per body, and that file's comments cut to what is not already in the code. `prefers_deepgemm_linear` moves to `deepgemm.py`, where its measured verdicts belong, so the dispatcher just asks. `_FineGrainedModule.local(name)` is the one accessor every forward reads operands through, so the DTensor unwrap appears once rather than at seven call sites. `_WeightFormat` is public (it is in `resolve_weight_format`'s signature), the `TRANSFORMERS_FINEGRAINED_NO_SWIZZLE` debug hatch is gone (tests patch the predicate), and `_set_optional_parameter` replaces six register-or-assign blocks. `check_arch` is passed straight through: `kernels` is pinned >= 0.16, where it landed, so the compatibility shim was dead. The scalar-shard branch becomes `_owns_expert`, stated as the general fact it is rather than as an NVFP4 anecdote. `FineGrainedConfig`'s docstring opens with a format table. The frozen fp8 modules' `stacklevel=2` is dropped: at module scope it named importlib rather than the module being deprecated. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…they served Every checkpoint `quant_method` the two answer to — `fp8`, `mxfp8`, `mxfp4`, `nvfp4`, `modelopt` — already routes to `FineGrainedHfQuantizer`, so both carry the same DeprecationWarning the fp8 pair does. Auditing that supersession found one real gap. `dequantize=True` on a GPT-OSS MXFP4 checkpoint produced no converter at all: the `_blocks`/`_scales` keys were wired only for the quantized path, so they went unmatched, where `Mxfp4Config(dequantize=True)` had handled them. The finegrained chain now dequantizes them, bit-equal to `mxfp4.convert_moe_packed_tensors` — the implementation it replaces — which the new test compares against directly. Two shared ops needed repair to get there: `FineGrainedPackedBlocks` regrouped a sibling `_scales` entry as if it were blocks, and `FineGrainedDequantize` left the scale it had consumed in the chain. The same path on NVFP4 or modelopt would have folded one block scale and dropped the second-level global, handing back a weight scaled by `1 / global`. Neither frozen integration supported that either, but it failed quietly — including where `dequantize` turns itself on (no GPU, compute capability < 8.9). It raises now. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # src/transformers/integrations/deepgemm.py # src/transformers/integrations/moe.py
It had become `FineGrainedFP8Config = FineGrainedConfig`, so the name still imported but the class was gone. That makes "frozen for backward compatibility" a fiction — the old name silently gained the new behaviour (`activation_format`, the modelopt payload parsing, the wider `post_init`) — and it rendered the same class under two `[[autodoc]]` headings. Restored verbatim, with a docstring line pointing at the config that serves the whole family. Splitting them means the two `isinstance` sites in `auto.py` no longer cover the new config by accident. `LOADING_ATTRIBUTES_CONFIG_TYPES` is the load-bearing one: it is what carries `dequantize` / `modules_to_not_convert` from a config passed to `from_pretrained` onto the checkpoint's own, so leaving it alone would have quietly stopped `FineGrainedConfig(dequantize=True)` from taking effect on every fine-grained checkpoint. Both configs are in it now. `quant_method="fp8"` still builds `FineGrainedConfig`: the frozen class is what you get by naming it, not what a checkpoint resolves to. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`c66564416e` restored the class and added a line saying it "configures the frozen `finegrained_fp8` integration". That is not true: it sets `quant_method="fp8"`, which `AUTO_QUANTIZER_MAPPING` routes to `FineGrainedHfQuantizer` — the frozen quantizer is not reachable through `from_pretrained` at all. Removing it also leaves the class byte-identical to main, which is what a class frozen for backward compatibility should look like in a diff. Which config serves which family is already stated where a reader meets it, in `FineGrainedConfig`'s own docstring on the same docs page. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three CI failures on a CPU-only runner, three causes: `post_expert_norm_name` is set by `use_experts_implementation`, but the shared forwards are the ones every backend adapts onto — including duck-typed stand-ins built without the decorator. The norm is optional, so read it as one. `_assert_stacked_gate_up` already read its second attribute defensively; `has_gate` now matches. `prefers_deepgemm_linear` calls `is_sm100()`, and `get_device_capability()` asserts rather than answering False on a CPU-only build. Caught there rather than gated on `is_available()`, so that a test faking the capability alone still answers — which is the convention the existing deepgemm tests are written to. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A converter whose target no module holds is a load failure, and the activation global is exactly that under a weight-only run: the experts allocate one only when they quantize activations, so the declarations are the only thing standing between the two runs. Checked against a real module's parameter names in both. Verified it catches the failure by removing the conditional: the weight-only run then converts `gate_up_proj_input_global_scale` onto a module without that slot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The integration reads better as `is_available() and capability >= 10`; the reason it wasn't was that three test sites faked the capability alone, which that form ignores. That is the tests constraining the integration, so the tests move instead: they now patch availability alongside the capability, and the module docstring says so. Replaces the try/except that was catching what the availability check now prevents. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… name A sharded parameter is already local by the time a forward runs, so the unwrap was doing nothing: FSDP2 unshards in its pre-forward hook, and a TP/EP plan swaps each DTensor for its local shard around the call (`MoeExpertsParallel` -> `_use_local_dtensor_params`). Verified both: under `fully_shard` alone a param reads as a plain Parameter inside forward, not a DTensor. So the invariant belongs to the distributed layer that wrapped the parameter, not to every backend that reads one. That removes `to_local` and the `local()` accessor. Absent slots still read back as None because `_set_optional_parameter` registers them as None parameters, which is what the kernels take for a missing bias or global. `_moe_operands` also stops building its dict through a `both(suffix)` helper: the kernels' arguments are now written out, so `down_proj_scale_inv=` can be found by grep. The five `getattr`s that remain are load-bearing — the kernels always say `gate_up_proj`, the module says `up_proj` when the model has no gate. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two branches decided behaviour by matching a NAME, which is how the Qwen3 plan override stayed invisible — the default arm silently absorbs a miss: - the grouped-linear swap tested `"GroupedLinear" in type(module).__name__`. It now tests `hasattr(module, "n_groups")`, the attribute the branch consumes, so a grouped linear named anything else is still swapped and one named right but shaped wrong no longer reaches an AttributeError. ConvBERT's `GroupedLinearLayer` was never caught (it is an `nn.Module`), so this is hardening rather than a fix. - `_global_role` classified with `"w2" in key`, an unanchored substring for what is a path SEGMENT. Since the default is gate_up, a stray match merged a down global into the wrong bucket silently. Matched on segments now. Readability: - `FineGrainedQuantize` and `FineGrainedDequantize` still carried their own `__init__`; the base-class collapse only matched the form with a default. - `_weight_holder`'s pair was indexed as `holder[0]` / `holder[1]` in four places, and is now unpacked into `module, scale_name`. - `_recontain` -> `_as_container` and `_swizzles_scales` -> `_holds_swizzled_scales`: both were coined, and the codebase already says "container" (`scale_container_dtype`) and "holds" (`holds_interleaved_gate_up`). - `q` / `s` -> `blocked` / `per_block` in the dequant reshape. - each file opens with a short statement of what it is and where its neighbour lives, rather than none at all. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two axes that were tangled into one `quant_method`. GROUPS. A checkpoint can be more than one format — DeepSeek-V4 is W4A4 mxfp4 experts over block-FP8 linears — and a flat config cannot say so, which is why its expert format arrives on the MODEL config as `expert_dtype` and is read back out deep inside `FineGrainedExperts.__init__`. `FineGrainedGroup` names a format per module subset; `group_for(name)` resolves one; `replace_with_finegrained_layer` asks per module, so the side-channel read is gone. Nothing about this is DeepSeek-shaped. Verified across three models and three combinations, swap and in-flight quantization both: Qwen3-MoE with nvfp4 experts over mxfp8 dense, and GPT-OSS with nvfp4 attention, static block-FP8 experts and mxfp8 for the rest — three groups, the exotic format on the dense side, and differing activation SCHEMES per group. Backwards compatibility is normalisation, the same move `quant_algo` already gets: the flat fields become a single group, and a legacy `expert_dtype` becomes two (`split_out_experts`). Existing checkpoints load unchanged and nothing downstream learns which spelling it came from. `config_groups` is normalised too — GLM-5.2-NVFP4 ships it and we were dropping it, along with the `input_activations` and `dynamic: false` it carries. A group naming a format we do not serve falls back to the flat fields rather than guessing at a partial translation. Resolution is order-INDEPENDENT: `to_json_string` sorts keys, so a config that round-tripped through `config.json` has lost its declaration order. Targeted groups win over the catch-all, and two targeted groups claiming one module raises. ARMS. Split per PRODUCER key layout rather than per format, which is what actually differs: GPT-OSS's packed `_blocks`, modelopt's two-level scales, the calibrated `input_scale`. MXFP8 ships the plain `weight`/`weight_scale_inv` pair and gets no arm — the base is its handler. The base keeps no `_quant_method() == ...` test at all. Also: a finegrained module holding a full-precision weight now fails the load instead of silently computing in bf16 and handing back a model the caller believes is quantized. Checked once in the post-load pass, not per forward — the dtype cannot change between them. No shipping checkpoint relied on the old fallback: gpt-oss-120b lists attention in `modules_to_not_convert`, so those modules are never swapped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
vasqu
self-requested a review
September 22, 2026 13:12
… simplification The three defects found on `finegrained-unified` were all in code this branch shares, so they were all live here -- and the middle one bites harder, because this branch has no `element_size() > 1` fallback and `assert_modules_are_quantized` turns it into a failed load rather than a silent bf16 forward: - the experts' static `activation_scale` was allocated without an explicit dtype and came out bf16 under `from_pretrained`; - it was never WRITTEN. The loader materialises a key no checkpoint supplies with `torch.empty_like` and `_init_weights` has no branch for a scale, so it reached the kernels as uninitialised memory -- a zero divides the activations by zero; - `_quantize_block_fp8` refused a shape the block does not divide, though the format's own producers do not: DeepSeek-V3 ships `kv_a_proj_with_mqa` as (576, 7168) E4M3 against a 128x128 block with a (5, 56) grid. It pads now. That refusal was also the filter keeping 2-D expert biases out, so they are turned away against `_weight_holder` instead. With shapes no longer refused, `assert_modules_are_quantized` stops being a trap: what is left for it to catch is a checkpoint that targets a module and ships no scales, which is the thing it was written for. The swap simplification lands here too, so the two branches stay comparable: `FineGrainedGroupedLinear` takes the same storage mapping as the dense linear (it was hardcoding fp8), and `_quantized_experts` lifts 27 lines out of the three-way dispatch. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
vasqu
reviewed
Sep 22, 2026
vasqu
left a comment
vasqu
left a comment
Collaborator
There was a problem hiding this comment.
Sorry it is a long one but I think we are getting close. My main concern on the configs is done now and this was a major blocker to me
It is quite a lot at once tho so bear with me. And the claude comments can get a lot lol, would appreciate if we could trim these down
| if tied and "embedding_rowwise" in self._tp_plan.values(): | ||
| head = self.get_output_embeddings() | ||
| if head is not None: | ||
| self._tp_plan.setdefault(next(n for n, m in self.named_modules() if m is head), "colwise_rep") |
Collaborator
There was a problem hiding this comment.
need to recheck but i was under the assumption that this was already automatically done cc @3outeille
Member
Author
There was a problem hiding this comment.
leaving this for @3outeille to confirm. FSDP resolves the tied-head pair; TP had no equivalent, which is why init_parallel_plans gives it colwise_rep here.
Member
There was a problem hiding this comment.
every models declares the lm_head plan in their ForCausalLm cf https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py#L586 so this is not necessary
Member
Author
There was a problem hiding this comment.
ah okay i see thanks ! that's the better way of doing it, but it's missing on composite models, i used Glm4vMoeForConditionalGeneration for some tests and it had this problem, i can add _tp_plan = {"lm_head": "colwise_gather_output"} to Glm4vMoeForConditionalGeneration only but other composite models have the same problem ? lmk what you think
Member
There was a problem hiding this comment.
can you elaborate on what you mean by composite models ? model that have multiple modality mixed in ?
Member
Author
There was a problem hiding this comment.
yes models like vlms (ForConditionalGeneration)
Structure, as asked twice in the thread:
integrations/finegrained/{__init__,core,conversions}.py
quantizers/finegrained/{__init__,base,blockfp8,mxfp4,mxfp8,nvfp4}.py
`quantizer_finegrained_fp8.py` stays where it is: that is the deprecated
`FineGrainedFP8HfQuantizer`, which this PR does not own. Each package re-exports
its public names, so `_import_structure` and every `from ...finegrained import X`
are untouched; the tests name `finegrained.core` directly, because they patch
module internals and a re-export cannot stand in for that.
`_apply_post_norm` is gone. `post_expert_norm` is an `nn.Module` and so was
already callable, no backend overrode the hook, and `_apply_gate` one branch
above was already DEFAULTED where this one raised `TypeError` -- the asymmetry
was the tell. Every call site now calls the norm, and the callable handed to the
kernels for an unfusable norm is the norm itself.
Comments: the blocks marked "suggested change" are deleted rather than shortened,
as are the ones called a nothingburger and self-explanatory, plus three more that
only restated their own line and two that were teaching `nn.Module.__setattr__`
and `Tensor.contiguous`. The rest are trimmed where the thread asked for trimming.
What survives says something the code cannot: why a lookahead excludes the expert
keys, why each companion needs its own plan entry, why a `shared_experts.up_proj`
is skipped despite matching the word.
Smaller ones: `_fused_experts_forward` removed and both arms spelled out; two
`for name, tensor` loops unrolled; the gate_up/up name resolved once as
`gate_up_name`; `is_concatenated` moved in with the layout metadata; the
`input_global_scale` lookup collapsed from three copies of one condition to one
(weight-only never registers the slot, so the `bf16` check was testing a fact the
parameter already encoded); helpers hoisted to the top of `conversions.py`;
`as_container` and `group_from_config_groups` made public; `rpartition(".")[-1]`;
`chunk` for the gate|up split; `_CT_FORMATS` and `_MODELOPT_ALGOS` folded into
their one reader; `_read_modelopt` and a general `subtree_patterns` split out;
`split_out_experts` replaced by a pure `groups_with_expert_dtype`;
`supports_post_expert_norm` and a `supports_dequantize` bool; one named quantizer
per key in `auto.py`; four import-time `warnings.warn` shims moved to
`logger.warning_once` at construction.
Tests: the two `core_model_loading` fp8 tests now exercise this integration
instead of the frozen module, which surfaced three bugs in them -- a
`PermuteForRope()` built with no `permute_layer_names` (it iterates them
unconditionally, so it raised on any input, and every shipped call site passes
`["q_proj", "k_proj"]`), an expectation that `v_proj` is rope-permuted, and
state-dict keys prefixed `model.` against a model that never nested that way.
The float-path test keeps its bit-equality guard but against OUR formula: we
store `amax / MAX` and divide by it, where the frozen path stores
`1.0 / (MAX / amax)` and lands off `amax / MAX` on 51 of 200 random blocks. Real
DeepSeek-V3 agrees with ours -- recomputing `amax(q * scale) / 448` on a block
reproduces its shipped `weight_scale_inv` exactly.
`TestZeroDimParameterSharding` becomes two methods on `TestDtensorShardOperation`
beside the other `shard_tensor` tests, using their builder. It needs no ranks and
no devices, so behind `@is_tensor_parallel_test` it was skipped in CI outright.
The triton skip on the qkv test goes too: that path is pure torch.
Not done, and answered on the thread instead: no new `QuantizationMethod` member.
`quant_method` is at once the quantizer key, the config key and the weight FORMAT
fed to `resolve_weight_format`, so a `FINEGRAINED` member has no format to resolve
to; splitting those two meanings is its own change. Separately, pointing `fp8` and
`mxfp8` at the new arms leaves `FineGrainedFP8HfQuantizer` orphaned -- defined,
never imported, never mapped -- so nothing mis-routes and what remains is dead
code someone else should decide about.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
6 tasks
`FineGrainedGroupsTest` was written against the unified branch and never came across, so `groups` -- the thing this PR is for -- had no unit test here at all: nothing touched `FineGrainedGroup`, `group_for` or `group_from_config_groups`, and only the seven-model load-path run would have caught a regression. Six tests now cover the flat-config normalisation, the legacy `expert_dtype` split, the `config.json` round trip (where `to_json_string` sorts keys and a catch-all that won over a targeted group would swallow everything), an ambiguous claim, a producer's `config_groups`, and a format we do not serve. `FineGrainedGroup` was not exported, so a caller could not build what `groups=` asks for without reaching into `transformers.utils.quantization_config` -- and the autodoc entry beside `FineGrainedConfig` would have failed the docs build. The two loader compile-safety tests come over from the frozen suite: the opaque loader node must return None under `torch.compile`, cold and warm, and nothing here was checking it. Its sentinel and in-place-bias tests do NOT come over -- that glue moved into the kernels package when the experts forwards became a single call, so there is no longer anything on this side to test. Also two `PLW1514` fixes CI catches and a per-file `ruff check` does not. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…n files
Stacked on the review PR so that one stays reviewable. The frozen modules stay
where they are and keep working -- this only moves the callers over and says
loudly that the files are on the way out.
Deprecation moves from the CLASS to the FILE. A module-level `DeprecationWarning`
fires when anything is imported out of `integrations/{finegrained_fp8,mxfp4,nvfp4}`
or `quantizers/quantizer_{finegrained_fp8,mxfp4,nvfp4}`, which is the scope that
matches the intent -- the whole file is frozen, not one config class. That is only
safe because `_LazyModule` never executes these unless someone asks for them:
`import transformers` does not reach them, nor does `from transformers import
Mxfp4Config` (the config lives in `utils/quantization_config`), nor does
`integrations.compressed_tensors`, which has its own `CompressedTensorsFP8Linear`.
All four verified. The per-class notices go, and `FrozenFp8ShimTest` flips to the
new contract.
Callers, everywhere the frozen path was still in use for a model this integration
serves:
- `tests/models/{deepseek_v4,glm_moe_dsa,gpt_oss,qwen4_exp}` build `FineGrainedConfig`.
The qwen4-exp embedding test is REPLACED rather than duplicated -- the frozen
path keeps its own coverage in `tests/quantization/finegrained_fp8`, and a second
copy of it in a model file was only ever testing the deprecated route. It also
gains the case that the table is FP8 under any format, which used to be silent.
- the ministral3 and mistral4 conversion scripts. Both already built a
`FineGrainedConfig` -- `"fp8"` maps to it now -- and then swapped modules with the
frozen `replace_with_fp8_linear`, so they wrote checkpoints whose config said
fine-grained while the modules were `FP8Linear`. Now both halves agree.
- the docs: the model pages, `experts_interface`, and a warning at the top of each
frozen method page pointing at `FineGrainedConfig`.
Not done here: `quantization/finegrained.md` does not exist. The three frozen
methods each have a user guide and the integration replacing them has none, so
the warnings point somewhere there is nothing to read.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`utils/check_noisy_comments.py` caps a block at 5 lines and 500 characters, and these were 10/679 and 6/504. Both keep the part that is not in the code: why the sharded comparison uses each model's measured floor rather than a fixed tolerance, and why the tied-head check reads the config flag instead of testing `head.weight is embed.weight` (tying happens later, in `post_init`). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`utils/check_noisy_comments.py` caps a block at 5 lines and 500 characters, and these were 10/679 and 6/504. Both keep the part that is not in the code: why the sharded comparison uses each model's measured floor rather than a fixed tolerance, and why the tied-head check reads the config flag instead of testing `head.weight is embed.weight` (tying happens later, in `post_init`). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`init_parallel_plans` runs from `post_init`, so it is on the constructor path of every model in the library -- and it called `config.get_text_config()` bare. That RAISES on a composite model carrying more than one text sub-config: musicgen has `text_encoder` and `decoder`, so every musicgen and musicgen-melody test died before doing anything, along with `TestGetEncoder`. Falls back to the top-level config when the call cannot disambiguate, which is the same answer for `tie_word_embeddings` in every case that used to work. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`init_parallel_plans` runs from `post_init`, so it is on the constructor path of every model in the library -- and it called `config.get_text_config()` bare. That RAISES on a composite model carrying more than one text sub-config: musicgen has `text_encoder` and `decoder`, so every musicgen and musicgen-melody test died before doing anything, along with `TestGetEncoder`. Falls back to the top-level config when the call cannot disambiguate, which is the same answer for `tie_word_embeddings` in every case that used to work. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
vasqu
reviewed
Sep 23, 2026
vasqu
left a comment
vasqu
left a comment
Collaborator
There was a problem hiding this comment.
Core looks good! Thats on my side but im too tired to check the tests atm, will check tomorrow
| post_expert_norm, norm_weight, norm_eps = None, None, 1e-6 | ||
| if module.has_post_expert_norm: | ||
| fusable = module.post_expert_norm_name in kernel.get_supported_norms() | ||
| if not fusable or getattr(module.post_expert_norm, "_hf_tp_input_reduce", False): |
Collaborator
There was a problem hiding this comment.
ah ok I already forgot what I did myself lol 🤦 yea that makes sense
| @@ -26,6 +26,11 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): | |||
|
|
|||
| def __init__(self, quantization_config, **kwargs): | |||
| super().__init__(quantization_config, **kwargs) | |||
| logger.warning_once( | |||
Collaborator
There was a problem hiding this comment.
yea answered on the other file but agree its probably better to be aggressive here
| groups: dict[str, FineGrainedGroup] | None = None, | ||
| **kwargs, | ||
| ): | ||
| self.quant_method = kwargs.pop("quant_method", QuantizationMethod.FP8) |
Collaborator
There was a problem hiding this comment.
arg yea fair enough, its just a bit weird no? since its now a unified entry point and no just FP8
just thinking whether this convention might not be confusing again
Collaborator
There was a problem hiding this comment.
I feel like we have a lot of conversions we might want to check as per tests like this
I would potentially move these into their own file tho so that the fp conversions live on their own and we can block the whole class then (would also move that one fp8 conversion out to there then)
Conversions: * `_dequantize_one` inverted the scale grid to recover the block, which is only valid when the shape tiles. DeepSeek-V3's `(576, 7168)` against a `(5, 56)` grid raised, and `(192, 256)` silently used block 96 in place of 128. It now takes the block from the config when that config's ceil-grid reproduces the scale it was handed, and pads-then-crops the way the quantizer already did. * helpers used across ops are public (`keyed_by_target`, `held_scale`, `as_container`); `_global_role` is a static method of the one op that reads it. Tests move to `tests/integrations/finegrained/`, next to the integration they cover rather than under `tests/kernels/`, split by what each file answers: `test_core` what the integration passes, `test_conversions` what it loads, `test_forwards` what the kernels return, `test_models` what it shards. The frozen fp8 integration's tests follow their module to `tests/integrations/`. `test_forwards` ran nowhere it should have: `@require_torch_gpu` plus a capability check kept it off Hopper, where these kernels work, and off XPU, which the integration supports; and `setUpClass` turned any `ImportError` into a skip, so a broken loader or a missing build passed silently. It is `@require_kernels` + `@require_torch_accelerator` now, device-agnostic, with the load unguarded. Also: `_CATCH_ALL` inlined, `_owns_expert` folded into the placement check it duplicated, the wording the review suggested, and a note in the quantization docs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
One PR rather than a stack. The only conflict was `tests/kernels/test_finegrained.py`, deleted on one side (the tests moved to `tests/integrations/finegrained/`) and comment-trimmed on the other; the trim is already in the moved copy, so the deletion stands. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The review asked to shorten one four-line comment and I shortened only that one. The same cut applies to the three others this branch adds, and to the block in `sharding_utils` that two new lines had grown to four. The two `# noqa: E731` lambdas in the plan tests are plain functions now, which also settles their disagreeing parameter names. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…tures `tests/kernels` had no `__init__.py`, so pytest put the directory on `sys.path` and a bare `from test_utils import ...` resolved. Moving `test_finegrained_fp8.py` to `tests/integrations`, which IS a package, took that away and CI failed to collect it. Most of `tests/` is packages already, so make this one too and import through it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`init_parallel_plans` gave a tied head a style when the config had sharded the embedding it reads. Every model is expected to declare that itself in its `ForCausalLM` -- 81 already do -- so the hook only papered over the composites that keep their plan on `text_config`, and it had to reach for `get_text_config()`, which raises on a config with two text sub-configs. Of this branch's load-path fixtures only GLM-4v-MoE ties (the others set `tie_word_embeddings=False` or default to it), so it gets the declaration and the hook goes. The remaining composites missing one are a separate change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
vasqu
reviewed
Sep 24, 2026
vasqu
left a comment
vasqu
left a comment
Collaborator
There was a problem hiding this comment.
Another round, I tried to go into a bit more of the tests overall but also some smaller comments here and there
| # declaration, so this updates it rather than deriving it. | ||
| experts.post_expert_norm = module.post_expert_norm | ||
| experts.has_post_expert_norm = True | ||
| experts.post_expert_norm_name = getattr(module, "post_expert_norm_name", None) |
Collaborator
There was a problem hiding this comment.
Ok general comment as it's quite a lot and I just went for patterns but 2-3 things
- There are a lot things that actually check for the call structure of the kernel which imo could be a separate file
- We have a lot of classes and minimal tests where we should consider which classes could be fused and what tests actually make sense / are not overkill. E.g. the serialization of a config might be not needed
Collaborator
There was a problem hiding this comment.
Imo the idea is kind of nice no? We could build for each quant and simulate each forward - not sure if we could combine with the core x kernel forwards a bit
Also we can run parametrized quite a bit for certain types where we expect it fails; imo the exact message might not be the most important for example
|
|
||
|
|
||
| @require_torch_multi_accelerator | ||
| class FineGrainedLoadPathEquivalenceTest(TestCasePlus): |
Collaborator
There was a problem hiding this comment.
memory testing mixin I guess we want to clean up a bit to make sure we dont overdo it e.g. see the llama integration tests
| with torch.no_grad(): | ||
| logits = model(ids).logits.float().cpu() | ||
| del model | ||
| backend_empty_cache(torch_device) |
Collaborator
There was a problem hiding this comment.
not sure if we really need to be that aggressive and let the mixin clean up on a per test basis 🤔
Collaborator
There was a problem hiding this comment.
So these serve as pseudo integration tests on the model loading
Imo maybe we should have tiny random models that resemble the original models and we run the tests through them -> then we can also check against logits which is much simpler than this imo
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: deepseek_v4, glm4v_moe, glm_moe_dsa, gpt_oss, ministral3, mistral4, openai_privacy_filter, qwen4_exp, finegrained_fp8, mxfp4, nvfp4 |
Contributor
CI recapDashboard: View test results in Grafana |
This branch has not been deployed
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
4 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.
What does this PR do?
bench from kernels pr
Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice.
As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues.
We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.
Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we
might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not
want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise
for reviewers and other users and will almost certainly get you blocked.
For more information, please read
CONTRIBUTING.md.Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.