Skip to content

module_target

module_target

Torch Module-level targets for quantization workflows.

QuantizationModuleTarget

QuantizationModuleTarget(module_path)

Base class for a quantized policy submodule target.

Initialize the target.

Parameters:

Name Type Description Default
module_path str

Dotted path to the target module, or "" for root.

required
Source code in src/versatil/quantization/module_target.py
def __init__(self, module_path: str) -> None:
    """Initialize the target.

    Args:
        module_path: Dotted path to the target module, or ``""`` for root.
    """
    self.module_path = module_path

label property

label

Return a readable module label for logs and errors.

Returns:

Type Description
str

module_path for submodule targets, or "(root)" for the

str

full-policy target.

contains_module

contains_module(module_name)

Return whether a named module is inside this target.

Parameters:

Name Type Description Default
module_name str

Fully qualified module name from named_modules().

required

Returns:

Type Description
bool

Whether module_name is the target module itself or a child of

bool

the target module. The root target contains every module.

Source code in src/versatil/quantization/module_target.py
def contains_module(self, module_name: str) -> bool:
    """Return whether a named module is inside this target.

    Args:
        module_name: Fully qualified module name from ``named_modules()``.

    Returns:
        Whether ``module_name`` is the target module itself or a child of
        the target module. The root target contains every module.
    """
    if self.module_path == "":
        return True
    return module_name == self.module_path or module_name.startswith(
        self.module_path + "."
    )

overlaps

overlaps(other)

Return whether two targets can select the same submodule.

Parameters:

Name Type Description Default
other QuantizationModuleTarget

Target to compare against this target.

required

Returns:

Type Description
bool

Whether either target is root, both targets are the same path, or

bool

one target is nested under the other.

Source code in src/versatil/quantization/module_target.py
def overlaps(self, other: "QuantizationModuleTarget") -> bool:
    """Return whether two targets can select the same submodule.

    Args:
        other: Target to compare against this target.

    Returns:
        Whether either target is root, both targets are the same path, or
        one target is nested under the other.
    """
    if self.module_path == "" or other.module_path == "":
        return True
    return (
        self.module_path == other.module_path
        or self.module_path.startswith(other.module_path + ".")
        or other.module_path.startswith(self.module_path + ".")
    )

EagerQuantizationModuleTarget

EagerQuantizationModuleTarget(module_path, quantize_config)

Bases: QuantizationModuleTarget

Target using an eager quantization config.

Initialize an eager quantization target.

Parameters:

Name Type Description Default
module_path str

Dotted path to the target module, or "" for root.

required
quantize_config AOBaseConfig

torchao eager quantization config applied to this target.

required
Source code in src/versatil/quantization/module_target.py
def __init__(
    self,
    module_path: str,
    quantize_config: AOBaseConfig,
) -> None:
    """Initialize an eager quantization target.

    Args:
        module_path: Dotted path to the target module, or ``""`` for root.
        quantize_config: torchao eager quantization config applied to this
            target.
    """
    super().__init__(module_path=module_path)
    self.quantize_config = quantize_config

PT2EQuantizationModuleTarget

PT2EQuantizationModuleTarget(module_path, pt2e_backend)

Bases: QuantizationModuleTarget

Target using a PyTorch 2 Export backend quantizer config.

Initialize a PT2E quantization target.

Parameters:

Name Type Description Default
module_path str

Dotted path to the target module, or "" for root.

required
pt2e_backend BasePT2EBackend

PT2E backend that creates the quantizer for this target.

required
Source code in src/versatil/quantization/module_target.py
def __init__(
    self,
    module_path: str,
    pt2e_backend: BasePT2EBackend,
) -> None:
    """Initialize a PT2E quantization target.

    Args:
        module_path: Dotted path to the target module, or ``""`` for root.
        pt2e_backend: PT2E backend that creates the quantizer for this
            target.
    """
    super().__init__(module_path=module_path)
    self.pt2e_backend = pt2e_backend

needs_calibration property

needs_calibration

Return whether this target requires calibration batches.

Returns:

Type Description
bool

True for static PT2E backends and False for dynamic PT2E

bool

backends.