llmcompressor.modifiers.awq.base

def on_end(self, state: State, event: Event, **kwargs):
    """
    Finish calibrating by setting scales and zero-points,
     removing observers and calibration hooks
    """
    self._assert_all_activations_consumed()

    self.ended_ = True

    named_modules = list(
        match_named_modules(state.model, self.resolved_targets, self.ignore)
    )

    # For TENSOR_GROUP (nvfp4), calculate global scales after smoothing
    for _, module in tqdm(named_modules, desc="Updating global scales"):
        update_weight_global_scale(module)

    # For TENSOR_GROUP (nvfp4), fuse global scales for attention and MLP layers
    # This is a requirement for vLLM inference.
    for module in tqdm(state.model.modules(), desc="Fusing global scales"):
        update_fused_layer_weight_global_scales(module)

    # Calculate scales and zero points using the fused global scales
    for _, module in tqdm(named_modules, desc="Calibrating weights"):
        update_weight_zp_scale(module)

    QuantizationMixin.end_calibration(self, state.model)

    # remove activation hooks
    self.remove_hooks()

def on_finalize(self, state: State, **kwargs) -> bool:
    """
    Clean up by clearing the activations and mapping data

    :param state: unused
    :return: True
    """
    if not self.ended_:
        self.on_end(state, None)

    self._log_error_metrics()

    self._parent_args_cache.clear()
    self._smooth_activation_means.clear()
    self._resolved_mappings.clear()
    self._error_metrics.clear()

    return True

def on_initialize(self, state: State, **kwargs) -> bool:
    """
    Initialize AWQ on the given state
    Initialize quantization, resolve mappings, cache module kwargs

    :param state: state to run AWQ on
    :return: True on a successful run, False otherwise
    """

    # apply config to model and prepare calibration hooks
    if QuantizationMixin.has_config(self):
        QuantizationMixin.initialize_quantization(self, state.model)

    # Validate that duo_scaling is only used with per-channel quantization
    if self.duo_scaling is not False:
        for _, module in match_named_modules(
            state.model, self.resolved_targets, self.ignore
        ):
            if (
                hasattr(module, "quantization_scheme")
                and hasattr(module.quantization_scheme, "weights")
                and module.quantization_scheme.weights.strategy
                == QuantizationStrategy.TENSOR
            ):
                raise ValueError(
                    "duo_scaling is only supported with per-channel quantization "
                    "strategies (group or channel), but found TENSOR strategy. "
                    "Please set duo_scaling=False or use a per-channel "
                    "quantization strategy."
                )

    if self.mappings is None:
        logger.info("No AWQModifier.mappings provided, inferring from model...")
        self.mappings = get_layer_mappings_from_architecture(
            architecture=state.model.__class__.__name__
        )

    # Set default offload_device
    if self.offload_device == Sentinel("not_provided"):
        # Check if we have a MoE model
        if is_moe_model(state.model):
            self.offload_device = torch.device("cpu")
            logger.info(
                "MoE model detected: setting offload_device to 'cpu' by default "
                "to reduce memory usage. You can override this by explicitly "
                "setting offload_device in your recipe."
            )
        else:
            # For non-MoE models, convert sentinel to None
            # (no offloading by default)
            self.offload_device = None

    self._set_resolved_mappings(state.model)

    return True

@field_validator("duo_scaling")
@classmethod
def validate_duo_scaling(cls, v):
    """Validate that duo_scaling is either True, False, or 'both' (lowercase)"""
    if v not in (True, False, "both"):
        raise ValueError(f"duo_scaling must be True, False, or 'both', got {v!r}")
    return v