diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 13c06973b4a9..720cad5131b1 100644 --- a/src/transformers/modeling_utils.py +++ b/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]: """ 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/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 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))