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

    modules = list(state.model.modules())
    for module in tqdm(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._parent_args_cache.clear()
    self._smooth_activation_means.clear()
    self._resolved_mappings.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)

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

    self._set_resolved_mappings(state.model)

    return True

@model_validator(mode="after")
def validate_model_after(model: "AWQModifier") -> "AWQModifier":
    """
    Confirm only one configuration for group_size, symmetric, and num_bits,
    as AWQ algorithm depends on it
    Confirm no activation quantization, as AWQ only works with WNA16
    """
    config = model.resolve_quantization_config()

    num_bits_set = set(
        group.weights.num_bits
        for group in config.config_groups.values()
        if group.weights is not None
    )
    assert (
        len(num_bits_set) == 1
    ), "In AWQ, all config groups must use the same configuration for num_bits"

    model._num_bits = next(iter(num_bits_set))

    symmetric_set = set(
        group.weights.symmetric
        for group in config.config_groups.values()
        if group.weights is not None
    )
    assert (
        len(symmetric_set) == 1
    ), "In AWQ, all config groups must use the same configuration for symmetric"

    model._symmetric = next(iter(symmetric_set))

    group_size_set = set(
        group.weights.group_size
        for group in config.config_groups.values()
        if group.weights is not None
    )
    assert (
        len(group_size_set) == 1
    ), "In AWQ, all config groups must use the same configuration for group_size"

    model._group_size = next(iter(group_size_set))

    in_num_bits_set = set(
        group.input_activations.num_bits
        for group in config.config_groups.values()
        if group.input_activations is not None
    )
    assert len(in_num_bits_set) == 0 or in_num_bits_set == {16}, (
        "AWQ activations must be 16-bit precision, "
        f"input activations {in_num_bits_set} not allowed"
    )

    out_num_bits_set = set(
        group.output_activations.num_bits
        for group in config.config_groups.values()
        if group.output_activations is not None
    )
    assert len(out_num_bits_set) == 0 or out_num_bits_set == {16}, (
        "AWQ activations must be 16-bit precision, "
        f"output activations {out_num_bits_set} not allowed"
    )

    return model