Repository navigation
Finegrained kernels: unified quantization kernels api (separate quantizers) - #49015
Draft
IlyasMoutawwakil wants to merge 75 commits into
Draft
IlyasMoutawwakil wants to merge 75 commits into
IlyasMoutawwakil wants to merge 75 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>
…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
…ey found Tests - tiny random copies of the shipped checkpoints (Qwen3, Ministral-3, Mistral-Small-4, Llama NVFP4, DeepSeek-V3/V4, GPT-OSS, GLM-5.2 NVFP4, MiniMax-M3), shrunk from the real configs and layouts: clean load, logits against the dequantized weights and a golden slice, and a byte-exact save back to the checkpoint layout - test_conversions: scheme-agnostic op tests, then one dummy-model chain per scheme (load through the real quantizer, check the held tensors, save back) - kernel call-structure tests move to tests/kernels/test_finegrained.py, the sharded load paths to test_distributed.py; per-pairing forwards; activation formats a weight format can't take fail at config construction - test_reverse_loading_mapping runs for DeepSeek-V4 and MiniMax-M3 (head models) Fixes - save: a parameter the quantizer creates follows its weight's reversed converter (`quantization_param_suffixes`), replacing the swizzle-only workaround; converter outputs get the renamings after conversion; SplitModulelist squeezes only its own dim - load: a pre-quantized checkpoint's scales keep their dtype (`param_keeps_checkpoint_dtype`) - the MXFP4 weight-only default reaches the groups the modules read (GPT-OSS ran W4A4) - modelopt: `dynamic: false` is NVFP4's calibrated global, not static activations; dense NVFP4 linears; FP8 KV-cache scales ignored; save keys and storage - VLM skip lists match renamed modules (MiniMax-M3's vision tower) - DeepSeek-V4 and MiniMax-M3 mappings reverse faithfully Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The quantizer's `dtype_plan` joins the model's in `from_pretrained` (`None` keeps the checkpoint's dtype), replacing the `param_keeps_checkpoint_dtype` hook in the loader. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
ArthurZucker
left a comment
ArthurZucker
left a comment
Collaborator
There was a problem hiding this comment.
Just like we have:
- tp_plan: glob → sharding style, matched by _get_parameter_tp_plan (distributed/tensor_parallel.py:74).
- torchao FqnToConfig: fqn/regex → per-layer config (quantizers/quantizer_torchao.py:133).
- compressed-tensors config_groups + targets (utils/quantization_config.py:1128). And https://github.com/huggingface/transformers/blob/omni-ci-investigate/docs/source/en/quantization/compressed_tensors.md#L133-L133 for an example.
On top of those, the PR adds a fourth version (FineGrainedGroup.targets).
The module-swap side is the same. There are 19 replace_with_*_linear functions in integrations/ (replace_with_fp8_linear, replace_with_mxfp4_linear, replace_with_nvfp4_linear, replace_with_gguf_modules, …), each walking the model with should_convert_module.
The PR just adds another one + a config, 4 quantizer subclasses (one empty), and a new ops hierarchy. That's about 2.3k src lines and about 3.6k test lines.
Why is this "FineGrained" at all? Nothing here is specific to these formats. The API we want is QuantizationConfig(plan={pattern: scheme}), where a scheme is any method (fp8, nvfp4, mxfp4, gguf, torchao.xxx, etc, …), matched with the same globs as tp_plan. We already have this idea three times (tp_plan, torchao FqnToConfig, compressed-tensors config_groups). Let's generalize one of those rather than adding a fourth (FineGrainedGroup.targets).
- It's better if we detect this from the dtype, as this is almost a "direct" mapping.
Per-module plans need per-scheme conversions, not per-class ones.
Conversions should be static data. Each scheme ships a list ofWeightConverter/WeightRenaming, always registered.
Patterns that don't match are ignored: an op that needs config should read it from the plan entry of the key it converts.
Module swapping should be one API. apply_plan(model, plan) swaps every matching module for scheme.module_cls. That replaces the different per-integration replace_with_* and allows standardization.
RE the scope: please split this.
Comment on lines
+166
to
+168
| # a 0-dim per-expert value has no axis to slice, so this rank takes it whole | ||
| if not source_shape: | ||
| return source[...].to(device=device, dtype=dtype) |
Collaborator
There was a problem hiding this comment.
why not do this even earlier ? (whether there is tensor idx or not?
Comment on lines
+195
to
+197
| # a 0-dim scale is one value for the whole tensor: no axis to split, every rank keeps it | ||
| if meta.ndim == 0: | ||
| placement = Replicate() |
Collaborator
There was a problem hiding this comment.
Mmmm this does make sense, of course, but its kind of a departure from the "explicit" API of replacing actual stuff.
The else here is gonna be removed anyways (as it does not make sense to "handle" so many different stuff on the same "colwise".
To me, we should not do this.
The quantization config creates these parameters, and knows their shape. It should update the quantization config to add "replicate" to the scale.
It does make sense to suppose 1 dim tensors are replicated, but it's not documented in this class.
I am fine with a # TODO make the plan update this if needed
We don't have other scales format I suppose, so this is fine for now. Just documented
Comment on lines
+790
to
+792
| "moe_experts_colwise": MoEParamShard(Shard(1)), | ||
| "moe_experts_packed_colwise": MoEParamShard(_StridedShard(dim=1, split_factor=2)), | ||
| "moe_experts_rowwise": MoEParamShard(Shard(2)), |
Collaborator
There was a problem hiding this comment.
not sure we even neeed them as rowise, packwed rowise and packe colwise support that. But I'll keep reading.
Collaborator
There was a problem hiding this comment.
why do we default to TP when EP is the actual thing people use now?
Member
Author
There was a problem hiding this comment.
we don't default to TP. EP stays opt-in through enable_expert_parallel, and this config already had its base_model_ep_plan. what this change does is fill in the expert and shared-expert entries the TP plan was missing (TP previously left every expert replicated), so both paths shard the experts. having both working also gives us a good consistency check: EP logits == TP logits == device_map logits which catches sharding mistakes that neither path would show on its own, like the new model that adds post expert norm and needs a collective after the post norm in TP, if that collective is lost in the integration (fused kernel doesn't do it) the logits diverge.
Collaborator
There was a problem hiding this comment.
unrelated not to merge with this
Member
Author
There was a problem hiding this comment.
opening separate PR
Member
Author
There was a problem hiding this comment.
for context this needed fixing after i added tp support
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
A mixed checkpoint keeps some layers' experts unquantized; the gate|up interleave and the input-scale split read the model's first FineGrainedExperts and applied its layout to every layer. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This reverts commit 5ea2a3e. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The gate_up holds one activation global per expert, as calibrated, and the kernels quantize each routed row against its own expert's. TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL (`global` or `rank`) shares their max instead, so the rows quantize once per token: faster at prefill, at some block-scale precision for experts below the max. The NVFP4 quantizer registers it after loading; `_projection_operands` hands the shared global to both forwards. EP plans shard the gate_up's activation globals with its other per-expert companions. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
vasqu
reviewed
Oct 5, 2026
| ``TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL`` turns it on and picks the experts the max | ||
| runs over under expert parallelism: ``global`` all of them, so the numerics do not depend on the EP | ||
| size, or ``rank`` this rank's, a tighter global. Unset keeps every expert's own.""" | ||
| scope = os.environ.get("TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL") |
Collaborator
There was a problem hiding this comment.
maybe a quant config value instead? Not sure how I feel about introducing an env here
Member
Author
There was a problem hiding this comment.
i tried to follow the sglang convention, imo it's better if we keep the config only for core quantization stuff instead of enabling/disabling a specific format's lossy optims, kinda like the disable deepgemm env var (which vllm does as well)
6 tasks
Conflicts: mxfp4 imports (both kept), DeepSeek-V4 RMSNorm (main's fp32 multiply), MiniMax-M3-VL tests (on the VLM tester now, so no config_tester). Regenerated deepseek_v4 and glm4v_moe from modular (a duplicate lm_head `_tp_plan`), and `_get_parameter_tp_plan` is `_get_parameter_plan` on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…erands helper on the mock Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: deepseek_v4, glm4v_moe, glm_moe_dsa, gpt_oss, minimax_m3_vl, ministral3, mistral4, qwen4_exp, finegrained_fp8, mxfp4, nvfp4 |
Contributor
CI recapDashboard: View test results in Grafana |
Member
Author
|
run-slow: deepseek_v4, glm4v_moe, glm_moe_dsa, gpt_oss, minimax_m3_vl, ministral3, mistral4, qwen4_exp, finegrained_fp8, mxfp4, nvfp4 |
Contributor
CI (Nvidia + AMD)Nvidia CICommit Info
✅ No failing test specific to this PR 🎉 👏 ! AMD CICommit Info
Model CI Report❌ 1 new failed tests from this PR 😭
|
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
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.