eager

eager

Eager torchao quantization workflow.

EagerQuantizationWorkflow

EagerQuantizationWorkflow(targets, is_qat=False, auto_filter_incompatible_linears=True)

Bases: BaseQuantizationWorkflow

Eager torchao quantization workflow for PTQ and QAT.

Initialize eager torchao quantization.

Parameters:

Name	Type	Description	Default
targets	list[EagerQuantizationModuleTarget]	Module-level eager quantization targets.	required
is_qat	bool	Whether this workflow is used for QAT checkpoint training and conversion.	False
auto_filter_incompatible_linears	bool	Whether to skip linears whose in_features are incompatible with the config group size.	True

Source code in src/versatil/quantization/workflows/eager.py

def __init__(
    self,
    targets: list[EagerQuantizationModuleTarget],
    is_qat: bool = False,
    auto_filter_incompatible_linears: bool = True,
) -> None:
    """Initialize eager torchao quantization.

    Args:
        targets: Module-level eager quantization targets.
        is_qat: Whether this workflow is used for QAT checkpoint training
            and conversion.
        auto_filter_incompatible_linears: Whether to skip linears whose
            ``in_features`` are incompatible with the config group size.
    """
    if not targets:
        raise ValueError("EagerQuantizationWorkflow requires at least one target.")
    self._targets = targets
    self._is_qat = is_qat
    self.auto_filter_incompatible_linears = auto_filter_incompatible_linears
    self._prepared_targets: list[_PreparedEagerTarget] = []

targets property

targets

Return eager quantization targets.

quantization_mode property

quantization_mode

Return quantization mode name.

is_qat property

is_qat

Return whether this eager workflow handles QAT checkpoints.

load_policy_context

load_policy_context(checkpoint_path, checkpoint_name)

Load a float or QAT-prepared checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	Directory containing the training checkpoint.	required
checkpoint_name	str	Checkpoint filename to load from the directory.	required

Returns:

Type	Description
PolicyContext	Float policy context for PTQ, or QAT-prepared policy context for
PolicyContext	QAT conversion.

Source code in src/versatil/quantization/workflows/eager.py

def load_policy_context(
    self,
    checkpoint_path: str,
    checkpoint_name: str,
) -> PolicyContext:
    """Load a float or QAT-prepared checkpoint.

    Args:
        checkpoint_path: Directory containing the training checkpoint.
        checkpoint_name: Checkpoint filename to load from the directory.

    Returns:
        Float policy context for PTQ, or QAT-prepared policy context for
        QAT conversion.
    """
    if self.is_qat:
        return load_qat_policy_context(
            checkpoint_path=checkpoint_path,
            checkpoint_name=checkpoint_name,
            quantization=self,
        )
    return load_float_policy_context(
        checkpoint_path=checkpoint_path,
        checkpoint_name=checkpoint_name,
    )

quantize

quantize(context, exportable, calibration_steps)

Apply eager quantization and export the policy.

Parameters:

Name	Type	Description	Default
context	PolicyContext	Loaded policy context containing the eager policy to mutate.	required
exportable	ExportablePolicy	Export wrapper around the same policy.	required
calibration_steps	int	Unused for eager quantization.	required

Returns:

Type	Description
QuantizedContext	Exported eager-quantized model and example inputs for deployment.

Source code in src/versatil/quantization/workflows/eager.py

def quantize(
    self,
    context: PolicyContext,
    exportable: ExportablePolicy,
    calibration_steps: int,
) -> QuantizedContext:
    """Apply eager quantization and export the policy.

    Args:
        context: Loaded policy context containing the eager policy to
            mutate.
        exportable: Export wrapper around the same policy.
        calibration_steps: Unused for eager quantization.

    Returns:
        Exported eager-quantized model and example inputs for deployment.
    """
    example_inputs = build_example_inputs(
        exportable=exportable,
        observation_space=context.observation_space,
        observation_horizon=context.observation_horizon,
        tokenizer=context.tokenizer,
    )
    # Export the float baseline before quantize_() mutates the policy.
    float_exported = export_policy(
        exportable=exportable, example_inputs=example_inputs
    )
    if self.is_qat:
        self.convert_model(model=context.policy)
    else:
        self._apply_ptq(model=context.policy)
    exported = export_policy(exportable=exportable, example_inputs=example_inputs)
    return QuantizedContext(
        float_model=float_exported,
        quantized_model=exported,
        example_inputs=example_inputs,
        quantization_workflow=QuantizationWorkflow.EAGER.value,
    )

prepare_model

prepare_model(model)

Apply fake quantization modules in-place before QAT training.

Parameters:

Name	Type	Description	Default
model	Module	Policy model to prepare for QAT.	required

Raises:

Type	Description
ValueError	If the workflow is not a QAT workflow, if a target path is invalid, or if a target selects no eligible nn.Linear modules.

Source code in src/versatil/quantization/workflows/eager.py

def prepare_model(self, model: nn.Module) -> None:
    """Apply fake quantization modules in-place before QAT training.

    Args:
        model: Policy model to prepare for QAT.

    Raises:
        ValueError: If the workflow is not a QAT workflow, if a target path
            is invalid, or if a target selects no eligible ``nn.Linear``
            modules.
    """
    if not self.is_qat:
        raise ValueError("prepare_model() requires is_qat=True.")
    self.validate_targets(model=model)
    self._prepared_targets = []
    for target in self.targets:
        selected, skipped = self._select_linear_modules(
            model=model,
            target=target,
        )
        if not selected:
            skipped_text = "; ".join(
                f"{name}: {reason}" for name, reason in skipped
            )
            raise ValueError(
                f"QAT target '{target.label}' selected zero eligible "
                "nn.Linear modules. "
                f"Skipped modules: {skipped_text or 'none'}."
            )
        selected_names = set(selected)
        self._prepared_targets.append(
            _PreparedEagerTarget(
                target=target,
                module_names=selected_names,
            )
        )
        for name, reason in skipped:
            logger.info(f"Skipping QAT module {name}: {reason}")
        logger.info(
            f"Preparing {len(selected_names)} nn.Linear modules in "
            f"{target.label} for QAT with "
            f"{type(target.quantize_config).__name__}."
        )
        quantize_(
            model=model,
            config=QATConfig(base_config=target.quantize_config, step="prepare"),
            filter_fn=lambda module, fqn, names=selected_names: fqn in names,
        )

convert_model

convert_model(model)

Convert prepared fake-quant modules to quantized modules in-place.

Parameters:

Name	Type	Description	Default
model	Module	Policy model previously prepared by prepare_model().	required

Raises:

Type	Description
ValueError	If the workflow is not a QAT workflow, or if prepare_model() has not captured prepared targets.

Source code in src/versatil/quantization/workflows/eager.py

def convert_model(self, model: nn.Module) -> None:
    """Convert prepared fake-quant modules to quantized modules in-place.

    Args:
        model: Policy model previously prepared by ``prepare_model()``.

    Raises:
        ValueError: If the workflow is not a QAT workflow, or if
            ``prepare_model()`` has not captured prepared targets.
    """
    if not self.is_qat:
        raise ValueError("convert_model() requires is_qat=True.")
    if not self._prepared_targets:
        raise ValueError("QAT convert_model() requires prepare_model() first.")
    for prepared in self._prepared_targets:
        selected_names = set(prepared.module_names)
        quantize_(
            model=model,
            config=QATConfig(
                base_config=prepared.target.quantize_config,
                step="convert",
            ),
            filter_fn=lambda module, fqn, names=selected_names: fqn in names,
        )