Skip to content

Refuse quantized weights in weight-space code paths instead of reading them as matrices - #1669

Merged
jlarson4 merged 11 commits into
devfrom
bug/quantized-weight-guards
Aug 14, 2026
Merged

Refuse quantized weights in weight-space code paths instead of reading them as matrices#1669
jlarson4 merged 11 commits into
devfrom
bug/quantized-weight-guards

Conversation

@jlarson4

Copy link
Copy Markdown
Collaborator

Description

  • Quantized forward passes continue working. The wrapped HF module dequantizes internally, and the bridge's forward paths already skip non-floating-point parameters when picking a compute dtype. What did not work is arithmetic on the stored weight: reshaping it per head, slicing a fused projection, folding LayerNorm into it. There the storage is packed, split from its scales, or not a tensor at all. Those reads returned plausible-looking matrices with no error.
  • New transformer_lens/utilities/quantization.py holds one predicate: a weight is readable if it is a torch.Tensor with a floating-point dtype of itemsize >= 2. It is future-proof against new narrow types without enumerating names. The error names the quantization method, resolved from the HF config's quant_method then the weight's class.
  • FP8 is often missed case. torch.float8_e4m3fn.is_floating_point is True, so an isinstance + is_floating_point check admits it and FP8 is the only family that survives both tensor_split and nn.Parameter() without complaint. Every guard test parametrizes over it for that reason.
  • Guarded sites: the bridge weight accessors (W_Q/W_K/W_V/W_O, W_in/W_gate/W_out), which returned packed bytes silently
  • use_split_qkv_input/use_attn_in, both boot-time fused-weight split functions, TransformerBridge.process_weights, the OLMoE converter, and the BitNet adapter, whose flagship packed checkpoint the model registry already records at 0% on the forward-pass phase.
  • The refusal covers the ordinary load path, not just user-supplied models. The original gate sat inside if hf_model is not None:, so from_pretrained("<name>") reached the converters unguarded, and only 3 of 28 converters guard their reads. Measured on a tiny Llama: same-shape int8 and FP8 convert silently and survive load_state_dict, which casts them to float32, so an int8 code of 107 lands as the weight 107.0 with no error anywhere; packed 4-bit is caught late by a shape mismatch, and GPTQ/AWQ fail loudly only by accident (their QuantLinear has no .weight at all). _refuse_unsupported_quantization now runs once in get_pretrained_state_dict, covering all 28 converters. It reads the loaded model's config rather than cfg.quantization_method, because convert_hf_model_config infers llama and gemma from the model name and never fetches a config for them. And it refuses on the stored weights rather than the declaration, so a checkpoint loaded with dequantize=True keeps working.
  • The boot-time guard added to the default fused splitters had no reachable caller for its family: Phi-3, GLM and GLM-4V all override split_gate_up_matrix with Phi-3's own splitter, and phi3.py's _split_gate_up / _split_phi3_qkv had no guard at all.
  • Reads adjacent to guarded ones are guarded too: the Mixtral and OLMoE router weights and gpt-oss's router weight and expert biases all sat beside guarded expert reads without guards of their own..
  • describe_quantization's class-name fallback was gated on not isinstance(weight, nn.Parameter), which excluded the classes it was written to identify: bitsandbytes' Params4bit and Int8Params are nn.Parameter subclasses. It now excludes the two uninformative names instead, matching how transformers itself keys off those classes. This only ever degraded the error text, never the guard's verdict.
  • process_weights filtered named_parameters() with name.endswith(".weight"), which skips every batched-MoE expert tensor on transformers 5.x (mlp.experts.gate_up_proj). The filter is gone and meta parameters now produce the shared helper's "load with real weights" message instead of being silently skipped, since folding on meta yields meta.
  • Three placements are deliberate, because the obvious one regresses a working flow. The HookedTransformer gate keeps the supported bitsandbytes 4-bit Llama path and rejects only unrecognized quant_methods, which previously fell through because the existing assertion was nested inside if load_in_4bit. The bridge guard sits inside process_weights rather than at enable_compatibility_mode, whose no_processing=True path is exercised by meta-device tests. And the split-function guards run before the nn.Parameter construction, or the existing opaque RuntimeError would win and the guard would never fire.
  • Meta-device tensors are reported as an unmaterialized load, not a quantization. They are float and full-width, so they pass the dtype checks, and naming a quantization there would send the reader after the wrong cause.
  • The MXFP4 gate no longer costs a fetch, so its architecture check is gone. _mxfp4_dequantize_config short-circuited on original_architecture == "GptOssForCausalLM" purely to avoid an AutoConfig round trip. But convert_hf_model_config has already fetched that config moments earlier and discarded it, so the method is now captured there onto cfg.quantization_method and the gate reads it for free.
  • Both MoE converters that had bespoke guards are now on the shared helper, closing the same FP8 hole in each. Mixtral's test moved from asserting the old message to parametrizing over int8/uint8/float8; gpt-oss keeps its actionable Mxfp4Config(dequantize=True) recipe via a new optional remedy argument, reworded so it does not assert MXFP4 when the guard may equally have caught int8 or FP8. All three converters (mixtral, olmoe, openai) now guard both fused expert tensors. down_proj is sliced too, and a packed one would have dropped its scales just as silently.
  • Tests build every quantized shape in memory covering packed integer storage, a torch.Tensor subclass holding packed bytes, a non-tensor wrapper, FP8, and meta. Each carries a positive control asserting float16/bfloat16/float32/float64 still pass, since a guard that fires on those would break the quantized-forward paths.

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4
jlarson4 merged commit 13a28dc into dev Aug 14, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant