base

base

Shared quantization workflow types.

BaseQuantizationWorkflow

Bases: ABC

Base class for ordered quantization workflows.

quantization_mode abstractmethod property

quantization_mode

Return the quantization path name: none, pt2e or eager.

is_qat property

is_qat

Return whether the workflow expects QAT-trained weights.

targets property

targets

Return module-level quantization targets.

pt2e_backend_names property

pt2e_backend_names

Return serialized PT2E backend names used by this workflow.

load_policy_context abstractmethod

load_policy_context(checkpoint_path, checkpoint_name)

Load the policy context required by this workflow.

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	Policy context with the loaded policy, metadata, normalizer, and
PolicyContext	optional tokenizer.

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

@abstractmethod
def load_policy_context(
    self,
    checkpoint_path: str,
    checkpoint_name: str,
) -> PolicyContext:
    """Load the policy context required by this workflow.

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

    Returns:
        Policy context with the loaded policy, metadata, normalizer, and
        optional tokenizer.
    """

quantize abstractmethod

quantize(context, exportable, calibration_steps)

Run the workflow and return exported deployment inputs.

Parameters:

Name	Type	Description	Default
context	PolicyContext	Loaded policy context used by the workflow.	required
exportable	ExportablePolicy	Policy wrapper exposing positional tensor inputs for torch.export.	required
calibration_steps	int	Maximum calibration batches for workflows that require calibration.	required

Returns:

Type	Description
QuantizedContext	Quantized context containing the float export, selected deployment
QuantizedContext	model, example inputs, and serialized workflow name.

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

@abstractmethod
def quantize(
    self,
    context: PolicyContext,
    exportable: ExportablePolicy,
    calibration_steps: int,
) -> "QuantizedContext":
    """Run the workflow and return exported deployment inputs.

    Args:
        context: Loaded policy context used by the workflow.
        exportable: Policy wrapper exposing positional tensor inputs for
            ``torch.export``.
        calibration_steps: Maximum calibration batches for workflows that
            require calibration.

    Returns:
        Quantized context containing the float export, selected deployment
        model, example inputs, and serialized workflow name.
    """

prepare_model

prepare_model(model)

Prepare fake-quant modules before loading or training QAT weights.

Parameters:

Name	Type	Description	Default
model	Module	Eager policy model to mutate in-place.	required

Raises:

Type	Description
NotImplementedError	If the workflow does not implement QAT preparation.

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

def prepare_model(self, model: nn.Module) -> None:
    """Prepare fake-quant modules before loading or training QAT weights.

    Args:
        model: Eager policy model to mutate in-place.

    Raises:
        NotImplementedError: If the workflow does not implement QAT
            preparation.
    """
    raise NotImplementedError(
        f"{type(self).__name__} does not support QAT preparation."
    )

validate_targets

validate_targets(model)

Validate this workflow's target module paths and overlaps.

Parameters:

Name	Type	Description	Default
model	Module	Policy model whose submodule tree should contain every target.	required

Raises:

Type	Description
ValueError	If a target path does not exist or two targets overlap.

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

def validate_targets(self, model: nn.Module) -> None:
    """Validate this workflow's target module paths and overlaps.

    Args:
        model: Policy model whose submodule tree should contain every
            target.

    Raises:
        ValueError: If a target path does not exist or two targets overlap.
    """
    validate_quantization_targets(model=model, targets=self.targets)

QuantizedContext dataclass

QuantizedContext(float_model, quantized_model, example_inputs, quantization_workflow)

Quantized export result plus inputs and metadata needed for deployment.

validate_quantization_targets

validate_quantization_targets(model, targets)

Validate quantization target paths and overlap.

Parameters:

Name	Type	Description	Default
model	Module	Policy model whose submodule tree should contain every target.	required
targets	list[QuantizationModuleTarget]	Module targets configured for one quantization workflow.	required

Raises:

Type	Description
ValueError	If a target path does not exist or two targets overlap.

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

def validate_quantization_targets(
    model: nn.Module,
    targets: list[QuantizationModuleTarget],
) -> None:
    """Validate quantization target paths and overlap.

    Args:
        model: Policy model whose submodule tree should contain every target.
        targets: Module targets configured for one quantization workflow.

    Raises:
        ValueError: If a target path does not exist or two targets overlap.
    """
    for target in targets:
        if target.module_path == "":
            continue
        try:
            model.get_submodule(target.module_path)
        except AttributeError as error:
            available = list(dict(model.named_children()).keys())
            raise ValueError(
                f"Quantization target '{target.module_path}' not found in model. "
                f"Available top-level modules: {available}."
            ) from error
    for index, target in enumerate(targets):
        for other in targets[index + 1 :]:
            if target.overlaps(other=other):
                raise ValueError(
                    "Quantization targets overlap: "
                    f"'{target.module_path}' and '{other.module_path}'."
                )