Coverage for transformer_lens/utilities/quantization.py: 100%

1"""Guards for weight-space code paths that cannot read quantized weights. 

2 

3TransformerLens supports quantized *forward* passes: the wrapped HF module 

4dequantizes internally, and the bridge's forward paths deliberately skip 

5non-floating-point parameters when picking a compute dtype. Those paths must 

6keep working. 

7 

8What does not work is reading a quantized ``.weight`` and doing arithmetic on it 

9directly — reshaping it into per-head matrices, slicing a fused projection, 

10folding LayerNorm into it. There the storage is packed (bitsandbytes 4-bit keeps 

11a ``[N, 1]`` uint8 buffer), split from its scales (FP8 keeps a separate 

12``weight_scale_inv``), or not a tensor at all (MXFP4 wraps a triton-kernels 

13object). Slicing those yields plausible-looking garbage rather than an error, 

14which is the failure mode this module exists to prevent. 

15""" 

16 

17from __future__ import annotations 

18 

19from typing import Any, Optional 

20 

21import torch 

22 

23 

24def unreadable_weight_reason(weight: Any) -> Optional[str]: 

25 """Why ``weight`` cannot be read as a plain matrix (a fragment completing 

26 "... cannot be read because <reason>"), or None if it can. 

27 

28 Readable = exactly {fp16, bf16, fp32, fp64}: every 1-byte torch dtype is 

29 integer storage or a scale-less narrow float. 

30 """ 

31 if not isinstance(weight, torch.Tensor): 

32 # MXFP4 and similar hand back a wrapper object holding the packed 

33 # payload plus scales. It is often *named* Tensor, which makes the 

34 # resulting TypeError look like a torch bug — hence the qualified name, 

35 # which is the only thing distinguishing it from the real one. 

36 cls = type(weight) 

37 return f"it is a {cls.__module__}.{cls.__qualname__}, not a torch.Tensor" 

38 if weight.device.type == "meta": 

39 return "it is on the meta device, so no values have been materialized" 

40 if not weight.dtype.is_floating_point: 

41 # bitsandbytes int8/4-bit, GPTQ/AWQ int32, BitNet packed uint8. Note 

42 # that bnb's Params4bit IS a torch.Tensor subclass, so only the dtype 

43 # distinguishes it — and its shape is the packed [N, 1], not [out, in]. 

44 return f"its dtype is {weight.dtype}, which is packed integer storage" 

45 if weight.dtype.itemsize < 2: 

46 # float8_e4m3fn and friends report is_floating_point=True, so a plain 

47 # float check lets them through. They are the only family that slices 

48 # and even survives nn.Parameter() silently, so omitting this branch 

49 # leaves the worst case uncaught. 

50 return ( 

51 f"its dtype is {weight.dtype}, a narrow float whose values are " 

52 "meaningless without the separate scale tensor stored beside them " 

53 "(e.g. weight_scale_inv)" 

54 ) 

55 return None 

56 

57 

58def _is_meta(weight: Any) -> bool: 

59 return isinstance(weight, torch.Tensor) and weight.device.type == "meta" 

60 

61 

62def quantization_method(config: Any) -> Optional[str]: 

63 """The ``quant_method`` declared on an HF config, or None if unquantized. 

64 

65 Accepts a ``PretrainedConfig`` or its ``to_dict()`` form, and tolerates the 

66 nested ``quantization_config`` being either shape — both appear in the wild, 

67 depending on whether the config was loaded or round-tripped through JSON. 

68 """ 

69 if config is None: 

70 return None 

71 if isinstance(config, dict): 

72 quant_config = config.get("quantization_config") 

73 else: 

74 quant_config = getattr(config, "quantization_config", None) 

75 if quant_config is None: 

76 return None 

77 if isinstance(quant_config, dict): 

78 method = quant_config.get("quant_method") 

79 else: 

80 method = getattr(quant_config, "quant_method", None) 

81 # HF stores this as a str or a str-valued QuantizationMethod enum; require 

82 # that rather than str()-ing whatever turned up, or a stub object's repr 

83 # ends up quoted back at the user as a method name. 

84 method = getattr(method, "value", method) 

85 return method if isinstance(method, str) else None 

86 

87 

88def describe_quantization(owner: Any) -> str: 

89 """Best-effort name for how ``owner``'s weights are quantized. 

90 

91 Resolution order: the HF config's declared ``quant_method``, then the 

92 weight's class name (bitsandbytes subclasses are identifiable that way), 

93 then a generic fallback. 

94 """ 

95 for holder in (owner, getattr(owner, "original_component", None)): 

96 method = quantization_method(getattr(holder, "config", None)) 

97 if method is not None: 

98 return method 

99 weight = getattr(owner, "weight", None) 

100 if weight is not None: 

101 # Exclude the two uninformative names rather than nn.Parameter itself: 

102 # bitsandbytes' Params4bit and Int8Params ARE Parameter subclasses, so 

103 # an isinstance test excluded exactly the classes this identifies. 

104 # transformers keys off the same class names (integrations/bitsandbytes). 

105 name = type(weight).__name__ 

106 if name not in ("Tensor", "Parameter"): 

107 return name 

108 return "an unknown quantization" 

109 

110 

111_GENERIC_REMEDY = ( 

112 "Weight-space operations need dequantized weights — reload the model " 

113 "without a quantization_config (or with dequantize enabled). Quantized " 

114 "*forward* passes remain supported." 

115) 

116 

117 

118def require_readable_weight( 

119 weight: Any, *, operation: str, owner: Any = None, remedy: Optional[str] = None 

120) -> torch.Tensor: 

121 """Return ``weight`` if it can be read as a plain matrix, else raise loudly. 

122 

123 ``operation`` completes "TransformerLens cannot <operation> because ..."; 

124 ``remedy`` replaces the generic advice — phrase it to hold for ANY 

125 quantization, since the caller cannot know which one it caught. 

126 """ 

127 reason = unreadable_weight_reason(weight) 

128 if reason is None: 

129 assert isinstance(weight, torch.Tensor) 

130 return weight 

131 if _is_meta(weight): 

132 # Not a quantization problem — naming one here would send the reader 

133 # after the wrong cause. This is an offloaded or never-materialized load. 

134 raise NotImplementedError( 

135 f"TransformerLens cannot {operation} because {reason}. Load the " 

136 "model with real weights (no meta-device or disk offload) before " 

137 "reading them." 

138 ) 

139 method = describe_quantization(owner) if owner is not None else "an unknown quantization" 

140 raise NotImplementedError( 

141 f"TransformerLens cannot {operation} because {reason} " 

142 f"(quantization: {method}). {remedy or _GENERIC_REMEDY}" 

143 )