compressor

compressor

Post-training compression orchestration.

PostTrainingCompressor

PostTrainingCompressor(checkpoint_path, modules, preparation, calibration_steps=32, checkpoint_name=value, output_directory=None, generate_report=False, pruning=None, quantization=None, deployment_backend=None)

Post-training compression pipeline for a trained policy.

Orchestrates loading, validation, preparation, pruning, quantization, deployment export, and serialization.

Initialize the compression pipeline.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	Path to the training checkpoint directory.	required
modules	list[CompressionTarget]	Per-module compression schemes (empty = global).	required
preparation	PreparationConfig	Global preparation settings.	required
calibration_steps	int	Number of calibration batches for static quantization.	32
checkpoint_name	str	Checkpoint filename inside the directory.	value
output_directory	str | None	Where to save compressed output. Defaults to checkpoint_path/compressed/.	None
generate_report	bool	Whether to generate a quantization report after saving. Disabled by default since it runs additional forward passes for benchmarking.	False
pruning	list[BasePruner] | None	Global pruning strategies (inherited by modules).	None
quantization	BaseQuantizationWorkflow | None	Quantization workflow. None exports the float model without quantization.	None
deployment_backend	DeploymentBackend | None	Deployment backend that owns artifact format and lowering. Defaults to torch inductor.	None

Source code in src/versatil/post_training_compression/compressor.py

def __init__(
    self,
    checkpoint_path: str,
    modules: list[CompressionTarget],
    preparation: PreparationConfig,
    calibration_steps: int = 32,
    checkpoint_name: str = CheckpointFilename.DEFAULT_CHECKPOINT.value,
    output_directory: str | None = None,
    generate_report: bool = False,
    pruning: list[BasePruner] | None = None,
    quantization: BaseQuantizationWorkflow | None = None,
    deployment_backend: DeploymentBackend | None = None,
) -> None:
    """Initialize the compression pipeline.

    Args:
        checkpoint_path: Path to the training checkpoint directory.
        modules: Per-module compression schemes (empty = global).
        preparation: Global preparation settings.
        calibration_steps: Number of calibration batches for
            static quantization.
        checkpoint_name: Checkpoint filename inside the directory.
        output_directory: Where to save compressed output.
            Defaults to checkpoint_path/compressed/<timestamp>.
        generate_report: Whether to generate a quantization
            report after saving. Disabled by default since
            it runs additional forward passes for benchmarking.
        pruning: Global pruning strategies (inherited by modules).
        quantization: Quantization workflow. ``None`` exports the
            float model without quantization.
        deployment_backend: Deployment backend that owns artifact format and
            lowering. Defaults to torch inductor.
    """
    self.checkpoint_path = checkpoint_path
    self.checkpoint_name = checkpoint_name
    self.output_directory = output_directory
    self.generate_report = generate_report
    self.calibration_steps = calibration_steps
    self.modules = modules
    self.preparation = preparation
    self.pruning: list[BasePruner] = pruning or []
    self.quantization = quantization
    self.deployment_backend = deployment_backend or TorchInductorBackend()

compress

compress(hydra_config)

Run the full compression pipeline.

Parameters:

Name	Type	Description	Default
hydra_config	DictConfig	Raw Hydra config for serialization into the compressed checkpoint directory.	required

Returns:

Type	Description
str	Path to the saved compressed model directory.

Source code in src/versatil/post_training_compression/compressor.py

def compress(self, hydra_config: DictConfig) -> str:
    """Run the full compression pipeline.

    Args:
        hydra_config: Raw Hydra config for serialization
            into the compressed checkpoint directory.

    Returns:
        Path to the saved compressed model directory.
    """
    modules = self.resolve_modules()
    quantization_workflow = self._resolve_quantization_workflow()
    self._validate_deployment_backend_compatibility(
        deployment_backend_name=self.deployment_backend.name,
        mode=quantization_workflow.quantization_mode,
        pt2e_backend_names=quantization_workflow.pt2e_backend_names,
    )
    context = quantization_workflow.load_policy_context(
        checkpoint_path=self.checkpoint_path,
        checkpoint_name=self.checkpoint_name,
    )
    policy = context.policy
    self.validate(policy=policy, modules=modules)
    quantization_workflow.validate_targets(model=policy)
    self._prepare_and_prune(policy=policy, modules=modules)
    exportable = ExportablePolicy.from_policy(policy)
    logging.info(f"Input keys: {exportable.observation_keys}")
    logging.info(f"Output keys: {exportable.action_keys}")
    quantized = quantization_workflow.quantize(
        context=context,
        exportable=exportable,
        calibration_steps=self.calibration_steps,
    )
    deployment_artifact = self.deployment_backend.export(
        model=quantized.quantized_model,
        example_inputs=quantized.example_inputs,
    )
    output_directory = self._resolve_output_directory()
    pt2e_backend_config = self._pt2e_backend_config(hydra_config=hydra_config)
    save_compressed_model(
        converted_model=deployment_artifact.converted_model,
        example_inputs=deployment_artifact.example_inputs,
        save_directory=output_directory,
        input_keys=policy.input_keys,
        output_keys=policy.output_keys,
        normalizer=policy.normalizer,
        training_checkpoint_path=self.checkpoint_path,
        quantization_config=hydra_config,
        quantization_workflow=quantized.quantization_workflow,
        model_filename=deployment_artifact.model_filename,
        artifact_format=deployment_artifact.artifact_format.value,
        backend_name=deployment_artifact.backend_name,
        model_bytes=deployment_artifact.model_bytes,
        denoising_thresholds=policy.get_denoising_thresholds(),
        pt2e_backend_config=pt2e_backend_config,
    )
    logging.info(f"Compressed model saved to {output_directory}")
    if self.generate_report:
        report = QuantizationReport(
            float_model=quantized.float_model,
            quantized_model=quantized.quantized_model,
            example_inputs=quantized.example_inputs,
            action_keys=policy.output_keys,
            quantization_workflow=quantized.quantization_workflow,
        )
        logging.info(f"\n{report.generate_report()}")
    return output_directory

resolve_modules

resolve_modules()

Return the compression targets for this run.

Supports two configuration modes: per-module (explicit modules list targeting specific submodules) and global (modules is empty, applying the top-level preparation and pruning to the entire policy).

Returns:

Type	Description
list[CompressionTarget]	Non-empty list of CompressionTarget instances.

Source code in src/versatil/post_training_compression/compressor.py

def resolve_modules(self) -> list[CompressionTarget]:
    """Return the compression targets for this run.

    Supports two configuration modes: per-module (explicit
    ``modules`` list targeting specific submodules) and global
    (``modules`` is empty, applying the top-level preparation and
    pruning to the entire policy).

    Returns:
        Non-empty list of CompressionTarget instances.
    """
    if self.modules:
        return self.modules
    return [
        CompressionTarget(
            module_path="",
            preparation=self.preparation,
            pruning=self.pruning,
        )
    ]

validate

validate(policy, modules)

Validate compression target module paths and overlaps.

Parameters:

Name	Type	Description	Default
policy	Module	The loaded policy model.	required
modules	list[CompressionTarget]	Resolved compression targets from resolve_modules().	required

Raises:

Type	Description
ValueError	If a module_path doesn't match a submodule, or if two targets overlap and would compound pruning on the same module.

Source code in src/versatil/post_training_compression/compressor.py

def validate(self, policy: nn.Module, modules: list[CompressionTarget]) -> None:
    """Validate compression target module paths and overlaps.

    Args:
        policy: The loaded policy model.
        modules: Resolved compression targets from resolve_modules().

    Raises:
        ValueError: If a module_path doesn't match a submodule, or if two
            targets overlap and would compound pruning on the same module.
    """
    for module in modules:
        if module.module_path == "":
            continue
        try:
            policy.get_submodule(module.module_path)
        except AttributeError as error:
            available = list(dict(policy.named_children()).keys())
            raise ValueError(
                f"Module path '{module.module_path}' not found in "
                f"policy. Available top-level modules: {available}"
            ) from error
    for index, module in enumerate(modules):
        for other in modules[index + 1 :]:
            if module.overlaps(other=other):
                raise ValueError(
                    "Compression targets overlap: "
                    f"'{module.module_path}' and '{other.module_path}'."
                )