Refuse quantized weights in weight-space code paths instead of reading them as matrices - #1669
Merged
Conversation
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.
Description
transformer_lens/utilities/quantization.pyholds one predicate: a weight is readable if it is atorch.Tensorwith a floating-point dtype ofitemsize >= 2. It is future-proof against new narrow types without enumerating names. The error names the quantization method, resolved from the HF config'squant_methodthen the weight's class.torch.float8_e4m3fn.is_floating_pointisTrue, so anisinstance + is_floating_pointcheck admits it and FP8 is the only family that survives bothtensor_splitandnn.Parameter()without complaint. Every guard test parametrizes over it for that reason.W_Q/W_K/W_V/W_O,W_in/W_gate/W_out), which returned packed bytes silentlyuse_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.if hf_model is not None:, sofrom_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 surviveload_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 (theirQuantLinearhas no.weightat all)._refuse_unsupported_quantizationnow runs once inget_pretrained_state_dict, covering all 28 converters. It reads the loaded model's config rather thancfg.quantization_method, becauseconvert_hf_model_configinfers 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 withdequantize=Truekeeps working.split_gate_up_matrixwith Phi-3's own splitter, andphi3.py's_split_gate_up/_split_phi3_qkvhad no guard at all.describe_quantization's class-name fallback was gated onnot isinstance(weight, nn.Parameter), which excluded the classes it was written to identify: bitsandbytes'Params4bitandInt8Paramsarenn.Parametersubclasses. 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_weightsfilterednamed_parameters()withname.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.HookedTransformergate keeps the supported bitsandbytes 4-bit Llama path and rejects only unrecognizedquant_methods, which previously fell through because the existing assertion was nested insideif load_in_4bit. The bridge guard sits insideprocess_weightsrather than atenable_compatibility_mode, whoseno_processing=Truepath is exercised by meta-device tests. And the split-function guards run before thenn.Parameterconstruction, or the existing opaqueRuntimeErrorwould win and the guard would never fire._mxfp4_dequantize_configshort-circuited onoriginal_architecture == "GptOssForCausalLM"purely to avoid anAutoConfiground trip. Butconvert_hf_model_confighas already fetched that config moments earlier and discarded it, so the method is now captured there ontocfg.quantization_methodand the gate reads it for free.Mxfp4Config(dequantize=True)recipe via a new optionalremedyargument, 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_projis sliced too, and a packed one would have dropped its scales just as silently.torch.Tensorsubclass 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
Checklist: