llmcompressor.modifiers.autoround.base

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,
        # Keep fp_outputs and best_params on CPU so that the previous block's
        # cache_device tensors don't accumulate on GPU between blocks.
        "low_gpu_mem_usage": True,
    }

    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
        # Remove this layer's input-capture hook before optimization begins
        # so subsequent forward passes during the SignSGD loop do not
        # re-populate _all_module_input with hidden_states tensors.
        layer_name = decoding_layer._tmp_name
        if layer_name in self._capture_hooks:
            self.remove_hooks({self._capture_hooks.pop(layer_name)})
        cur_inputs = self._all_module_input.pop(layer_name, None)
        if not cur_inputs:
            raise RuntimeError(
                f"No calibration inputs captured for layer {layer_name}. "
                "This can happen if calibration data is missing or the "
                "forward pass did not execute for this layer."
            )
        cur_inputs, padded_masks = self._pad_captured_inputs(cur_inputs)
        self._set_attention_masks(
            ar, decoding_layer, cur_inputs, padded_masks=padded_masks
        )
        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
        elif torch.distributed.is_initialized():
            # Standard DDP (1 GPU per rank): pre-load the block onto the rank's
            # local GPU before quantize_block is called.  auto_round 0.14.2+
            # calls _move_block_to_device inside setup_ddp_if_needed_, which
            # triggers a slow reload when the block was evicted from page cache
            # during the sequential calibration pass.
            # Moving it here makes that call a no-op.
            # Use current_device_index() (local GPU index) rather than global
            # rank to avoid invalid device ordinal errors on multi-node setups.
            if hasattr(torch, "accelerator") and torch.accelerator.is_available():
                device = torch.device(
                    torch.accelerator.current_accelerator().type,
                    torch.accelerator.current_device_index(),
                )
            else:
                device = torch.device("cpu")
            decoding_layer.to(device)

        # cur_inputs remain on CPU; auto_round's block_forward handles
        # per-batch CPU->GPU transfer automatically.
        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
    self._capture_hooks.clear()

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