serialization

serialization

Save and load compressed models with metadata.

save_compressed_model

save_compressed_model(converted_model, example_inputs, save_directory, input_keys, output_keys, normalizer, training_checkpoint_path, quantization_config, quantization_workflow, model_filename=value, normalizer_filename=value, artifact_format=value, backend_name=value, model_bytes=None, denoising_thresholds=None, pt2e_backend_config=None)

Save compressed model artifact with normalizer and metadata.

Saves the deployment artifact, normalizer, quantization config, training config, optional tokenizer files, and compression metadata.

Parameters:

Name	Type	Description	Default
converted_model	Module | None	The converted model, used for Torch Export artifacts.	required
example_inputs	tuple[Tensor, ...]	Example input tensors for torch.export.	required
save_directory	str	Directory to save into (created if needed).	required
input_keys	list[str]	Sorted input (observation) key ordering.	required
output_keys	list[str]	Sorted output (action) key ordering.	required
normalizer	LinearNormalizer	The policy's normalizer module.	required
training_checkpoint_path	str	Path to the original training checkpoint directory used as the source for compression.	required
quantization_config	Any	The QuantizationConfig used for quantization.	required
quantization_workflow	str	The workflow used (QuantizationWorkflow value).	required
model_filename	str	Filename for the saved deployment artifact.	value
normalizer_filename	str	Filename for the saved normalizer state.	value
artifact_format	str	Serialized artifact format identifier.	value
backend_name	str	Serialized deployment backend identifier.	value
model_bytes	bytes | None	Optional pre-lowered artifact bytes, used for .pte.	None
denoising_thresholds	dict[str, float] | None	Per-action-key denoising thresholds from the source policy, persisted so compressed deployments zero small deltas exactly like the float runtime.	None
pt2e_backend_config	dict[str, Any] | None	Instantiable config node of the PT2E quantizer backend, persisted so inference can rebuild the backend without depending on the full compressor config schema.	None

Returns:

Type	Description
Path	Path to the save directory.

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

def save_compressed_model(
    converted_model: nn.Module | None,
    example_inputs: tuple[torch.Tensor, ...],
    save_directory: str,
    input_keys: list[str],
    output_keys: list[str],
    normalizer: LinearNormalizer,
    training_checkpoint_path: str,
    quantization_config: Any,
    quantization_workflow: str,
    model_filename: str = CompressionFilename.COMPRESSED_MODEL.value,
    normalizer_filename: str = CompressionFilename.NORMALIZER.value,
    artifact_format: str = ArtifactFormat.TORCH_EXPORT_PT2.value,
    backend_name: str = DeploymentBackendName.TORCH_INDUCTOR.value,
    model_bytes: bytes | None = None,
    denoising_thresholds: dict[str, float] | None = None,
    pt2e_backend_config: dict[str, Any] | None = None,
) -> Path:
    """Save compressed model artifact with normalizer and metadata.

    Saves the deployment artifact, normalizer, quantization config, training
    config, optional tokenizer files, and compression metadata.

    Args:
        converted_model: The converted model, used for Torch Export artifacts.
        example_inputs: Example input tensors for torch.export.
        save_directory: Directory to save into (created if needed).
        input_keys: Sorted input (observation) key ordering.
        output_keys: Sorted output (action) key ordering.
        normalizer: The policy's normalizer module.
        training_checkpoint_path: Path to the original training checkpoint
            directory used as the source for compression.
        quantization_config: The QuantizationConfig used for quantization.
        quantization_workflow: The workflow used (QuantizationWorkflow value).
        model_filename: Filename for the saved deployment artifact.
        normalizer_filename: Filename for the saved normalizer state.
        artifact_format: Serialized artifact format identifier.
        backend_name: Serialized deployment backend identifier.
        model_bytes: Optional pre-lowered artifact bytes, used for .pte.
        denoising_thresholds: Per-action-key denoising thresholds from the
            source policy, persisted so compressed deployments zero small
            deltas exactly like the float runtime.
        pt2e_backend_config: Instantiable config node of the PT2E quantizer
            backend, persisted so inference can rebuild the backend without
            depending on the full compressor config schema.

    Returns:
        Path to the save directory.
    """
    save_path = Path(save_directory)
    save_path.mkdir(parents=True, exist_ok=True)
    model_path = save_path / model_filename
    if model_bytes is not None:
        model_path.write_bytes(model_bytes)
    else:
        if converted_model is None:
            raise ValueError(
                "converted_model is required when model_bytes is not provided."
            )
        exported_program = _export_with_dynamic_batch(
            model=converted_model,
            example_inputs=example_inputs,
        )
        torch.export.save(exported_program, str(model_path))
    torch.save(normalizer.state_dict(), save_path / normalizer_filename)
    config_omega = OmegaConf.structured(quantization_config)
    OmegaConf.save(
        config=config_omega,
        f=save_path / CompressionFilename.QUANTIZATION_CONFIG.value,
    )
    tokenizer_source = (
        Path(training_checkpoint_path) / CheckpointFilename.TOKENIZER_DIR.value
    )
    tokenizer_dest = save_path / CompressionFilename.TOKENIZER_DIR.value
    if tokenizer_source.exists():
        if tokenizer_dest.exists():
            shutil.rmtree(tokenizer_dest)
        shutil.copytree(tokenizer_source, tokenizer_dest)
    config_source = Path(training_checkpoint_path) / CheckpointFilename.CONFIG.value
    if config_source.exists():
        shutil.copy2(config_source, save_path / CheckpointFilename.CONFIG.value)
    metadata = {
        CompressionMetadataKey.MODEL_FILE.value: model_filename,
        CompressionMetadataKey.NORMALIZER_FILE.value: normalizer_filename,
        CompressionMetadataKey.ARTIFACT_FORMAT.value: artifact_format,
        CompressionMetadataKey.DEPLOYMENT_BACKEND.value: backend_name,
        CompressionMetadataKey.INPUT_KEYS.value: input_keys,
        CompressionMetadataKey.OUTPUT_KEYS.value: output_keys,
        CompressionMetadataKey.TORCHAO_VERSION.value: _get_torchao_version(),
        CompressionMetadataKey.TORCH_VERSION.value: torch.__version__,
        CompressionMetadataKey.TRAINING_CHECKPOINT_PATH.value: training_checkpoint_path,
        CompressionMetadataKey.QUANTIZATION_WORKFLOW.value: quantization_workflow,
        CompressionMetadataKey.DENOISING_THRESHOLDS.value: denoising_thresholds or {},
        CompressionMetadataKey.PT2E_BACKEND.value: pt2e_backend_config,
    }
    with open(save_path / CompressionFilename.COMPRESSION_METADATA.value, "w") as file:
        json.dump(metadata, file, indent=2)

    return save_path

load_compression_metadata

load_compression_metadata(metadata_path)

Load compression metadata from a checkpoint directory.

Parameters:

Name	Type	Description	Default
metadata_path	str	Path to compression_metadata.json.	required

Returns:

Type	Description
dict[str, Any]	Dict with runtime metadata.

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

def load_compression_metadata(metadata_path: str) -> dict[str, Any]:
    """Load compression metadata from a checkpoint directory.

    Args:
        metadata_path: Path to compression_metadata.json.

    Returns:
        Dict with runtime metadata.
    """
    with open(metadata_path) as file:
        metadata: dict[str, Any] = json.load(file)
    return metadata