DeepGEMM BF16 + mixed FP8/FP4 + MegaMoE + refactor#45634
Merged
Conversation
|
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. |
ArthurZucker
approved these changes
Jun 1, 2026
ArthurZucker
left a comment
Collaborator
There was a problem hiding this comment.
Nice!
Main concern is on the changes to src/transformers/quantizers/quantizer_finegrained_fp8.py
Comment on lines
+765
to
+771
| elif isinstance(module, nn.Linear) and "GroupedLinear" in type(module).__name__: | ||
| # Block-diagonal grouped linear (e.g. DSv4's `DeepseekV4GroupedLinear`): | ||
| # one underlying weight conceptually split into `n_groups` independent | ||
| # sub-matmuls fed by disjoint input slices. Vanilla `FP8Linear` would | ||
| # collapse those groups into one giant linear and yield the wrong | ||
| # output dim, so swap to `FP8GroupedLinear` which keeps the per-group | ||
| # bmm contract and runs each block as its own FP8 matmul. |
Collaborator
There was a problem hiding this comment.
Let's define MODULE_TO_FP8_MODULE ? this way its easy to know which are supported etcc? WDYT?
Member
Author
There was a problem hiding this comment.
kinda tricky, nn.Linear -> FP8Linear is simple, but what maps to FP8GroupedLinear and FP8Experts is not a specific class but rather a pattern of classes
|
|
||
| module, tensor_name = get_module_from_name(model, param_name) | ||
| if isinstance(module, (FP8Linear, FP8Experts)): | ||
| if isinstance(module, (FP8Linear, FP8Experts, FP8GroupedLinear)): |
Collaborator
There was a problem hiding this comment.
here we can reuse the support values MODULE_TO_FP8MODULE
…rkaround PR #46265 fixes the FP8 weight/scale loading at the right layer: scales are allocated with at least 2 in the fused dim (so the receiver matches the post-concat shape of split gate+up sources), and the model converter's substring match + suffix-preserving rename now routes both `*.weight` and `*.weight_scale_inv` keys through the same converter onto separate target buckets (`*_proj` and `*_proj_scale_inv`) with the same merge ops. That makes our dequantize=False sibling-converter logic redundant — the model's `mlp.experts.*.w1.weight → gate_up_proj` converter automatically handles `*.w1.weight_scale_inv → gate_up_proj_scale_inv` too. Dropped. Folded `max(_cdiv(...), 2)` into `_alloc_expert_proj` so the dim>=2 fix applies wherever we allocate scales (PR inlines it only for gate_up_proj; this is a no-op for typical shapes and a defensive default for small intermediate dims). Kept on our side: - `_process_model_after_weight_loading` running `setup_megamoe_weights` on each FP8Experts when `_experts_implementation == 'deepgemm_megamoe'`. - `update_tp_plan` swapping `moe_tp_experts`/`ep_router` to the megamoe variants across both `base_model_tp_plan` and `base_model_ep_plan`. Co-Authored-By: Claude Opus 4.7 <[email protected]>
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: deepseek_v4, finegrained_fp8 |
Merged
khushali9
pushed a commit
to khushali9/transformers
that referenced
this pull request
Jun 8, 2026
* init * style * full support * support EP better using offsets ! * comments * get rid of neutralize_ep_sentinels * remove deepgemm stuff * fix * prefix * move * fix * remove comment * fix unintilized outputs leaking * revert unnecessary changes * more unnecessary changes * revert downcast * keep it simple * guard deepgemm cuda version * fix style * update * add deepgemm testing * moe sentinel support * fix * compilable sonicmoe * mega moe kernel support attempt * use package for now * skip ep router and experts pre/post processing * simpler * fix Co-authored-by: Copilot <[email protected]> * fix * fix * dtensor support * more dtensor * simpler * remove comment * revert * bc order * revert extra indent * revert unnecessary change * update * less defensive * allow all kernels * alow all kernels * hub only * fix * fix * test * test * sync * check nvcc * probe * fix * test psum * test * test * probe * fix * test * nan issue * repro * repro * fix * simplifications * fix * fix * fix * fix * fix * fix * fix * empty * simplify * test deepseek * dsv4 only * download dsv4 * fix test * push * test * fix * fix ep plan * fix attempt * debug * attempt * debug * fixes in modeling * more modeling changes * more modeling changes * initial megamoe works but wrong output, deepgemm fully functional * mega moe works ! * simplification * comments * fix merge * fix rank divergence * use _keep_in_fp32_modules * simplify quantizer * introduce indexerscorer * remove testing file * style * weight rename for scorer * style * address review comments * style * align with laguna * stale comments and styling * use grouped finegrained fp8 for grouped linear Co-authored-by: sywangyi [email protected] * comments and docs * Revert "[DeepSeek V4] Fix MoE converter substring-matching FP8 scale companions (huggingface#45930)" This reverts commit ecc7e32. * fix quantizer and use upstream deepgemm * remove force al kernels * [deepseek_v4] fix FP8 save-side substring match + fp32 gate through softmax * core_model_loading: WeightConverter.reverse_transform now appends a (?=\.|$) token-boundary lookahead to each reversed source so they can't substring-match longer sibling tokens. Re-fixes huggingface#45794 (the FP8Experts scale companion `mlp.experts.gate_up_proj_scale_inv` was being routed into the weight-only [Chunk, SplitModulelist] ops on save) without needing `$` on every model-specific converter — that was huggingface#45930's approach and broke load via `process_source_pattern`. * Add `ConversionOps.supports_round_trip` opt-in flag; the on-the-fly `quantization_operation` check in reverse_transform now uses it instead of a blanket bail-out. Flagged the ops with audited reverse pairs (Chunk/Concatenate, MergeModulelist/SplitModulelist, Transpose, Conv3d↔Linear, Ernie pair, _IdentityOp, Fp8Quantize/Fp8Dequantize). * deepseek_v4 compressors (HCA / Indexer / CSA): drop the trailing `.to(chunk_gate.dtype)` so the gate stays fp32 from the `+ position_bias` through the softmax — matches the reference Compressor.forward which upcasts x to fp32 at entry. * deepseek_v4 config: only set `attention_factor=1.0` when the compress group resolves to `rope_type="yarn"`. Suppresses the noisy "Unrecognized keys in rope_parameters for rope_type=default" warning on the second post_init pass (where yarn params have already been consumed). * Refresh slow integration test expectations to match the fp32-gate output. Verified on the published DeepSeek-V4-Flash checkpoint: RUN_SLOW=1 pytest tests/models/deepseek_v4 -k Integration -> 2 passed. * always floats scal inv with triton * minimal changes i quantizer * force stable abi build for now * support transposed * fix yarn * unnecessary change * add flash base test and fix the incorrect and crash issue during enab… (huggingface#46076) * fix incorrect output in tests/models/deepseek_v4/test_modeling_deepseek_v4.py::DeepseekV4IntegrationTest::test_v4_flash_fp8_generation and add test for flash-base Signed-off-by: Wang, Yi <[email protected]> * update Signed-off-by: Wang, Yi <[email protected]> --------- Signed-off-by: Wang, Yi <[email protected]> * merge * better errors across deepgemm * more better errors for deepgemm * ep tests * Apply suggestions from code review Co-authored-by: Arthur <[email protected]> * scope supports_round_trip to quantization ops only Drop the flag from the regular shape ops (Chunk, Concatenate, MergeModulelist, SplitModulelist, Transpose, Conv3d↔Linear, _IdentityOp, the Ernie pair) — those reverse themselves via their own `reverse_op` property, no flag needed. The flag only gates `WeightConverter.reverse_transform`'s check on the converter's `quantization_operation`, so today it's just Fp8Quantize / Fp8Dequantize. Other quant backends (bnb, torchao, eetq, mxfp4, quanto, ...) stay opt-out until their round-trip is audited. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]> * Apply suggestions from code review Co-authored-by: Arthur <[email protected]> * Update src/transformers/core_model_loading.py * style: ruff format Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]> * custom ep/tp plan for megamoe * process megamoe after weight loading * fix megamoe process * up * the actual fix??? * up * nit comment for futur * style: drop unused exception variable in Fp8Dequantize._dequantize_one * up * style: apply make fix-repo docstring formatting * simplify after merge * review comments * revert unncessary * avoid name contraction * fixes * dynamic tp/ep plan swap * fix deepgemm single error and device assertion docs * properly support bias in grouped linear * simplify * fix dist error handling and min cuda compatbility * fix least cuda compat * get alignment from deepgemm and cache it --------- Signed-off-by: Wang, Yi <[email protected]> Co-authored-by: Copilot <[email protected]> Co-authored-by: Arthur Zucker <[email protected]> Co-authored-by: Wang, Yi <[email protected]> Co-authored-by: Arthur <[email protected]> Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
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
6 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?
Based on #45621
GroupedLinear co-authored-by @sywangyi
Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result,
we ask that new users do not submit pure code agent PRs at this time.
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.
PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously.
This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result,
this policy is likely to be updated regularly in the near future. For more information, please read
CONTRIBUTING.md.Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
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.