transformer_lens.utilities.quantization module

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

TransformerLens supports quantized forward passes: the wrapped HF module dequantizes internally, and the bridge’s forward paths deliberately skip non-floating-point parameters when picking a compute dtype. Those paths must keep working.

What does not work is reading a quantized .weight and doing arithmetic on it directly — reshaping it into per-head matrices, slicing a fused projection, folding LayerNorm into it. There the storage is packed (bitsandbytes 4-bit keeps a [N, 1] uint8 buffer), split from its scales (FP8 keeps a separate weight_scale_inv), or not a tensor at all (MXFP4 wraps a triton-kernels object). Slicing those yields plausible-looking garbage rather than an error, which is the failure mode this module exists to prevent.

transformer_lens.utilities.quantization.describe_quantization(owner: Any) → str

Best-effort name for how owner’s weights are quantized.

Resolution order: the HF config’s declared quant_method, then the weight’s class name (bitsandbytes subclasses are identifiable that way), then a generic fallback.

transformer_lens.utilities.quantization.quantization_method(config: Any) → str | None

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

Accepts a PretrainedConfig or its to_dict() form, and tolerates the nested quantization_config being either shape — both appear in the wild, depending on whether the config was loaded or round-tripped through JSON.

transformer_lens.utilities.quantization.require_readable_weight(weight: Any, *, operation: str, owner: Any = None, remedy: str | None = None) → Tensor

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

operation completes “TransformerLens cannot <operation> because …”; remedy replaces the generic advice — phrase it to hold for ANY quantization, since the caller cannot know which one it caught.

transformer_lens.utilities.quantization.unreadable_weight_reason(weight: Any) → str | None

Why weight cannot be read as a plain matrix (a fragment completing “… cannot be read because <reason>”), or None if it can.

Readable = exactly {fp16, bf16, fp32, fp64}: every 1-byte torch dtype is integer storage or a scale-less narrow float.