llmcompressor.modifiers.quantization.quantization

def end_calibration(self, model: torch.nn.Module):
    """
    Remove calibration hooks and observers, and set the model status to frozen.
    Keep quantization enabled for future operations

    :param model: model to end calibration for
    """
    self.remove_hooks(self._calibration_hooks)
    for _, module in match_named_modules(model, self.resolved_targets, self.ignore):
        freeze_module_quantization(module)  # remove observers

    model.apply(enable_quantization)  # keep quantization enabled

def has_config(self) -> bool:
    """
    Determine if the user has specified a quantization config on this modifier
    """
    return not (
        self.config_groups is None
        and self.targets == ["Linear"]
        and self.ignore == []
        and self.scheme is None
        and self.kv_cache_scheme is None
    )

def initialize_quantization(self, model: torch.nn.Module):
    """
    Attach quantization schemes to modules in the model according to
    the quantization config specified on this modifier

    :param model: model to attach schemes and observers to
    """

    for _, module in match_named_modules(model, self.resolved_targets, self.ignore):
        reset_quantization_status(module)  # reset any previously applied qconfigs

    apply_quantization_config(model, self.resolved_config)

    if not self.bypass_divisibility_checks:
        validate_group_size_divisibility(model, self.resolved_targets, self.ignore)

    # disable quantization until calibration
    model.apply(disable_quantization)

def resolve_quantization_config(self) -> QuantizationConfig:
    """
    Returns the quantization config specified by this modifier
    """
    scheme = self.scheme
    targets = self.targets
    config_groups = self.config_groups
    kv_cache_scheme = self.kv_cache_scheme
    ignore = self.ignore

    if scheme is not None and config_groups is not None:
        raise ValueError("Please specify either `scheme` or `config_groups`")

    if scheme is not None:
        # takes precedence over config_groups

        if isinstance(scheme, str) and is_preset_scheme(scheme):
            # attach targets to scheme
            scheme = {scheme: targets}

        config_groups = {}
        for idx, key in enumerate(scheme.keys()):
            if is_preset_scheme(key):
                scheme_obj = preset_name_to_scheme(key, scheme[key])
            else:
                scheme_obj = QuantizationScheme.model_validate(
                    {"targets": scheme[key], **scheme}
                )

            # Apply observer overrides if specified
            scheme_obj = self._apply_observer_overrides(scheme_obj)

            group_name = f"group_{idx}"
            config_groups[group_name] = scheme_obj

    if config_groups is None or len(config_groups) == 0:
        default_quant_scheme = QuantizationScheme(targets=targets)
        # Apply observer overrides to default scheme as well
        default_quant_scheme = self._apply_observer_overrides(default_quant_scheme)
        config_groups = {"group_0": default_quant_scheme}
    elif scheme is None:
        # Apply observer overrides to all config groups when config_groups
        # was provided directly (not derived from scheme)
        for scheme_obj in config_groups.values():
            self._apply_observer_overrides(scheme_obj)

    return QuantizationConfig(
        config_groups=config_groups,
        kv_cache_scheme=kv_cache_scheme,
        quantization_status=QuantizationStatus.INITIALIZED,
        ignore=ignore,
    )

def start_calibration(self, model: torch.nn.Module):
    """
    Attach observers, register activation calibration hooks (including
    kv_cache quantization) and enable quantization as we calibrate

    :param model: model to prepare for calibration
    """
    targets = match_named_modules(model, self.resolved_targets, self.ignore)
    if targets_embeddings(model, targets):
        untie_word_embeddings(model)

    for _, module in match_named_modules(model, self.resolved_targets, self.ignore):
        self._initialize_observers(module)
        self._calibration_hooks |= self._initialize_hooks(module)
        apply_calibration_status(module)

    model.apply(enable_quantization)  # quantize at the same time as calibrate

@field_validator("observer", mode="before")
def validate_observer(cls, value: Any) -> Optional[Dict[str, str]]:
    """
    Validate observer dictionary format. Accepts keys: 'weights', 'input', 'output'
    """
    if value is None:
        return value

    if not isinstance(value, dict):
        raise ValueError("`observer` must be a dictionary")

    valid_keys = {"weights", "input", "output"}
    for key in value.keys():
        if key not in valid_keys:
            raise ValueError(
                f"Invalid observer key '{key}'. Valid keys are: {valid_keys}"
            )
        if not isinstance(value[key], str):
            raise ValueError(f"Observer value for '{key}' must be a string")

    return value

def on_end(self, state: State, event: Event, **kwargs):
    """
    Finish calibrating by removing observers and calibration hooks
    """
    self.ended_ = True
    QuantizationMixin.end_calibration(
        self, state.model
    )  # keep quantization enabled

def on_initialize(self, state: State, **kwargs) -> bool:
    """
    Prepare to calibrate activations and weights

    According to the quantization config, a quantization scheme is attached to each
    targeted module. The module's forward call is also overwritten to perform
    quantization to inputs, weights, and outputs.

    Then, according to the module's quantization scheme, observers and calibration
    hooks are added. These hooks are disabled until the modifier starts.
    """
    if not QuantizationMixin.has_config(self):
        raise ValueError(
            "QuantizationModifier requires that quantization fields be specified"
        )
    QuantizationMixin.initialize_quantization(self, state.model)

    return True

def on_start(self, state: State, event: Event, **kwargs):
    """
    Begin calibrating activations and weights. Calibrate weights only once on start
    """
    self.started_ = True
    QuantizationMixin.start_calibration(self, state.model)

    named_modules = list(
        match_named_modules(state.model, self.resolved_targets, self.ignore)
    )
    # TODO: this step can be combined with update_weight_zp_scale
    # once update_fused_layer_weight_global_scales is removed
    # and not required by vLLM
    for _, module in named_modules:
        update_weight_global_scale(module)

    # NOTE: update_fused_layer_weight_global_scales operates on Attention
    # and MLP layers, not quantizable Linear layers. Rather than running
    # on targeted modules, we need to run on all modules.
    # Because this call is idempotent, setting all global_scales to the
    # min value, it is ok to run potentially multiple times for all modules
    for module in state.model.modules():
        update_fused_layer_weight_global_scales(module)

    for _, module in tqdm.tqdm(named_modules, desc="Calibrating weights"):
        update_weight_zp_scale(module)