From 9aa42b5a5b85db5f51e5ac36778cb753fcdbd221 Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 1 Jun 2026 14:32:21 +0200 Subject: [PATCH 1/2] Quantize models on `__init__` / `from_config` (not just `from_pretrained`) Constructing a model directly (`cls(config)` / `AutoModel.from_config`) with a `quantization_config` set on the config currently ignores it and returns a full-precision model. This wires up on-the-fly quantization for that path, mirroring `from_pretrained`. `PreTrainedModel.post_init` now calls `_maybe_quantize_on_init`, which (only on the construction root, and only for real non-meta weights) builds the `HfQuantizer`, captures the random-init weights, swaps modules via `preprocess_model`, and routes the weights back through the exact same `convert_and_load_state_dict_in_model` loader used by `from_pretrained` so the result is identical to loading those weights quantized. Gated by a new opt-in `HfQuantizer.supports_quantize_on_init` flag (enabled for FineGrainedFP8); quantizers needing calibration stay opt-out. A meta-device guard keeps the `from_pretrained` path untouched. --- src/transformers/modeling_utils.py | 71 ++++++++++++++++++- src/transformers/quantizers/base.py | 6 ++ .../quantizers/quantizer_finegrained_fp8.py | 1 + .../quantization/finegrained_fp8/test_fp8.py | 50 +++++++++++++ 4 files changed, 127 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 13c06973b4a9..ed079db18d40 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -95,7 +95,7 @@ from .monkey_patching import apply_patches, patch_output_recorders from .pytorch_utils import id_tensor_storage from .quantizers import HfQuantizer -from .quantizers.auto import get_hf_quantizer +from .quantizers.auto import AutoHfQuantizer, get_hf_quantizer from .quantizers.quantizers_utils import get_module_from_name from .safetensors_conversion import auto_conversion from .utils import ( @@ -1346,6 +1346,15 @@ def __init__(self, config: PreTrainedConfig, *inputs, **kwargs): self.config = config self.name_or_path = config.name_or_path + # Claim the "construction root" for on-the-fly quantization (see `_maybe_quantize_on_init`). The outermost + # model's `__init__` runs before its submodules', so the first `PreTrainedModel` built in the tree claims the + # root; nested submodels (which share the same `config` object) see the marker already set and won't re-claim it. + if ( + getattr(config, "quantization_config", None) is not None + and getattr(config, "_quantization_init_root", None) is None + ): + config._quantization_init_root = id(self) + # Check the attention implementation is supported, or set it if not yet set (on the internal attr, to avoid # setting it recursively) self.config._attn_implementation_internal = self._check_and_adjust_attn_implementation( @@ -1440,6 +1449,11 @@ def post_init(self): self.init_weights() self._backward_compatibility_gradient_checkpointing() + # If the config carries a `quantization_config`, quantize the freshly-initialized model in-place. This is a + # no-op unless `self` is the construction root (see `__init__`) and the model holds real (non-meta) weights, + # so `from_pretrained` - which builds on meta and runs its own quantization - is left untouched. + self._maybe_quantize_on_init() + @property def tp_plan(self) -> dict[str, str]: """ @@ -3768,6 +3782,61 @@ def _get_dtype_plan(self, dtype: torch.dtype) -> dict: return dtype_plan + def _maybe_quantize_on_init(self): + """Quantize a freshly-constructed model in-place when `config.quantization_config` is set. + + This mirrors what `from_pretrained` does for on-the-fly quantization (build the `HfQuantizer`, swap the + modules via `preprocess_model`, quantize the weights, then `postprocess_model`), but for the + construct-from-scratch path (`cls(config)` / `AutoModel.from_config`) where there is no checkpoint. The + randomly-initialized weights are routed back through the exact same quantization-aware loader used by + `from_pretrained`, so the resulting model is byte-for-byte equivalent to loading those weights quantized. + + It is a no-op unless all of the following hold: + - `self` is the construction root (the outermost model, see `__init__`); + - the model holds real (non-meta) weights - `from_pretrained` builds on meta and quantizes itself, so we + must not touch that path; + - the quantizer opts in via `HfQuantizer.supports_quantize_on_init` (quantizers that need calibration + data, e.g. AWQ/GPTQ, stay opt-out). + """ + quantization_config = getattr(self.config, "quantization_config", None) + if quantization_config is None or getattr(self.config, "_quantization_init_root", None) != id(self): + return + try: + # `from_pretrained` builds the skeleton on meta and runs its own quantization during weight loading; + # detect that case and bail out so we only handle the construct-from-scratch path. + if any(param.device.type == "meta" for param in self.parameters()): + return + + hf_quantizer = AutoHfQuantizer.from_config(quantization_config, pre_quantized=False) + if not getattr(hf_quantizer, "supports_quantize_on_init", False): + return + hf_quantizer.validate_environment(device_map=None, weights_only=True) + + dtype = self.dtype + # Capture the full-precision random weights *before* swapping modules - `preprocess_model` replaces the + # `nn.Linear`/experts with fresh (empty) quantized modules, which would otherwise discard them. + state_dict = self.state_dict() + with torch.device("meta"): + hf_quantizer.preprocess_model( + model=self, dtype=dtype, device_map=None, checkpoint_files=None, use_kernels=False + ) + + load_config = LoadStateDictConfig( + dtype=dtype, + dtype_plan=self._get_dtype_plan(dtype), + hf_quantizer=hf_quantizer, + weight_mapping=get_model_conversion_mapping(self, None, hf_quantizer), + ) + convert_and_load_state_dict_in_model(self, state_dict, load_config, tp_plan=None) + + self.hf_quantizer = hf_quantizer + hf_quantizer.postprocess_model(self) + finally: + # Always release the marker so the next independent construction can re-claim the root, even if the + # quantizer bailed out (no GPU, not opted in, ...) or raised. + if hasattr(self.config, "_quantization_init_root"): + del self.config._quantization_init_root + def set_use_kernels(self, use_kernels, kernel_config: KernelConfig | None = None): """ Set whether or not to use the `kernels` library to kernelize some layers of the model. diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index fbebea2977be..d377915c64e6 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -81,9 +81,15 @@ class HfQuantizer(ABC): The quantization config that defines the quantization parameters of your model that you want to quantize. requires_calibration (`bool`): Whether the quantization method requires to calibrate the model before using it. + supports_quantize_on_init (`bool`): + Whether the quantization method can quantize a freshly-constructed model in-place (i.e. directly from + `cls(config)` / `AutoModel.from_config`, with no checkpoint), by quantizing the randomly-initialized + weights. Quantizers that need calibration data or otherwise assume a pre-quantized checkpoint should + leave this `False`. See `PreTrainedModel._maybe_quantize_on_init`. """ requires_calibration = False + supports_quantize_on_init = False def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): self.quantization_config = quantization_config diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index be10624d4842..0e1074e3f483 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -22,6 +22,7 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): """ requires_calibration = False + supports_quantize_on_init = True quantization_config: "FineGrainedFP8Config" def __init__(self, quantization_config, **kwargs): diff --git a/tests/quantization/finegrained_fp8/test_fp8.py b/tests/quantization/finegrained_fp8/test_fp8.py index 8783fdb5830d..02d04aff632a 100644 --- a/tests/quantization/finegrained_fp8/test_fp8.py +++ b/tests/quantization/finegrained_fp8/test_fp8.py @@ -453,3 +453,53 @@ def test_linear_with_diff_feature_size_preserves_shape(self): x_ = linear(x) self.assertEqual(x_.shape, (1, 5, 256)) + + +@require_accelerate +@require_torch_accelerator +@unittest.skipIf( + get_device_properties()[0] == "cuda" + and (get_device_properties()[1] < 8 or (get_device_properties()[1] == 8 and get_device_properties()[2] < 9)), + "Skipping FP8QuantizeOnInitTest because it is not supported on GPU with capability < 8.9", +) +class FP8QuantizeOnInitTest(unittest.TestCase): + """Quantize a freshly-constructed model (no checkpoint) when `config.quantization_config` is set.""" + + model_name = "meta-llama/Llama-3.2-1B" + + def _build_config(self): + config = AutoConfig.from_pretrained(self.model_name) + config.quantization_config = FineGrainedFP8Config().to_dict() + return config + + def test_direct_init_quantizes(self): + """`cls(config)` with a `quantization_config` returns a quantized model with fp8 weights + scales.""" + config = self._build_config() + model = AutoModelForCausalLM.from_config(config).to(torch_device) + + self.assertTrue(getattr(model, "is_quantized", False)) + self.assertIsInstance(model.hf_quantizer, FineGrainedFP8HfQuantizer) + # Marker used to run quantization only on the construction root must not leak onto the config. + self.assertFalse(hasattr(model.config, "_quantization_init_root")) + + weight = model.model.layers[0].self_attn.q_proj.weight + weight_scale_inv = model.model.layers[0].self_attn.q_proj.weight_scale_inv + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(weight_scale_inv.dtype, torch.float32) + + def test_save_then_load_round_trip(self): + """A model quantized on init can be saved and reloaded with `from_pretrained`.""" + config = self._build_config() + model = AutoModelForCausalLM.from_config(config).to(torch_device) + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname) + reloaded = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=torch_device) + self.assertTrue(getattr(reloaded, "is_quantized", False)) + self.assertEqual(reloaded.model.layers[0].self_attn.q_proj.weight.dtype, torch.float8_e4m3fn) + + def test_no_quantization_config_is_noop(self): + """Without a `quantization_config`, construction is unchanged (no quantizer attached).""" + config = AutoConfig.from_pretrained(self.model_name) + model = AutoModelForCausalLM.from_config(config) + self.assertFalse(getattr(model, "is_quantized", False)) + self.assertIsNone(getattr(model, "hf_quantizer", None)) From 28f63581af114184a9e7564b4aa77e3e18bd073a Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 2 Jun 2026 11:58:50 +0200 Subject: [PATCH 2/2] Move quantize-on-init into quantizers_utils, gate on requires_calibration Drop the `supports_quantize_on_init` opt-in flag: all non-calibration quantizers can quantize freshly-initialized weights, so reuse the existing `requires_calibration` to gate instead. `_maybe_quantize_on_init` moves out of `PreTrainedModel` into `quantizers_utils.maybe_quantize_on_init`. Resolve the quantizer class without instantiating (new `AutoHfQuantizer.resolve`, factored out of `from_config`) so we can check `requires_calibration` before any backend import. Calibration-only methods (AWQ, GPTQ, ...) now emit a warning instead of silently returning a full-precision model. Read the dtype off the freshly-initialized parameters with a default-dtype fallback rather than `self.dtype`, which raises StopIteration on a param-less model; the parameters are the source of truth that matches the captured state_dict. --- src/transformers/modeling_utils.py | 63 +------------- src/transformers/quantizers/auto.py | 13 ++- src/transformers/quantizers/base.py | 6 -- .../quantizers/quantizer_finegrained_fp8.py | 1 - .../quantizers/quantizers_utils.py | 83 +++++++++++++++++++ 5 files changed, 98 insertions(+), 68 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index ed079db18d40..720cad5131b1 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -95,8 +95,8 @@ from .monkey_patching import apply_patches, patch_output_recorders from .pytorch_utils import id_tensor_storage from .quantizers import HfQuantizer -from .quantizers.auto import AutoHfQuantizer, get_hf_quantizer -from .quantizers.quantizers_utils import get_module_from_name +from .quantizers.auto import get_hf_quantizer +from .quantizers.quantizers_utils import get_module_from_name, maybe_quantize_on_init from .safetensors_conversion import auto_conversion from .utils import ( ADAPTER_SAFE_WEIGHTS_NAME, @@ -1346,7 +1346,7 @@ def __init__(self, config: PreTrainedConfig, *inputs, **kwargs): self.config = config self.name_or_path = config.name_or_path - # Claim the "construction root" for on-the-fly quantization (see `_maybe_quantize_on_init`). The outermost + # Claim the "construction root" for on-the-fly quantization (see `maybe_quantize_on_init`). The outermost # model's `__init__` runs before its submodules', so the first `PreTrainedModel` built in the tree claims the # root; nested submodels (which share the same `config` object) see the marker already set and won't re-claim it. if ( @@ -1452,7 +1452,7 @@ def post_init(self): # If the config carries a `quantization_config`, quantize the freshly-initialized model in-place. This is a # no-op unless `self` is the construction root (see `__init__`) and the model holds real (non-meta) weights, # so `from_pretrained` - which builds on meta and runs its own quantization - is left untouched. - self._maybe_quantize_on_init() + maybe_quantize_on_init(self) @property def tp_plan(self) -> dict[str, str]: @@ -3782,61 +3782,6 @@ def _get_dtype_plan(self, dtype: torch.dtype) -> dict: return dtype_plan - def _maybe_quantize_on_init(self): - """Quantize a freshly-constructed model in-place when `config.quantization_config` is set. - - This mirrors what `from_pretrained` does for on-the-fly quantization (build the `HfQuantizer`, swap the - modules via `preprocess_model`, quantize the weights, then `postprocess_model`), but for the - construct-from-scratch path (`cls(config)` / `AutoModel.from_config`) where there is no checkpoint. The - randomly-initialized weights are routed back through the exact same quantization-aware loader used by - `from_pretrained`, so the resulting model is byte-for-byte equivalent to loading those weights quantized. - - It is a no-op unless all of the following hold: - - `self` is the construction root (the outermost model, see `__init__`); - - the model holds real (non-meta) weights - `from_pretrained` builds on meta and quantizes itself, so we - must not touch that path; - - the quantizer opts in via `HfQuantizer.supports_quantize_on_init` (quantizers that need calibration - data, e.g. AWQ/GPTQ, stay opt-out). - """ - quantization_config = getattr(self.config, "quantization_config", None) - if quantization_config is None or getattr(self.config, "_quantization_init_root", None) != id(self): - return - try: - # `from_pretrained` builds the skeleton on meta and runs its own quantization during weight loading; - # detect that case and bail out so we only handle the construct-from-scratch path. - if any(param.device.type == "meta" for param in self.parameters()): - return - - hf_quantizer = AutoHfQuantizer.from_config(quantization_config, pre_quantized=False) - if not getattr(hf_quantizer, "supports_quantize_on_init", False): - return - hf_quantizer.validate_environment(device_map=None, weights_only=True) - - dtype = self.dtype - # Capture the full-precision random weights *before* swapping modules - `preprocess_model` replaces the - # `nn.Linear`/experts with fresh (empty) quantized modules, which would otherwise discard them. - state_dict = self.state_dict() - with torch.device("meta"): - hf_quantizer.preprocess_model( - model=self, dtype=dtype, device_map=None, checkpoint_files=None, use_kernels=False - ) - - load_config = LoadStateDictConfig( - dtype=dtype, - dtype_plan=self._get_dtype_plan(dtype), - hf_quantizer=hf_quantizer, - weight_mapping=get_model_conversion_mapping(self, None, hf_quantizer), - ) - convert_and_load_state_dict_in_model(self, state_dict, load_config, tp_plan=None) - - self.hf_quantizer = hf_quantizer - hf_quantizer.postprocess_model(self) - finally: - # Always release the marker so the next independent construction can re-claim the root, even if the - # quantizer bailed out (no GPU, not opted in, ...) or raised. - if hasattr(self.config, "_quantization_init_root"): - del self.config._quantization_init_root - def set_use_kernels(self, use_kernels, kernel_config: KernelConfig | None = None): """ Set whether or not to use the `kernels` library to kernelize some layers of the model. diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index a972c6f637bb..ba03198e9c27 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -182,7 +182,12 @@ class AutoHfQuantizer: """ @classmethod - def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwargs): + def resolve(cls, quantization_config: QuantizationConfigMixin | dict): + """Resolve a (possibly dict) quantization config to its `(config_object, HfQuantizer subclass)`. + + Does not instantiate the quantizer, so callers can inspect class-level attributes (e.g. + `requires_calibration`) without triggering backend imports or `__init__` side effects. + """ # Convert it to a QuantizationConfig if the q_config is a dict if isinstance(quantization_config, dict): quantization_config = AutoQuantizationConfig.from_dict(quantization_config) @@ -207,7 +212,11 @@ def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwar f" {list(AUTO_QUANTIZER_MAPPING.keys())}" ) - target_cls = AUTO_QUANTIZER_MAPPING[quant_method] + return quantization_config, AUTO_QUANTIZER_MAPPING[quant_method] + + @classmethod + def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwargs): + quantization_config, target_cls = cls.resolve(quantization_config) return target_cls(quantization_config, **kwargs) @classmethod diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index d377915c64e6..fbebea2977be 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -81,15 +81,9 @@ class HfQuantizer(ABC): The quantization config that defines the quantization parameters of your model that you want to quantize. requires_calibration (`bool`): Whether the quantization method requires to calibrate the model before using it. - supports_quantize_on_init (`bool`): - Whether the quantization method can quantize a freshly-constructed model in-place (i.e. directly from - `cls(config)` / `AutoModel.from_config`, with no checkpoint), by quantizing the randomly-initialized - weights. Quantizers that need calibration data or otherwise assume a pre-quantized checkpoint should - leave this `False`. See `PreTrainedModel._maybe_quantize_on_init`. """ requires_calibration = False - supports_quantize_on_init = False def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): self.quantization_config = quantization_config diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 0e1074e3f483..be10624d4842 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -22,7 +22,6 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): """ requires_calibration = False - supports_quantize_on_init = True quantization_config: "FineGrainedFP8Config" def __init__(self, quantization_config, **kwargs): diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 0e90e238ec4a..d83d26d9cda7 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -14,6 +14,11 @@ import re from typing import Any +from ..utils import logging + + +logger = logging.get_logger(__name__) + def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: if "." in tensor_name: @@ -39,3 +44,81 @@ def should_convert_module(full_name, patterns: list[str] | None = None): for key in patterns ) return not should_not_convert + + +def maybe_quantize_on_init(model): + """Quantize a freshly-constructed model in-place when `config.quantization_config` is set. + + Mirrors the on-the-fly quantization that `from_pretrained` performs (build the `HfQuantizer`, swap the + modules via `preprocess_model`, route the weights through the quantization-aware loader, then + `postprocess_model`), but for the construct-from-scratch path (`cls(config)` / `AutoModel.from_config`) + where there is no checkpoint. The randomly-initialized weights are routed through the exact same loader + used by `from_pretrained`, so the result is equivalent to loading those weights quantized. + + No-op unless `model` is the construction root (the outermost model, see `PreTrainedModel.__init__`) and + holds real (non-meta) weights - `from_pretrained` builds on meta and quantizes itself, so that path is + left untouched. Quantizers that need calibration data / a pre-quantized checkpoint (AWQ, GPTQ, ...) cannot + quantize random weights, so the model is left in full precision with a warning. + """ + import torch + + from ..conversion_mapping import get_model_conversion_mapping + from ..core_model_loading import convert_and_load_state_dict_in_model + from ..modeling_utils import LoadStateDictConfig + from .auto import AutoHfQuantizer + + config = model.config + quantization_config = getattr(config, "quantization_config", None) + if quantization_config is None or getattr(config, "_quantization_init_root", None) != id(model): + return + try: + # `from_pretrained` builds the skeleton on meta and quantizes during weight loading; bail out here so + # we only handle the construct-from-scratch path. + if any(param.device.type == "meta" for param in model.parameters()): + return + + # Resolve the quantizer class without instantiating it, so we can check `requires_calibration` before + # triggering any backend import or `__init__` side effect. + quantization_config, quantizer_cls = AutoHfQuantizer.resolve(quantization_config) + if quantizer_cls.requires_calibration: + # Calibration-only methods (AWQ, GPTQ, ...) cannot quantize freshly-initialized weights. Warn instead + # of silently returning a full-precision model - the user set a `quantization_config` and expects a + # quantized model. + logger.warning( + f"A `quantization_config` ({quantization_config.quant_method}) was set on a model built from " + "scratch (`from_config` / `cls(config)`), but this method requires a pre-quantized checkpoint or " + "calibration data and cannot quantize randomly-initialized weights. The model was left in full " + "precision and was NOT quantized; load a pre-quantized checkpoint with `from_pretrained` instead." + ) + return + + hf_quantizer = quantizer_cls(quantization_config, pre_quantized=False) + hf_quantizer.validate_environment(device_map=None, weights_only=True) + + # Read the dtype off the freshly-initialized parameters so it matches the captured `state_dict`, falling + # back to the default float dtype for a (degenerate) param-less model. `cls(config)` builds in the global + # default (fp32) unless `from_config` opened a dtype context, so this tracks the real weights either way. + dtype = next((p.dtype for p in model.parameters() if p.is_floating_point()), torch.get_default_dtype()) + # Capture the full-precision random weights *before* swapping modules - `preprocess_model` replaces the + # `nn.Linear`/experts with fresh (empty) quantized modules, which would otherwise discard them. + state_dict = model.state_dict() + with torch.device("meta"): + hf_quantizer.preprocess_model( + model=model, dtype=dtype, device_map=None, checkpoint_files=None, use_kernels=False + ) + + load_config = LoadStateDictConfig( + dtype=dtype, + dtype_plan=model._get_dtype_plan(dtype), + hf_quantizer=hf_quantizer, + weight_mapping=get_model_conversion_mapping(model, None, hf_quantizer), + ) + convert_and_load_state_dict_in_model(model, state_dict, load_config, tp_plan=None) + + model.hf_quantizer = hf_quantizer + hf_quantizer.postprocess_model(model) + finally: + # Always release the marker so the next independent construction can re-claim the root, even if the + # quantizer bailed out (no GPU, calibration-only, ...) or raised. + if hasattr(config, "_quantization_init_root"): + del config._quantization_init_root