Clarify quantized dtype ownership; keep the whole-model skip (#1743) - #1754
Clarify quantized dtype ownership; keep the whole-model skip (#1743)#1754Canonik wants to merge 3 commits into
Conversation
transformers picks weight_scale_inv's storage dtype from the checkpoint's scale_fmt: one-byte ue8m0 for "ue8m0", float32 for "float", which is the default. Same parameter, same quantizer, same role, two widths. activation_scale is float32 under either format. So cast_floating_params_to_dtype's itemsize < 2 branch protects one spelling of a quantizer-owned scale and rewrites the other. Pins that against the real FP8Linear rather than a hand-rolled fake, since the point is what transformers actually does.
The skip is right, the reason given for it was wrong, and the wrong reason is what invited TransformerLensOrg#1743. cast_floating_params_to_dtype claimed its itemsize < 2 branch skips "quantizer-owned scale parameters". It does not. It skips one-byte floats, which is a subset. transformers' finegrained-FP8 stores weight_scale_inv as float32 whenever scale_fmt is "float" (the default) and keeps activation_scale float32 under every format, and fbgemm-fp8 stores weight_scale and the expert scales as float32. Reading that docstring, dropping the call-site guard and leaning on the dtype check looks safe. It reintroduces TransformerLensOrg#1713 on those checkpoints instead. So the guard stays, and it stays whole-model, which is also the line transformers draws: to(dtype=) raises for bitsandbytes and GPTQ, half()/float() raise for anything quantized, all keyed on one is_quantized flag. from_pretrained is the component responsible for applying the requested dtype to ordinary floating params, and it has already run by the time this helper is reached, so anything still off that dtype afterward may be quantizer-owned. It also releases at the right moment, because HF deletes quantization_config when a load is dequantized. No runtime behavior change. Restores the test to asserting the skip, now with the evidence for why, and adds a real FP8Linear at both scale widths so the float32 side is covered too.
Two things worth pinning about the whole-model guard. It releases: HF deletes quantization_config when a checkpoint is loaded dequantized, so quantization_method drops back to None and normalization runs again. Without that, skipping on the config would strand dequantized checkpoints in their load dtype, which is the one way the guard could be genuinely too broad. And it releases because transformers says so, not because we assume it, so the second test fails if that ever changes upstream instead of letting those checkpoints go quiet.
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for thinking this through thoroughly @Canonik, rather than taking the issue at face value. I agree with the ownership argument and the anti-tautology check reproduced exactly. Comments on the rationale in maybe_cast_floating_params + the fixture in the scale-preservation test + the upstream-pinning test.
| dtypes (e.g., FP8 scales) that must not be overwritten. This helper wraps | ||
| the cast with that check. | ||
| The skip is whole-model on purpose, and it is a division of responsibility rather | ||
| than a loss. ``from_pretrained`` is the component that applies the requested dtype |
There was a problem hiding this comment.
Three quantizers deliberately override the requested dtype in update_dtype, which from_pretrained calls before any parameter is materialised: AWQ turns bfloat16 into float16 on CUDA/XPU, and fbgemm-FP8 and FP-Quant turn anything that is not bfloat16 into bfloat16. On those checkpoints the layernorms and embeddings end up at the quantizer's dtype, not the caller's, and they are not quantizer-owned. A reader who tests this docstring's claim finds a counterexample and reopens #1743. The bitsandbytes measurement in the parenthetical can't surface it, because bitsandbytes is one of the quantizers with no update_dtype override.
Please update to mention that the skip defers to HF's effective load dtype, which the quantizer may have overridden on purpose, rather than asserting the requested dtype was applied. Mirror the same wording in the comment at transformer_lens/model_bridge/sources/transformers.py:851-855.
| than a loss. ``from_pretrained`` is the component that applies the requested dtype | ||
| to ordinary floating parameters, and it has already run by this point; whatever is | ||
| still not at the requested dtype afterward may be quantizer-owned storage, and | ||
| dtype alone cannot distinguish ownership safely. (bitsandbytes 4-bit and 8-bit |
There was a problem hiding this comment.
"bitsandbytes 4-bit and 8-bit checkpoints were checked… Not every quantization backend was reachable to test" describes only one test environment at one moment. It duplicates the PR description and can go become incorrect over time. At 23 lines this is now the longest docstring in the file by a wide margin. Can we trim this down? Cut the measurement parenthetical (lines 288-292) and the .half() / .float() paragraph (lines 294-301). Keep the ownership rule and the two issue links.
| before = {name: param.dtype for name, param in model.named_parameters()} | ||
| maybe_cast_floating_params(model, torch.bfloat16) | ||
|
|
||
| assert {name: param.dtype for name, param in model.named_parameters()} == before |
There was a problem hiding this comment.
FP8Linear(..., has_bias=True) gives a float32 bias, and the whole-dict assertion pins it as uncast. Line 157-159 calls a bias an ordinary parameter the quantizer does not own, and the docstring says the ordinary parameter is set up at the requested dtype, but only ln_weight is. A future narrowing that preserves every quantizer-owned scale still turns this test red, on quantized.bias alone, so it fails for a reason its name disclaims. test_skips_quantized_model already covers ordinary parameters.
Please put the bias at the target dtype in the fixture near line 339 (model.quantized.bias = nn.Parameter(model.quantized.bias.to(torch.bfloat16))). The test will still fail on activation_scale and weight_scale_inv if the whole-model guard is removed.
| model.is_quantized = True | ||
| # Called unbound: remove_quantization_config only touches `model`, and building | ||
| # a real quantizer would need a live quantization config per method. | ||
| base.HfQuantizer.remove_quantization_config(None, model) |
There was a problem hiding this comment.
The docstring says the test fires if transformers stops clearing quantization_config on a dequantized load, but calling remove_quantization_config directly only pins that method's body. Delete the call from HfQuantizer.postprocess_model and the test still passes.
Lets drive HfQuantizer.postprocess_model instead, with a stub carrying pre_quantized=True and a quantization_config with dequantize=True, then assert quantization_config is gone.
|
|
||
| def _fp8_linear(scale_fmt: str) -> nn.Module: | ||
| """A real transformers finegrained-FP8 ``Linear``, or skip if the integration moved.""" | ||
| integration = pytest.importorskip( |
There was a problem hiding this comment.
pyproject.toml declares torch>=2.6, but scale_fmt="ue8m0" reaches _get_ue8m0_dtype(), which raises RuntimeError when torch.float8_e8m0fnu is missing, which will happen on torch 2.6 (this was changed in 2.7). Both ue8m0 parameterisations error rather than skip, despite the helper's docstring promising a skip. Please add if not hasattr(torch, "float8_e8m0fnu"): pytest.skip(...) before constructing the ue8m0 variant.
2 participants
Description
Fixes #1743, by correcting its premise rather than its proposed fix.
#1743 is right about the mechanism:
maybe_cast_floating_paramsskips dtype normalization for an entire model the moment aquantization_configis present. Its proposed fix is to drop that call-site guard and rely on theitemsize < 2check insidecast_floating_params_to_dtype. That would reintroduce #1713.itemsize < 2is not an ownership test. transformers' finegrained-FP8 picksweight_scale_inv's storage dtype from the checkpoint:Same parameter, same quantizer, same role, two widths, and float32 is the default.
activation_scaleis float32 under either format, and fbgemm-fp8 storesweight_scaleand its expert scales as float32nn.Parameters. Running the mergedcast_floating_params_to_dtypeagainst a realFP8Linearrewrites the float32weight_scale_invto bf16, which is #1713's corruption in a scale format the guard cannot see. It also corruptsactivation_scaleonue8m0checkpoints, so even #1713's own format is not fully protected by dtype alone.scale_fmt"ue8m0"activation_scale(float32)"float"(default)weight_scale_inv,activation_scale(both float32)Nothing is lost by skipping.
dtypealready goes intofrom_pretrained(model_kwargs["torch_dtype"]), which is the component responsible for applying it to ordinary floating parameters. Measured on bitsandbytes 4-bit (pre-quantized) and 8-bit (on-the-fly): every ordinary floating parameter already at the requested dtype, 0 mismatches, including with a deliberately conflictingconfig.dtype=float32. On the 4-bit checkpoint that is 100 distinct floating tensors against 48 packedParams4bit(a naive module walk reports 101, because GPT-2 tieswte.weightandlm_head.weight).This also matches the line transformers draws itself:
PreTrainedModel.to(dtype=...)raises for bitsandbytes and GPTQ, and.half()/.float()raise for any quantized model ("the model has already been casted to the correct dtype"), all gated on the same whole-modelis_quantizedflag. And the gate releases in step with HF:HfQuantizer.postprocess_modelcallsremove_quantization_configon a dequantized load, deletingconfig.quantization_config, soquantization_methodreturns None and normalization resumes. That is the one way the guard could have been genuinely too broad, and it isn't.No runtime behavior change. The contribution is corrected rationale plus regression tests for the actual ownership invariant. The previous docstring claimed the
itemsize < 2branch skips "quantizer-owned scale parameters", which is false and is what made removing the call-site guard look safe.Type of change
Tests
test_itemsize_guard_does_not_identify_quantizer_owned_scales[ue8m0|float]uses a realFP8Linear. The expected set is derived from the cast's own predicate (floating,itemsize >= 2, not already at target) rather than dtype width, so it stays correct if transformers movesweightto wide packed-integer storage like GPTQ's int32.test_preserves_quantizer_owned_scales_of_any_width[ue8m0|float]: both scale widths survive an active quantizer. The float32 side fails outright if the cast is ever re-enabled behind only the one-byte guard.test_skips_quantized_model: restored, with the evidence for why.test_casts_once_hf_releases_quantizer_ownership: casting resumes once ownership is released.test_hf_still_clears_quantization_config_when_dequantizing: pins the upstream behavior this design depends on, so a transformers change fails loudly instead of silently stranding dequantized checkpoints.Anti-tautology check: applying #1743's proposed fix fails 4 of these tests.
Checklist
make unit-test: 5581 passed. The 8 failures are pre-existing ondevand reproduce unchanged atb652cfa8in a clean worktree. 6 are CPU/CUDA device-mixing failures that pass underCUDA_VISIBLE_DEVICES=""(CI runs CPU-only pertests/QUARANTINES.md), and 2 aretest_attention.pybitsandbytes stub tests that CI skips viaskipif(not is_bitsandbytes_available()).uv run mypy .clean,black/isort/pyclnclean.