Quantize models on __init__ / from_config

src/transformers/modeling_utils.py

@@ -96,7 +96,7 @@
 from .pytorch_utils import id_tensor_storage
 from .quantizers import HfQuantizer
 from .quantizers.auto import get_hf_quantizer
-from .quantizers.quantizers_utils import get_module_from_name
+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,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.
+        maybe_quantize_on_init(self)
+
     @property
     def tp_plan(self) -> dict[str, str]:
         """

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

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

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))