llmcompressor.modifiers.autoround

def apply_autoround(self, state, modules):
    """
    Applies AutoRound quantization tuning on the current decoding layer.

    The tuning logic is as follows:
    for iter in range(iters):
        quant_output = forward(layer, cached_inputs)
        loss = mse_loss(quant_output, original_output)
        loss.backward()
        optimizer.step()
        if loss < best_loss:
            best_params = update_params(layer)

    For more details, please refer to the AutoRound repository:
    https://github.com/intel/auto-round/
    """
    modules = modules or []

    decoding_layers = [m for m in modules if self._is_decoding_layer(m)]
    if len(decoding_layers) == 0:
        return
    if len(decoding_layers) != 1:
        raise ValueError(
            "Only one decoding layer is expected in the modules list, "
            f"found {len(decoding_layers)}."
        )
    decoding_layer = decoding_layers[0]

    logger.info("Applying AutoRound on layer {}", decoding_layer._tmp_name)

    # Build wrapped_model for AutoRound initialization
    wrapped_model = _wrap_decoding_layer(decoding_layer)
    wrapped_model.name_or_path = state.model.name_or_path
    wrapped_model.config = state.model.config

    # Build kwargs for AutoRound initialization
    ar_quant_scheme = self._mapping_config_to_autoround()
    layer_config = self._build_layer_config_for_autoround(wrapped_model)
    ignore_layers = self.get_unquantized_layer_names(decoding_layer)
    kwargs = {
        "tokenizer": "",  # A placeholder
        "scheme": ar_quant_scheme,
        "layer_config": layer_config or None,
        "iters": self.iters,
        "lr": self.lr,
        "enable_torch_compile": self.enable_torch_compile,
        "batch_size": self.batch_size,
        "device_map": self.device_ids,
        "ignore_layers": ",".join(ignore_layers) if ignore_layers else "",
        "disable_opt_rtn": self.disable_opt_rtn,
    }

    llmc_registered_qparams = self._preprocess_qparams(decoding_layer)
    with (
        torch.enable_grad(),
        align_module_device(decoding_layer),
        suspend_offloading(wrapped_model),
    ):
        self._update_device_map_for_dp(kwargs)
        ar = AutoRound(
            model=wrapped_model,
            **kwargs,
        )
        ar.configure_layer_config(enable_gguf_official_mixed=False)
        ar.batch_dim = 0
        first_param = next(decoding_layer.parameters())
        device = first_param.device
        cur_inputs = self._all_module_input[decoding_layer._tmp_name]
        self._set_attention_masks(ar, decoding_layer, cur_inputs)
        decoding_layer.tuning_device = device
        # Only hand device placement to AutoRound when the caller explicitly
        # requested it or when a rank is configured to use a local GPU group.
        auto_offload = False
        needs_multi_gpu = (
            self.device_ids is not None or get_local_gpu_group_size() > 1
        )
        if needs_multi_gpu:
            # Let AutoRound own placement within the rank-local GPU group.
            device = get_main_device()
            decoding_layer.to("cpu")
            auto_offload = True

        # Ensure cached inputs are on the same device as the block.
        # Calibration forward may have run on a different GPU.
        cur_inputs = self._move_inputs_to(cur_inputs, device)
        ar_inputs = [((args, kwargs),) for args, kwargs in cur_inputs]

        q_input, _ = ar.quantize_block(
            block=decoding_layer,
            inputs=ar_inputs,
            q_input=self._q_input,
            device=str(device),
            auto_offload=auto_offload,
        )
        self._q_input = q_input

        decoding_layer = self._unwrapper_quantized_layer(decoding_layer)

    decoding_layer.eval()
    # Update offload parameters and remove temporary attributes
    self._postprocess_qparams(decoding_layer, llmc_registered_qparams)

def on_calibration_end(self, state: State, event: Event, **kwargs):
    """
    Finish calibrating by removing observers and calibration hooks
    """
    QuantizationMixin.end_calibration(self, state.model)
    self._remove_temporary_names(state.model)
    self.remove_hooks()
    self._q_input = None

def on_initialize(self, state: State, **kwargs) -> bool:
    """
    Initialize the model state for quantization and calibration.

    :param state: session state storing input model and calibration data
    """
    # apply config to model and prepare calibration hooks.
    # Wrap in disable_onloading to suppress DistributedCPUCache's
    # per-param broadcast+barrier when creating quant params (scale,
    # zero_point). With LLMCOMPRESSOR_GPUS_PER_GROUP > 1, modules have varying GPU
    # execution devices, causing GPU→CPU copy timing to vary between
    # ranks → broadcast deadlock. Quant params are deterministic —
    # each rank computes identical values, no sync needed.
    if QuantizationMixin.has_config(self):
        from compressed_tensors.offload import disable_onloading

        with disable_onloading():
            QuantizationMixin.initialize_quantization(self, state.model)

    # prepare module names
    self._add_temporary_names(state.model)
    # freeze all model parameters
    for _, param in state.model.named_parameters():
        param.requires_grad_(False)

    self._sequential_targets = infer_sequential_targets(
        state.model, sequential_targets=kwargs.get("sequential_targets")
    )
    return True

def start_calibration(self, model: torch.nn.Module):
    """
    Register activation calibration hooks and enable quantization as we calibrate

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

    for _, module in match_named_modules(model, self.targets, self.ignore):
        # skip register observers for auto-round
        apply_calibration_status(module)

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