torch_patches

torch_patches

In-memory compatibility patches for torchao version-specific issues.

The patches rewrite torchao module source as it is imported, through a meta-path finder registered by register_torchao_patches. Nothing is written to site-packages: other processes and projects sharing the environment see pristine torchao, read-only installs work, and uninstalling versatil leaves no trace.

SourcePatch dataclass

SourcePatch(original, replacement, required)

One source-fragment replacement applied at module load time.

Attributes:

Name	Type	Description
original	str	Source fragment the installed torchao version contains.
replacement	str	Fragment compiled in its place.
required	bool	Whether the module must contain either fragment. Required patches raise on unknown torchao source because their absence means silently wrong behavior; optional patches disable crashing statements, so their absence means upstream already removed them.

PatchingSourceLoader

PatchingSourceLoader(fullname, path, patches)

Bases: SourceFileLoader

Source loader applying string replacements before compilation.

Initialize the loader.

Parameters:

Name	Type	Description	Default
fullname	str	Fully qualified module name.	required
path	str	Filesystem path of the module source.	required
patches	list[SourcePatch]	Replacements to apply to the module source.	required

Source code in src/versatil/quantization/torch_patches.py

def __init__(self, fullname: str, path: str, patches: list[SourcePatch]) -> None:
    """Initialize the loader.

    Args:
        fullname: Fully qualified module name.
        path: Filesystem path of the module source.
        patches: Replacements to apply to the module source.
    """
    super().__init__(fullname, path)
    self._patches = patches

get_code

get_code(fullname)

Compile the patched source, bypassing the bytecode cache.

Reading the cache could load unpatched bytecode; writing it would leak patched bytecode to imports that bypass the hook.

Parameters:

Name	Type	Description	Default
fullname	str	Fully qualified module name.	required

Source code in src/versatil/quantization/torch_patches.py

def get_code(self, fullname: str) -> CodeType:
    """Compile the patched source, bypassing the bytecode cache.

    Reading the cache could load unpatched bytecode; writing it would
    leak patched bytecode to imports that bypass the hook.

    Args:
        fullname: Fully qualified module name.
    """
    data = self.get_data(self.path)
    return self.source_to_code(data, self.path)

source_to_code

source_to_code(data, path)

Apply the replacements and compile.

Parameters:

Name	Type	Description	Default
data	bytes	Raw module source bytes.	required
path	str	Filesystem path of the module source.	required

Raises:

Type	Description
RuntimeError	If the source contains neither the original nor the patched fragment, meaning the installed torchao version is not the one these patches target.

Source code in src/versatil/quantization/torch_patches.py

def source_to_code(self, data: bytes, path: str) -> CodeType:
    """Apply the replacements and compile.

    Args:
        data: Raw module source bytes.
        path: Filesystem path of the module source.

    Raises:
        RuntimeError: If the source contains neither the original nor the
            patched fragment, meaning the installed torchao version is not
            the one these patches target.
    """
    source = importlib.util.decode_source(data)
    for patch in self._patches:
        if patch.original in source:
            source = source.replace(patch.original, patch.replacement)
            continue
        if patch.required and patch.replacement not in source:
            raise RuntimeError(
                f"torchao module '{self.name}' does not match the version "
                "the versatil compatibility patches target. Install the "
                "torchao version pinned in pyproject.toml or update "
                "versatil.quantization.torch_patches."
            )
    code: CodeType = super().source_to_code(source.encode(), path)
    return code

TorchaoPatchFinder

TorchaoPatchFinder(module_patches)

Bases: MetaPathFinder

Meta-path finder routing patched torchao modules through the loader.

Initialize the finder.

Parameters:

Name	Type	Description	Default
module_patches	dict[str, list[SourcePatch]]	Mapping from module name to its source patches.	required

Source code in src/versatil/quantization/torch_patches.py

def __init__(self, module_patches: dict[str, list[SourcePatch]]) -> None:
    """Initialize the finder.

    Args:
        module_patches: Mapping from module name to its source patches.
    """
    self._module_patches = module_patches

find_spec

find_spec(fullname, path=None, target=None)

Return a patched spec for targeted modules.

Parameters:

Name	Type	Description	Default
fullname	str	Fully qualified module name being imported.	required
path	list[str] | None	Parent package search path from the import system.	None
target	object	Existing module for reloads, unused.	None

Source code in src/versatil/quantization/torch_patches.py

def find_spec(
    self,
    fullname: str,
    path: list[str] | None = None,
    target: object = None,
) -> importlib.machinery.ModuleSpec | None:
    """Return a patched spec for targeted modules.

    Args:
        fullname: Fully qualified module name being imported.
        path: Parent package search path from the import system.
        target: Existing module for reloads, unused.
    """
    patches = self._module_patches.get(fullname)
    if patches is None:
        return None
    spec = importlib.machinery.PathFinder.find_spec(fullname, path)
    if spec is None or spec.origin is None:
        return None
    loader = PatchingSourceLoader(
        fullname=fullname, path=spec.origin, patches=patches
    )
    return importlib.util.spec_from_file_location(
        fullname,
        spec.origin,
        loader=loader,
        submodule_search_locations=spec.submodule_search_locations,
    )

register_torchao_patches

register_torchao_patches()

Install the import hook applying the torchao compatibility patches.

Idempotent. Must run before torchao.quantization is imported; targeted modules that are already imported cannot be patched in place and are reported with a warning.

Source code in src/versatil/quantization/torch_patches.py

def register_torchao_patches() -> None:
    """Install the import hook applying the torchao compatibility patches.

    Idempotent. Must run before ``torchao.quantization`` is imported;
    targeted modules that are already imported cannot be patched in place
    and are reported with a warning.
    """
    if any(isinstance(finder, TorchaoPatchFinder) for finder in sys.meta_path):
        return
    module_patches = dict(_QAT_MODULE_PATCHES)
    if sys.version_info >= (3, 14):
        module_patches.update(_PT2E_MODULE_PATCHES)
    already_imported = [name for name in module_patches if name in sys.modules]
    if already_imported:
        logging.warning(
            f"torchao modules {already_imported} were imported before versatil "
            "registered its compatibility patches; the patches do not apply to "
            "this process."
        )
    sys.meta_path.insert(0, TorchaoPatchFinder(module_patches=module_patches))