pt2e

pt2e

PT2E quantization workflow.

PT2EQuantizationWorkflow

PT2EQuantizationWorkflow(targets)

Bases: BaseQuantizationWorkflow

PT2E graph quantization workflow.

Initialize with PT2E module targets.

Parameters:

Name	Type	Description	Default
targets	list[PT2EQuantizationModuleTarget]	module-level PT2E quantization targets.	required

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

def __init__(self, targets: list[PT2EQuantizationModuleTarget]) -> None:
    """Initialize with PT2E module targets.

    Args:
        targets: module-level PT2E quantization targets.
    """
    if not targets:
        raise ValueError("PT2EQuantizationWorkflow requires at least one target.")
    self._targets = targets
    if self.is_qat:
        raise NotImplementedError("PT2E QAT configuration is not supported yet.")

targets property

targets

Return PT2E quantization targets.

pt2e_backend property

pt2e_backend

Return the first target backend for runtime environment setup.

pt2e_backend_names property

pt2e_backend_names

Return serialized PT2E backend names used by all targets.

quantization_mode property

quantization_mode

Return pt2e because this workflow quantizes an exported graph.

needs_calibration property

needs_calibration

Static PT2E requires calibration, dynamic does not.

is_qat property

is_qat

Return whether the PT2E backend is configured for QAT.

load_policy_context

load_policy_context(checkpoint_path, checkpoint_name)

Load the policy checkpoint required by PT2E quantization.

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 used for PT2E export, preparation,
PolicyContext	calibration, and conversion.

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

def load_policy_context(
    self,
    checkpoint_path: str,
    checkpoint_name: str,
) -> PolicyContext:
    """Load the policy checkpoint required by PT2E quantization.

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

    Returns:
        Float policy context used for PT2E export, preparation,
        calibration, and conversion.
    """
    return load_float_policy_context(
        checkpoint_path=checkpoint_path,
        checkpoint_name=checkpoint_name,
    )

quantize

quantize(context, exportable, calibration_steps)

Export, prepare, optionally calibrate, and convert with PT2E.

Parameters:

Name	Type	Description	Default
context	PolicyContext	Loaded float policy context.	required
exportable	ExportablePolicy	Policy wrapper exposing positional tensor inputs for torch.export.	required
calibration_steps	int	Maximum number of training batches used for static PT2E calibration.	required

Returns:

Type	Description
QuantizedContext	Float exported model, PT2E-converted model, and example inputs.

Raises:

Type	Description
ValueError	If a target path is invalid or targets overlap.

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

def quantize(
    self,
    context: PolicyContext,
    exportable: ExportablePolicy,
    calibration_steps: int,
) -> QuantizedContext:
    """Export, prepare, optionally calibrate, and convert with PT2E.

    Args:
        context: Loaded float policy context.
        exportable: Policy wrapper exposing positional tensor inputs for
            ``torch.export``.
        calibration_steps: Maximum number of training batches used for
            static PT2E calibration.

    Returns:
        Float exported model, PT2E-converted model, and example inputs.

    Raises:
        ValueError: If a target path is invalid or targets overlap.
    """
    self.validate_targets(model=context.policy)
    calibration = self._build_calibration(
        context=context,
        exportable=exportable,
        targets=self.targets,
        calibration_steps=calibration_steps,
    )
    # Export always uses synthetic batch>=2 example inputs: a raw
    # training batch with batch_size=1 would trip torch.export's 0/1
    # specialization on the dynamic batch dimension. Calibration still
    # runs on real dataloader batches.
    example_inputs = build_example_inputs(
        exportable=exportable,
        observation_space=context.observation_space,
        observation_horizon=context.observation_horizon,
        tokenizer=context.tokenizer,
    )
    exported = export_policy(exportable=exportable, example_inputs=example_inputs)
    # prepare_pt2e/convert_pt2e mutate the exported graph in place; keep a
    # genuinely float copy so the report compares against the pre-
    # quantization model instead of the mutated graph itself.
    float_exported = copy.deepcopy(exported)
    converted = self._convert_exported_model(
        exported=exported,
        targets=self.targets,
        calibration=calibration,
    )
    return QuantizedContext(
        float_model=float_exported,
        quantized_model=converted,
        example_inputs=example_inputs,
        quantization_workflow=QuantizationWorkflow.PT2E.value,
    )