Repository navigation
[OMNIML-6062] [11/21] Dequantize block-FP8 checkpoints as they load expert-parallel - #2743
shengliangxu wants to merge 14 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 69.57% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 23 functions across 2 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns No listed security anti-pattern was introduced. The PR changes only modelopt/torch/utils/plugins/model_load_utils.pyand a unit test; no modelopt/example dependency manifests changed. The added code…Title check The title names the main change: block-FP8 checkpoints are dequantized during expert-parallel loading. The raw summary confirms that load_and_prepare_ep now requests FineGrainedFP8Config(dequantize=Tr… Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The PR changes only
modelopt/torch/utils/plugins/model_load_utils.pyand a unit test; no modelopt/example dependency manifests changed. The added code contains notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# noseccomment. Existingtrust_remote_codehandling remains caller-configurable with a default ofFalse.
Full details: Title check
-
Explanation
The title names the main change: block-FP8 checkpoints are dequantized during expert-parallel loading. The raw summary confirms that load_and_prepare_ep now requests FineGrainedFP8Config(dequantize=True) for FP8 checkpoints and applies a Transformers compatibility patch. The title is specific and matches the changeset. The bracketed ticket and series tags ([OMNIML-6062] [11/21]) are conventional and do not reduce clarity.
-
✨ Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-6b-ep-dense-sharding #2743 +/- ##
========================================================================
Coverage ? 61.70%
========================================================================
Files ? 649
Lines ? 72887
Branches ? 0
========================================================================
Hits ? 44972
Misses ? 27915
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
f354b3b to
ed1ebbf
Compare
|
|
/ok to test ed1ebbf |
ed1ebbf to
6db2eb0
Compare
|
/ok to test 6db2eb0 |
shengliangxu
changed the title
6db2eb0 to
6bcb1ca
Compare
6bcb1ca to
14f1610
Compare
|
/ok to test 14f1610 |
shengliangxu
changed the title
14f1610 to
f647855
Compare
shengliangxu
changed the title
|
/ok to test f647855 |
f647855 to
a40c380
Compare
shengliangxu
changed the title
a40c380 to
d407931
Compare
|
/ok to test 6805a2b |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). | ||
| # mxfp4 is for experts stored in the standard layout: GPT-OSS's are transposed, which the load refuses. | ||
| _EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) | ||
| _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] "fp8" here admits more checkpoints than the new path can dequantize.
_ckpt_quant_method returns the raw quantization_config.quant_method string, and "fp8" is not exclusive to block FP8. The llm-compressor / vLLM per-tensor FP8 layout — {"quant_method": "fp8", "activation_scheme": "static", "ignored_layers": [...]}, with scalar weight_scale per linear and no weight_block_size — carries exactly the same quant_method. Those checkpoints are common on the Hub.
Before this change such a checkpoint got the actionable refusal two hunks below:
The expert-parallel load cannot dequantize a 'fp8' checkpoint: it loads unquantized checkpoints and dequantizes mxfp4. Start from an unquantized copy of the model.
Now it is accepted, gets FineGrainedFP8Config(dequantize=True) with the default (128, 128) block, and fails somewhere inside transformers' block-FP8 loader — in the best case on scale_rows, scale_cols = scales.shape[-2:] in the new override, which raises a bare ValueError: not enough values to unpack for a 0-d/1-element weight_scale; more likely on missing/unexpected weight_scale_inv keys. Either way the user loses a message that told them what to do.
The distinguishing bit is in the config the loader already has. Suggest narrowing the admission test to block FP8 — e.g. treat quant_method == "fp8" without a weight_block_size in the checkpoint's quantization_config as an unhandled format, so it keeps falling into the existing quant_method not in _EP_SOURCE_FORMATS refusal ("...dequantizes fp8 (blockwise), mxfp4"). That also keeps _decide_distribution's new fp8 clause at line 1297 from refusing dense="fsdp" with a block-FP8 rationale for a per-tensor checkpoint.
There was a problem hiding this comment.
Fixed in 3c17f8a: _ckpt_quant_method now reports an fp8 checkpoint without a weight_block_size as fp8-per-tensor, so the expert-parallel load refuses it, before anything loads, with the format message ("cannot dequantize a 'fp8-per-tensor' checkpoint ... Start from an unquantized copy of the model"), and the decision's block-FP8 rule no longer applies to it. It was worse than a crash: on transformers 5.18 a tiny Llama saved in AutoFP8's per-tensor layout ({"quant_method": "fp8", "activation_scheme": "static"}, a scalar weight_scale and input_scale per projection) loads through FineGrainedFP8Config(dequantize=True) without an error: the dequant finds no weight_scale_inv, passes the FP8 weights through unscaled (max |w| 448 where the original has 0.079), and the scales only show up as unexpected keys in the load report. Per-tensor rows in the format-refusal and _ckpt_quant_method tables fail without the fix; the block-FP8 configs in the tests now carry their weight_block_size.
| if quant_method == "fp8": | ||
| from transformers import FineGrainedFP8Config | ||
|
|
||
| from_pretrained_kwargs["quantization_config"] = FineGrainedFP8Config(dequantize=True) |
There was a problem hiding this comment.
[SUGGESTION] This constructs a fresh FineGrainedFP8Config, so the checkpoint's own weight_block_size only reaches the dequant if transformers' AutoHfQuantizer.merge_quantization_configs happens to special-case FineGrainedFP8Config (keeping the checkpoint's config and copying only the loading attributes, as it does for the dequantize flag on Mxfp4Config). If a release instead takes the passed config wholesale, self.hf_quantizer.quantization_config.weight_block_size that both new branches read (lines 1676, 1721) is the default (128, 128), not the checkpoint's.
DeepSeek-V3 is 128x128, so neither the unit tests (which build the op directly from a SimpleNamespace) nor the GPU run can tell the two apart. For a checkpoint with a different block the consequence is at least an unexplained ValueError: a scale grid of N fits neither the configured block 128 nor an even split of D on any weight with a partial last block — the _blk division fallback silently rescues only the block-aligned ones.
Reading it off the config you already have makes it independent of the merge semantics and one line:
if quant_method == "fp8":
from transformers import FineGrainedFP8Config
_qc = getattr(_ep_config(hf_config), "quantization_config", None) or getattr(
hf_config, "quantization_config", None
)
_block = (
_qc.get("weight_block_size") if isinstance(_qc, dict) else getattr(_qc, "weight_block_size", None)
)
fp8_config = FineGrainedFP8Config(dequantize=True, **({"weight_block_size": _block} if _block else {}))
from_pretrained_kwargs["quantization_config"] = fp8_config(Also worth asserting in the docstring/comment that this is where the block comes from — the _patch_finegrained_fp8_dequant docstring says "block size from the real weight_block_size", which is only true given that plumbing.)
There was a problem hiding this comment.
Pushing back: the checkpoint's weight_block_size already reaches the dequant on every release this path runs on. get_hf_quantizer sets config.quantization_config = AutoHfQuantizer.merge_quantization_configs(<the checkpoint's quantization_config>, <ours>), which keeps the checkpoint's config and copies only our loading attributes (dequantize, and from 5.8 modules_to_not_convert), since FineGrainedFP8Config is in LOADING_ATTRIBUTES_CONFIG_TYPES; the quantizer whose config both branches read is built from that merged config. The merge is the same in every release from 5.6 to 5.19. On 5.18, merging {"quant_method": "fp8", "weight_block_size": [64, 64]} with FineGrainedFP8Config(dequantize=True) gives weight_block_size [64, 64] and dequantize True, and a tiny Llama saved with 64x64 blocks loads with the dequant reading (64, 64) and matches the original to FP8 rounding (max error 0.0027). A block set on our config would be dropped by that merge, so plumbing it would be dead code. A comment at the call now says where the block comes from (3c17f8a).
| " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or " | ||
| "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." | ||
| ) | ||
| elif ep_size > 1 and _ckpt_quant_method(hf_config) == "fp8": |
There was a problem hiding this comment.
[SUGGESTION] Two things about this clause.
1. The ep_size > 1 gate leaves dense="auto" able to pick a plan no loader honors. The reason given — a sharded dense load cannot dequantize — has nothing to do with expert parallelism: _load_sharded refuses any checkpoint with a quantization_config (line 477, "The FSDP2 load cannot dequantize this ... checkpoint; load it unsharded"). With ep_size == 1 (e.g. MODELOPT_EP_DEGREE=1, or an expert count that doesn't divide the world) unshardable stays None, so too_big makes dense = "fsdp" — and test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated now asserts exactly that for the block-FP8 config (resolve_distribution_plan(..., ep_size=1)[0] == DistributionPlan(4, dense="fsdp")). The user then gets told to "load it unsharded" for a model the decision just judged too big to replicate: an impossible instruction, where replicate was at least attemptable. Dropping ep_size > 1 and (and widening the condition to "the checkpoint is quantized at all") makes the one decision point consistent with what both loaders accept.
2. "F8_E8M0 scales" is a narrower rationale than the real one, and it is repeated four times (this message, the comment above it, the load_and_prepare_ep docstring, the NotImplementedError at line 2166). DCP doesn't fail on the scale dtype — the FSDP2/DCP route reads tensors as stored and has no dequantization step at all, which is why line 477 rejects every quantized checkpoint regardless of scale format. It's also confusing next to the _patch_finegrained_fp8_dequant docstring, which says the override deliberately does not decode E8M0: a reader comes away thinking E8M0 is the blocker in both places when it is the blocker in neither. Something like "a sharded load reads tensors as stored and cannot dequantize them" would be true for every version and every scale layout, and still names the remedy.
There was a problem hiding this comment.
Fixed in 3c17f8a, point 2: the comment, the decision's reason, the loader's docstring and its refusal now say that Distributed Checkpoint, which loads a sharded dense part, reads a block-FP8 checkpoint as stored, with no dequantization; none names a scale format any more (DeepSeek-V3's scales are fp32). Postponed, point 1: at EP 1 the plan goes to load_pretrained's FSDP2 load ([5/21], [6/21]), whose own refusal covers every quantized format, block FP8 or not. Having [5/21]'s decision replicate any quantized checkpoint there changes the non-expert-parallel plan for every format, so it belongs with that route as a follow-up. Dropping the ep_size > 1 gate for 'quantized at all' would also refuse, at EP > 1, checkpoints whose dense part is stored unquantized, such as MXFP4, which quantizes only the experts.
There was a problem hiding this comment.
Claude review — block-FP8 dequant on the expert-parallel load
Scope: full review (the trigger comment carried no scoping instructions). All 5 changed files reviewed against the new base shengliangx/hfep-6b-ep-dense-sharding: modelopt/torch/utils/plugins/model_load_utils.py (+244/-5) plus the four test files. I also read _decide_distribution whole, the _route/_load_sharded guards at :410-:480 (where a sharded load's refusal of quantized checkpoints lives), and _ckpt_quant_method at :1922, since the dense-layout decision moved into this PR.
Findings: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 2
Resolved since the last round
- The full-size fp32 scale tile is gone.
repeat_interleavenow tiles over the rows only, giving a(rows, scale_cols)temporary — ~4 MB for DeepSeek-V3's(18432, 7168)MLP weights instead of 528 MB — and the whole-column-blockunflatten(-1, (whole, block_n)).mul_(...)is a genuine view ofq, so the scaling really is in place. The measured 6.0 B/elem in the PR table matchesq(fp32) + the bf16 result. The partial-last-column fallback and thewhole == 0case (a weight narrower than one block) both broadcast correctly. - The converter fuse is now idempotent in both directions. Because
_DequantThenMergeis aMergeModulelist, theisinstance(ops[1], _MergeModulelist)test matches a[Fp8Dequantize, _DequantThenMerge, Concatenate]list and collapses it back to one op, dropping the duplicate prepend rather than skipping the fuse and leavingoperations[0]anFp8Dequantize. - The master print reports
dtypeinstead of a hardcodedbf16.
Most impactful finding
[IMPORTANT Compatibility] _EP_SOURCE_FORMATS admits more than this path dequantizes (line 1919). quant_method: "fp8" is not exclusive to block FP8 — the llm-compressor / vLLM per-tensor layout (activation_scheme: static, scalar weight_scale, no weight_block_size) reports the same string, and those checkpoints are common on the Hub. They used to get the actionable "Start from an unquantized copy of the model" refusal; now they are handed FineGrainedFP8Config(dequantize=True) with a default 128x128 block and die inside transformers' block-FP8 loader — at best on scale_rows, scale_cols = scales.shape[-2:] with a bare "not enough values to unpack". The config the loader already holds distinguishes the two: gating admission on weight_block_size being present keeps the clear refusal for the rest.
Suggestions (non-blocking)
- The checkpoint's
weight_block_sizemay not reach the dequant (line 2251). A freshFineGrainedFP8Config(dequantize=True)only preserves it if transformers'merge_quantization_configsspecial-cases this config class the way it doesMxfp4Config; otherwise both branches read the default(128, 128). DeepSeek-V3 is 128x128, so neither the unit tests nor the GPU run can tell. Plumbing the block through explicitly makes the docstring's "block size from the realweight_block_size" true by construction. - The new
unshardableclause is gated onep_size > 1, and its stated reason is narrower than the real one (line 1297). A sharded dense load refuses every quantized checkpoint (:477), independent of expert parallelism — so atep_size == 1dense="auto"still picksfsdpfor a block-FP8 checkpoint (the updated test asserts it) and the user is told to "load it unsharded" for a model just judged too big to replicate. Separately, "F8_E8M0 scales" now appears as the rationale in four places, including two user-facing messages, while the actual blocker is that a sharded load reads tensors as stored and has no dequantization step at all.
What reads well
- The arithmetic in both version branches holds up against a careful re-derivation:
s[..., :whole, None]really does slice the scale columns (theEllipsisabsorbs the leading dims androws), the row tile is always long enough under both_blkbranches, andscale_cols - wholecan only ever be 0 or 1, so the partial-block fallback cannot silently drop columns. The in-place path also leaves the caller's tensor alone (to(float32, copy=True)). _blktrusting the configured block over division, with576 % 5and9 % 3both named as traps, plus thegrid-divides-the-rowscase that would pass under the naive rule.- The 5.10.1+ wrapper's bypass predicate gates on block alignment, not grid divisibility, which is the fix that closed the last round's CRITICAL;
partial-dividing-the-grid((192, 256) on a (2, 2) grid) pins it. Double-wrapping on a re-imported module is harmless there too, since the inner wrapper sees an already-padded fp32 tensor and falls straight through to stock. test_undoing_the_patch_restores_stock_transformersdiffingvars(cls)to assert the patch rebinds only its three targets, which is what makestest_an_expert_converter_it_cannot_fuse_stops_the_loadmeaningful. Both refusal paths fail loud on exactly the states 5.9'stensor_idxbug produced, and the patch can only raise identically on every rank, so no collective divergence.- The new
load_and_prepare_epguard sits before_resolve_checkpoint_dir, so a refused plan downloads nothing.
Risk
Low. The dequant arithmetic and the converter fusion both look right and are covered against a blockwise reference including the shapes that motivated the patch; the remaining finding is about which checkpoints the gate lets in, not about what it does with them, and it fails loud rather than quietly. All three items are local fixes.
🤖 Generated with Claude Code
|
/ok to test 9902c0d |
|
/claude review |
| " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or " | ||
| "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." | ||
| ) | ||
| elif ep_size > 1 and _ckpt_quant_method(hf_config) == "fp8": |
There was a problem hiding this comment.
[SUGGESTION] The two unshardable reasons are mutually exclusive by elif, but both can hold at once — and the branch that wins then hands out advice the other one refuses.
A block-fp8 MoE on 8 ranks with MODELOPT_EP_DEGREE=2 and dense="fsdp" takes the first branch (world_size // ep_size == 4 > 1) and gets its remedy: "Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or MODELOPT_FSDP_SIZE=1 …". Following the first half (MODELOPT_EP_DEGREE=8) collapses edp to 1, which is exactly what makes this new elif fire — so the next attempt is refused again, now for the block-fp8 reason, and only the second attempt surfaces the remedy (MODELOPT_FSDP_SIZE=1) that was going to be the answer all along. Two refusals and two edits to get one piece of advice.
Cheapest fix is to let the fp8 reason append rather than compete, so a checkpoint that hits both is told both at once:
unshardable, remedy = None, ""
if ep_size > 1 and world_size // ep_size > 1:
unshardable = (
f"within each of its {world_size // ep_size} expert-data-parallel replicas, a layout "
"the export cannot write yet"
)
remedy = (
" Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or "
"MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU."
)
if ep_size > 1 and _ckpt_quant_method(hf_config) == "fp8":
unshardable = (
"from a block-fp8 checkpoint: Distributed Checkpoint reads it as stored, with no "
"dequantization"
)
remedy = (
" Set MODELOPT_FSDP_SIZE=1 to replicate it, which dequantizes it as it loads, or start "
"from an unquantized copy of the model."
)(The fp8 reason overwriting the edp one is right: it is the one that survives every MODELOPT_EP_DEGREE the user can pick.)
| # mxfp4 is for experts stored in the standard layout: GPT-OSS's are transposed, which the load refuses. | ||
| _EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) | ||
| # fp8 is block FP8; _ckpt_quant_method calls per-tensor FP8 fp8-per-tensor, which the load refuses. | ||
| _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) |
There was a problem hiding this comment.
[SUGGESTION] The refusal an AutoFP8 user now gets reads as a contradiction. _EP_SOURCE_FORMATS is rendered straight into the message at :2195-:2199, so a per-tensor FP8 checkpoint produces:
The expert-parallel load cannot dequantize a
'fp8-per-tensor'checkpoint: it loads unquantized checkpoints and dequantizes fp8, mxfp4. Start from an unquantized copy of the model.
Their config.json says "quant_method": "fp8", and the message says it dequantizes fp8 — yet it refused. The distinction (weight_block_size present or not) is nowhere in the text, so the natural next step is to assume a bug and retry, not to go find a bf16 copy. The invented "fp8-per-tensor" label is doing the discriminating work but reads like a format name the user should recognize, and they won't find it in any checkpoint.
Since the set is both the gate and the message, a small display mapping keeps the gate exact and the message honest — e.g. render fp8 as block fp8 in dequantized, and say in the refusal that a fp8 checkpoint without weight_block_size is per-tensor FP8, which this load does not dequantize. One sentence, and the user knows which of their two FP8 checkpoints to use.
| # weight_scale_inv. Mirror get_model's mxfp4 dequant-on-load, using FineGrainedFP8Config(dequantize=True). | ||
| # (compressed-tensors pack-quantized experts, e.g. Kimi, are handled separately by the on-the-fly | ||
| # decompressor plugin; here we only cover the native fp8 case.) | ||
| if quant_method == "fp8": |
There was a problem hiding this comment.
[SUGGESTION] ModelOpt already has a block-FP8 path, and this comment doesn't say why it isn't the one being used here — while the repo's own PTQ skill tells readers the opposite of what this line does.
modelopt/torch/quantization/plugins/huggingface.py registers _QuantFP8Linear for FP8Linear (:1393) and _QuantFP8Experts for FP8Experts (:1390), both doing lazy per-use blockwise dequant — and _QuantFP8Experts.get_expert_weight (:1096-:1104) is the same tile-and-crop arithmetic this patch adds, partial-last-block crop included:
scales = scales.repeat_interleave(rows, 0).repeat_interleave(cols, 1)
return (weight.float() * scales[: weight.shape[0], : weight.shape[1]]).to(dtype)Meanwhile plugins/modelopt/skills/ptq/references/unsupported-models.md:247 says, in bold: "Do NOT use FineGrainedFP8Config(dequantize=True) — it expands the entire model to BF16 upfront, wasting ~2x GPU memory. The plugin approach is both more memory-efficient and simpler." An agent loading that skill and then reading this code gets contradictory instructions, and a human reviewer has to re-derive the answer.
I believe the answer is that the lazy path genuinely cannot serve this load — _QuantFP8Experts is _eager_only and keys off transformers' FP8Experts module, whereas the EP path needs materialized bf16 fused gate_up_proj/down_proj params for enable_expert_parallel's grouped-GEMM sharding and the torch_a2a _grouped_linear, which no FP8 expert module provides. If that's right, say it in the comment (one clause: "not _QuantFP8Linear/_QuantFP8Experts: the grouped-GEMM EP shard needs materialized bf16 fused expert params, which the lazy FP8 modules do not provide") and consider updating the skill reference to scope its "do not" to the non-EP path. If a shared helper for the tile-and-crop is cheap to extract, that removes the third copy of the arithmetic too — but the comment is the part that costs nothing and saves the next reader the trace.
|
|
||
| _FineGrainedFP8HfQuantizer.update_weight_conversions = _uwc_fixed | ||
| _FineGrainedFP8HfQuantizer._modelopt_uwc_patched = True | ||
| except Exception as e: |
There was a problem hiding this comment.
[SUGGESTION] except Exception turns every failure in this block into the same claim about the cause, and one of those failures isn't that cause.
The block's own guard — hasattr(_FineGrainedFP8HfQuantizer, "update_weight_conversions"), with the comment "A transformers without the hook prepends no dequant to the expert converters: nothing to fuse" — sits after the three imports it is meant to cover. On a release that has no transformers.core_model_loading at all (the EP floor here is require_transformers("5.6"), and the new test module skips on 5.5 for lacking Fp8Dequantize, so there is a window between the floor and where the conversion-op loader landed), the import raises first and the user is told:
block-FP8 load: could not fuse transformers' expert dequant into its merge (No module named 'transformers.core_model_loading').
On such a release nothing prepends a dequant and nothing would be row-sliced — the stated failure mode cannot occur there. The refusal itself is defensible (that transformers almost certainly can't do this load anyway), but it points at the wrong thing, and test_an_expert_converter_it_cannot_fuse_stops_the_load pins exactly this path by removing core_model_loading from sys.modules, so the test's name encodes the same mis-attribution.
Two options, either fine: name the version floor in the message ("this transformers has no conversion-op loader; block-fp8 expert-parallel loads need >= X"), or hoist the quantizer import and the hasattr check ahead of the core_model_loading imports so "nothing to fuse" is decided before anything that can fail for an unrelated reason. Also worth narrowing except Exception to (ImportError, AttributeError) — as written, a genuine bug introduced later in _DequantThenMerge's class body would be reported as a transformers-compatibility problem.
There was a problem hiding this comment.
Claude review — block-FP8 dequant on the expert-parallel load
Scope: full review (the trigger comment carried no scoping instructions). All 5 changed files reviewed against shengliangx/hfep-6b-ep-dense-sharding: modelopt/torch/utils/plugins/model_load_utils.py (+260/-7) plus the four test files. For composition I also read _decide_distribution whole, _load_sharded's guards at :437-:480, the dtype resolution at :2212, and — new this round — _QuantFP8Linear / _QuantFP8Experts in modelopt/torch/quantization/plugins/huggingface.py, since they are the project's other block-FP8 mechanism.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 4
Resolved since the last round
- Per-tensor FP8 no longer slips into the block-FP8 path.
_ckpt_quant_methodnow reads aquant_method: fp8without aweight_block_sizeas"fp8-per-tensor", which_EP_SOURCE_FORMATSrefuses before_resolve_checkpoint_dir, so an AutoFP8 checkpoint fails loud instead of dying inside transformers' block loader with weights scaled by 1. The discrimination is right for the formats in the wild: DeepSeek-V3 and the block-FP8 Qwen checkpoints carryweight_block_size; AutoFP8's per-tensor and per-channel configs do not, and an explicit"weight_block_size": nullfalls on the refusing side too. The 448-vs-0.079 table in the description is the right evidence for why this had to be a refusal rather than a warning. - The dense layout is now decided and enforced consistently.
_decide_distributionreplicates a block-fp8 dense part instead of leavingdense="auto"free to pickfsdp, refuses an explicitfsdpnamingMODELOPT_FSDP_SIZE=1, andload_and_prepare_eprepeats the refusal at:2188for a plan handed to it directly. The guard is correctlyfp8-only and not extended tomxfp4: gpt-oss's quantized tensors are all expert tensors, whichfrom_pretrainedloads over theepsubmesh, so only block FP8 — which covers the dense linears as well — breaks under a DCP-sharded dense load. - The
ep_size=1corner the last round raised is now pinned by a test, and it fails loud:_load_shardedrefuses anyquantization_configat:477before a weight is read.
Verification I ran on the arithmetic (no finding, recorded so it isn't re-derived)
s[..., :whole, None]really does slice the scale columns —Noneadds an axis without consuming one, so(…, rows, scale_cols)becomes(…, rows, whole, 1)against anunflattened(…, rows, whole, block_n).q[..., :whole * block_n]is basic slicing, andunflatten(-1, (whole, block_n))on strides(…, cols, 1)is expressible as a view, so themul_is genuinely in place onq's own copy.whole == 0(narrower than one block) andscale_cols - whole ∈ {0, 1}both broadcast correctly, and the row tile is long enough under both_blkbranches.- Double-wrapping on the >= 5.10.1 path is harmless in both directions, not just one: the outer wrapper's padded tensor is fp32, so the inner falls through on
quantized.dtype not in fp8_dtypes; and on a non-partial weight the outer bypasses and the inner independently re-derives the same bypass. The FP4 dtypes (int8/float4_e2m1fn_x2) are excluded fromfp8_dtypes, so upstream's unpack and ue8m0 decoding are never routed through the pad. dtype = getattr(hf_config, "dtype", None) or torch.bfloat16is always concrete at the call site, so_FP8_DEQUANT_DTYPEcan never be left at a previous load's value by aNoneargument.- The patch is called on every rank, unconditionally, and can only raise identically on all of them — the master-only
printis the sole rank-dependent statement, and it issues no collective. No divergence.
Suggestions (non-blocking, 4 inline)
- The two
unshardablereasons areelif-exclusive but can both hold; theedpbranch then wins and its remedy (MODELOPT_EP_DEGREE=<world>) leads straight into the fp8 refusal, so the user needs two failed attempts to learnMODELOPT_FSDP_SIZE=1. Letting the fp8 reason overwrite rather than compete is a three-character change. - The per-tensor refusal reads as a contradiction to the person receiving it: "cannot dequantize a
'fp8-per-tensor'checkpoint: … dequantizes fp8, mxfp4", when theirconfig.jsonsaysfp8. Renderingfp8asblock fp8in the message, and namingweight_block_sizeas the discriminator, costs one sentence. - The comment at
:2266doesn't say why_QuantFP8Linear/_QuantFP8Experts— ModelOpt's existing lazy block-FP8 dequant, whoseget_expert_weightis this same tile-and-crop including the partial-block crop — aren't the mechanism here, whileplugins/modelopt/skills/ptq/references/unsupported-models.md:247tells readers in bold not to useFineGrainedFP8Config(dequantize=True). The reason looks sound (the grouped-GEMM EP shard needs materialized bf16 fused params); it just isn't written down, so the code and the skill read as contradicting each other. - The fuse's
except Exceptionreports a missingtransformers.core_model_loadingas "could not fuse transformers' expert dequant into its merge" — a cause that cannot exist on a release without the conversion-op loader, which thehasattrguard was meant to cover but sits after the imports.
What reads well
_ckpt_quant_methoddoing the per-tensor discrimination in the one place both call sites already go through, rather than adding a second predicate, and keeping the refusal ahead of_resolve_checkpoint_dirso a refused load downloads nothing.- The
fp8-only dense guard, reasoned from which tensors each format quantizes rather than from the format name — that is why extending it tomxfp4would have been wrong, and the asymmetry is the kind of thing that usually gets "fixed" into a bug later. test_undoing_the_patch_restores_stock_transformersdiffingvars(cls)to assert the patch rebinds only its three targets, which is what makes the refusal test meaningful; and the rebind following the fuse, so a refused patch leaves transformers byte-identical.- The test table carrying its own reasoning (
grid-divides-the-rows: 9 rows of 4-row blocks give a 3-row grid that also divides 9) — the cases that would pass under the naive rule are the ones written down. - Both
< 5.10.1refusal paths fail loud on exactly the states 5.9'stensor_idxbug produced, instead of loading unscaled weights.
Risk
Low. The arithmetic in both version branches holds up against a re-derivation and is covered against a blockwise reference on the shapes that motivated the patch; the dense-layout decision and its two enforcement points agree; and the new refusal closes the one admission gap left from the last round. All four remaining items are message clarity, comment provenance, and one exception-handling narrowing — none of them change what the code computes.
🤖 Generated with Claude Code
|
/ok to test 37db5f8 |
|
/ok to test 12aee3f |
A native block-FP8 checkpoint (e.g. DeepSeek-V3) is dequantized to bf16 while it loads, through transformers' FineGrainedFP8Config, since quantizers attach to bf16 weights. Its dequant is corrected on transformers < 5.10.1 for weights whose last 128x128 block is partial, and made to emit the model's dtype. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…out the session Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The patch and its module flags sit in model_load_utils with the expert-parallel loader that applies them, so the test drives and resets them there. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The patch fuses the Fp8Dequantize that transformers prepends to the routed-expert converter into its MergeModulelist, so the loader still gives each expert its tensor_idx. Only the install was tested. The new test runs the patched update_weight_conversions on an expert converter and the resulting chain on per-expert block-FP8 weights. It fails if the converter no longer leads with a MergeModulelist, or if the fused op skips the dequant or the merge. transformers 5.18's chain passes an output_dtype the patched dequant does not take, so the chain runs that version's own dequant; the patched arithmetic has its own tests. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
- The two block-size cases, a grid finer than the configured block and no block configured, become rows of the blockwise-reference table, which gains configured and grid columns. - The dtype test keeps its unrequested (bf16) and requested (fp16) cases; requesting bf16 or fp32 took the same branch. - The version test keeps the 5.10.1 boundary; 5.17.0 took the same branch. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The pre-5.10.1 dequant patch emitted a non-empty weight unscaled when its scale grid was empty. That only happens when a weight and its scales were sliced apart, as transformers 5.9's expert tensor_idx bug did on expert-parallel loads, so it loaded wrong weights without a word; upstream's own dequant fails on such a grid. It also fell back to the configured block for a grid neither that block nor an even split gives, scaling rows wrongly. Raise in both cases; an empty weight still passes through its empty grid. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From 5.10.1 the patch kept transformers' own dequant untouched, on the assumption that no shipping model has a partial last block there. DeepSeek-V3, V3.1 and Kimi-K2-Instruct do: kv_a_proj_with_mqa has 576 rows under a 5-row grid of 128-row blocks, and Fp8Dequantize still raises "not divisible by scale grid" for it, so their block-FP8 loads failed. The stock dequant is kept, for its ue8m0 and FP4 decoding and output_dtype, and wrapped: an FP8 weight whose grid is the configured block's ceil grid is padded to whole blocks for it and cropped back. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
When fusing transformers' prepended expert dequant into its merge failed, the patch printed a warning and went on. Every routed expert then loaded sliced apart from its scales. The failure now raises. A transformers without the update_weight_conversions hook prepends no dequant, so there is nothing to fuse and the patch skips it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…tizes The load dequantizes native block-FP8 checkpoints as it loads, so they join the formats it accepts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…heck The hub-ID test now loads an unquantized, a block-FP8 and an MXFP4 config: each must reach the checkpoint's resolution, past the format check, so dropping either format from the accepted set fails it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
From transformers 5.10.1 the expert-parallel load wraps transformers' block-FP8 dequant, which derives the block size from the scale grid, and padded a weight to whole blocks only when it did not divide into that grid. A partial last block can still divide evenly: 448 rows under a 4-row grid of 128-row blocks go to the dequant as blocks of 112, and every row group after the first gets the wrong scale, silently. The wrapper now pads whenever the weight is not block-aligned, which is a no-op for whole blocks. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The pre-5.10.1 dequant multiplies its own fp32 copy of the weight by the tiled scales in place, so a weight holds two full-size fp32 temporaries instead of three while it loads. The dequant is rebound after the expert converter's fuse, so a refused patch leaves transformers as it was. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The rule that a block-FP8 checkpoint's dense part is not sharded lands here, with the load that makes such a checkpoint loadable, rather than in [10/21], where it preempted the format refusal. Left to decide, the plan replicates that dense part and says why; asked to shard it, the decision refuses, naming MODELOPT_FSDP_SIZE=1, and the loader refuses a plan that shards it. The pre-5.10.1 dequant scales its fp32 copy of a weight in place, through a view of its whole column blocks and then the partial last one, against scales tiled over the rows only. It builds no full-size fp32 scale tile: an (18432, 7168) weight peaks at 6 bytes per element instead of 8. A weight narrower than one block joins the reference table. The expert converter's fuse no longer skips a merge it already fused, a layout transformers' update_weight_conversions never produces (it builds fresh converters on every call); fusing it again would dequantize once all the same. The load's log names the dtype the dequant emits. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
quant_method "fp8" is not block FP8 alone: per-tensor FP8 checkpoints (AutoFP8's: a scalar weight_scale per weight, no weight_block_size) name the same method. The expert-parallel load admitted them, and transformers' dequant reads them as 128x128 blocks with no weight_scale_inv, so their weights load unscaled without an error (on transformers 5.18 a tiny Llama loaded that way has max |w| 448 where the original has 0.079). _ckpt_quant_method now calls an fp8 checkpoint without a weight_block_size "fp8-per-tensor", which the load refuses, before anything loads, as a format it cannot dequantize; the decision's block-FP8 rule no longer applies to it either. The decision, the loader's refusal and its docstring give the reason a block-FP8 dense part cannot be sharded as Distributed Checkpoint reading it as stored, with no dequantization: its scales' format was never the blocker (DeepSeek-V3's are fp32). A comment says where the dequant's block size comes from: transformers keeps the checkpoint's quantization_config and takes only the loading attributes from the one the load passes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 3c17f8a |
What does this PR do?
Type of change: New feature
A native block-FP8 checkpoint, such as DeepSeek-V3, is dequantized to bf16 while it loads expert-parallel, through transformers'
FineGrainedFP8Config(dequantize=True). Quantizers attach to bf16 weights.A block-FP8 checkpoint keeps its dense part replicated (moved here from [10/21], where block FP8 did not load at all): DCP, which loads a sharded dense part, cannot dequantize it: it reads the checkpoint as stored. Left to decide, the plan ([5/21]) replicates it and says so; an explicit
fsdpis refused by the decision, namingMODELOPT_FSDP_SIZE=1, and by the loader.Per-tensor FP8 is refused. AutoFP8's per-tensor checkpoints name the same
quant_methodwithout aweight_block_size; transformers' dequant would read them as 128x128 blocks and load their weights unscaled, without an error. The load refuses them, before anything loads, asfp8-per-tensor. A tiny per-tensor FP8 Llama, loaded that way on transformers 5.18:On transformers older than 5.10.1, the loader also patches that dequantization in three ways:
kv_a_proj_with_mqa: (576, 7168) with a (5, 56) grid. It also raises on an empty scale grid. The patch tiles each block's scale and crops to the weight, using the checkpoint'sweight_block_size.MergeModulelist. The loader then assigns no per-expert index, and the expert-parallel shard slices each expert's rows instead of whole experts: a DeepSeek-V3gate_up_projloads as (256, 32, 7168) and fails calibration. The patch fuses the dequantization into the merge, so the converter starts with aMergeModulelistagain and runs as upstream intends.The patch installs once and survives a re-imported module. From transformers 5.10.1 it keeps upstream's dequantization, which takes the destination dtype and decodes E8M0 scales, and only pads a weight with a partial last block to whole blocks for it, since upstream still infers the block from the scale grid: (448, 7168) under a (4, 56) grid would otherwise be scaled in blocks of 112 rows.
The pre-5.10.1 dequantization scales its own fp32 copy of the weight in place, through a view of its whole column blocks, against scales tiled over the rows only, so no full-size fp32 tile is built. It is rebound only once the expert converter's fuse is in, so a refused patch leaves transformers as it was. Against scaling by a full-size fp32 tile, on CPU with one thread:
Usage
Automatic:
load_and_prepare_epreads the checkpoint'squantization_configand picks this path forquant_method: fp8with aweight_block_size.Testing
tests/unit/torch/utils/test_fp8_dequant_patch.py, on CPU:MergeModulelist, and per-expert block-FP8 weights come out as the blockwise reference. It fails without the fusion, and if the fused op skips the dequantization or the merge. Skipped on transformers 5.5, which has noupdate_weight_conversions.test_ep_dense_sharding.py); the decision replicates it, or refuses an explicitfsdp, namingMODELOPT_FSDP_SIZE=1(test_resolve_distribution_plan.py).fp8-per-tensor(test_ep_loader.py; both rows fail without the fix); the block-FP8 configs in the tests carry theirweight_block_size.tests/unit/torch/utils,export,models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part 11 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [10/21].
Summary by CodeRabbit