Add public ExportBackend.Torch that folds weight compression

changelog.d/154.public-torch-backend.added

@@ -0,0 +1 @@
+Added the public `ExportBackend.Torch` backend for fast torch inference, e.g. running evals. `finalize(backend=ExportBackend.Torch)` folds weight compression (palettization, weight quantization, pruning) into dense weights and keeps activation fake quantization. The private `ExportBackend._TORCH` name is deprecated in favor of `ExportBackend.Torch`.

src/coreai_opt/_utils/torch_utils.py

@@ -362,6 +362,54 @@ def remove_compression_parametrizations(
     return num_removed
 
 
+def fold_compression_parametrizations(
+    model: torch.nn.Module,
+    parametrization_types: type[torch.nn.Module] | tuple[type[torch.nn.Module], ...],
+) -> int:
+    """Fold parametrized tensors carrying a compression parametrization into plain tensors.
+
+    Every tensor whose parametrization list contains an instance of
+    ``parametrization_types`` is unparametrized, keeping its current
+    parametrized value, computed in eval mode. The whole parametrization list on
+    that tensor is folded, since composed parametrizations cannot be folded
+    individually.
+
+    The value is written into the ``original`` tensor in place (see
+    ``torch.nn.utils.parametrize.remove_parametrizations``), so tensor identity
+    is preserved and tied tensors stay tied.
+
+    Args:
+        model (torch.nn.Module): The model to modify in place.
+        parametrization_types (type[torch.nn.Module] | tuple[type[torch.nn.Module], ...]):
+            Parametrization types that mark a tensor for folding.
+
+    Returns:
+        int: Number of tensors folded.
+    """
+    to_fold = [
+        (module, tensor_name)
+        for module in model.modules()
+        if P.is_parametrized(module)
+        for tensor_name, parametrizations in module.parametrizations.items()
+        if any(isinstance(p, parametrization_types) for p in parametrizations)
+    ]
+
+    folded_originals: set[int] = set()
+    for module, tensor_name in to_fold:
+        parametrizations = module.parametrizations[tensor_name]
+        # Only the parametrizations are moved to eval: they are discarded by the fold, so
+        # their mode needs no restoring and the rest of the model is left untouched.
+        parametrizations.eval()
+        original = parametrizations.original
+        # A tensor tied to one folded earlier shares its ``original``, which already
+        # holds the folded value; re-applying the parametrizations would compress it twice.
+        leave_parametrized = id(original) not in folded_originals
+        folded_originals.add(id(original))
+        P.remove_parametrizations(module, tensor_name, leave_parametrized=leave_parametrized)
+
+    return len(to_fold)
+
+
 def find_parametrization_matching_cls(
     module: torch.nn.Module,
     param_name: str,

src/coreai_opt/base_model_compressor.py

@@ -79,6 +79,21 @@ def _mark_model_as_prepared(model: torch.nn.Module) -> None:
         """
         model.register_buffer(_COREAI_OPT_PREPARED_ATTR, torch.tensor(True), persistent=False)
 
+    @staticmethod
+    def _unmark_model_as_prepared(model: torch.nn.Module) -> None:
+        """
+        Clear the prepared marker from a model.
+
+        Called at the end of ``finalize`` so the finalized model can be further
+        processed (e.g., by Quantizer for activation quantization after weight
+        palettization).
+
+        Args:
+            model: The model to unmark.
+        """
+        if hasattr(model, _COREAI_OPT_PREPARED_ATTR):
+            delattr(model, _COREAI_OPT_PREPARED_ATTR)
+
     @classmethod
     def supported_modules(cls) -> tuple[type[torch.nn.Module]]:
         """
@@ -129,10 +144,12 @@ def finalize(
         For example, for quantization:
         - For CoreAI backend (default), custom ops will be inserted.
         - For CoreML, compression metadata will be inserted as buffers in the model.
+        - For Torch, weight compression is folded into dense weights for fast torch
+          inference.
 
         Args:
             model: Optional model to finalize. If None, uses self._model
-            backend: Target export backend (CoreAI, CoreML)
+            backend: Target backend (CoreAI, CoreML, Torch)
             mmap_dir (str | None): If provided, serialize finalized compressed
                 weights to safetensors files under this directory and re-load
                 them via mmap, so large models can be finalized without holding

src/coreai_opt/common.py

@@ -140,19 +140,29 @@ class ExportBackend(_StrEnum, metaclass=_DeprecatedMemberEnumMeta):
     Each member is a string value representing the backend format.
 
     Attributes:
+        Torch: Plain PyTorch model for fast inference (e.g. running evals). Weight
+            compression is folded into dense weights; activation fake quantization
+            is kept.
         CoreML: Core ML format with compression metadata buffers.
         CoreAI: Core AI format with custom ops.
 
     """
 
-    _TORCH = auto()
+    Torch = auto()
     CoreML = auto()
     CoreAI = auto()
 
-    __deprecated_aliases__: ClassVar[dict[str, str]] = {"MIL": "CoreML", "MLIR": "CoreAI"}
+    __deprecated_aliases__: ClassVar[dict[str, str]] = {
+        "_TORCH": "Torch",
+        "MIL": "CoreML",
+        "MLIR": "CoreAI",
+    }
 
     if TYPE_CHECKING:
         # Surface the deprecated aliases above for static type checkers.
+        _TORCH: ExportBackend
+        """Deprecated. Use `ExportBackend.Torch` instead."""
+
         MIL: ExportBackend
         """Deprecated. Use `ExportBackend.CoreML` instead."""
 

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -29,6 +29,7 @@
 )
 from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
 from coreai_opt._utils.torch_utils import (
+    fold_compression_parametrizations as _fold_compression_parametrizations,
     move_model_to_train as _move_model_to_train,
     remove_compression_parametrizations as _remove_compression_parametrizations,
 )
@@ -45,6 +46,7 @@
 from coreai_opt.palettization.spec.fake_palettize import (
     _disable_fake_palett,
     _enable_fake_palett,
+    _FakePalettizeImplBase,
 )
 
 from ._prepare_for_export import (
@@ -477,15 +479,15 @@ def finalize(
     ) -> torch.nn.Module:
         """Convert palettized model to backend-specific representations.
 
-        Only call ``finalize`` when exporting to a target backend. For torch-based
-        evaluation, use the model returned by ``prepare()`` directly rather than
-        calling ``finalize``.
+        Use ``backend=ExportBackend.Torch`` for fast torch inference (e.g. running
+        evals): palettized weights are folded into dense weights, so the model no
+        longer reconstructs them from the LUT on every forward pass.
 
         Args:
             model (nn.Module | None): Model to finalize. If None, uses the
                 internal prepared model.
-            backend (ExportBackend): Target export backend for the palettized
-                model. Supports CoreAI (default) and CoreML backends.
+            backend (ExportBackend): Target backend for the palettized model.
+                Supports CoreAI (default), CoreML and Torch backends.
             mmap_dir (str | None): If provided, finalized palettized weights are
                 written under this directory and re-loaded as mmap-backed
                 tensors so they don't have to be held in RAM. Only supported
@@ -518,8 +520,8 @@ def finalize(
 
         # Backend-specific processing
         match backend:
-            case ExportBackend._TORCH:
-                pass
+            case ExportBackend.Torch:
+                _fold_compression_parametrizations(finalized_model, _FakePalettizeImplBase)
 
             case ExportBackend.CoreAI:
                 finalized_model = _prepare_for_mlir_export(finalized_model, mmap_dir=mmap_dir)
@@ -531,10 +533,7 @@ def finalize(
                 msg = f"Unsupported backend: {backend}"
                 raise NotImplementedError(msg)
 
-        # Clear the prepared flag so the finalized model can be further processed
-        # (e.g., by Quantizer for activation quantization after weight palettization)
-        if hasattr(finalized_model, "_coreai_opt_prepared"):
-            delattr(finalized_model, "_coreai_opt_prepared")
+        self._unmark_model_as_prepared(finalized_model)
 
         return finalized_model
 

src/coreai_opt/pruning/magnitude_pruner.py

@@ -20,6 +20,7 @@
 from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
 from coreai_opt._utils.torch_utils import (
     find_parametrization_matching_cls as _find_parametrization_matching_cls,
+    fold_compression_parametrizations as _fold_compression_parametrizations,
     move_model_to_eval as _move_model_to_eval,
 )
 from coreai_opt.common import ExportBackend
@@ -142,10 +143,14 @@ def finalize(
     ) -> torch.nn.Module:
         """Finalize the model to be lowered to the target backend.
 
+        Use ``backend=ExportBackend.Torch`` for fast torch inference (e.g. running
+        evals): pruned weights are folded into dense weights.
+
         Args:
             model (torch.nn.Module | None): Model to finalize. Uses the model
                 passed at construction time when ``None``.
-            backend (ExportBackend): Target export backend.
+            backend (ExportBackend): Target backend. Supports CoreAI (default),
+                CoreML and Torch backends.
 
         Returns:
             torch.nn.Module: The finalized model ready for the target backend.
@@ -161,8 +166,8 @@ def finalize(
             raise RuntimeError("Model must be prepared before finalization. Call prepare() first.")
 
         match backend:
-            case ExportBackend._TORCH:
-                pass
+            case ExportBackend.Torch:
+                _fold_compression_parametrizations(model, PruneImplBase)
 
             case ExportBackend.CoreAI:
                 model = _prepare_for_mlir_export(model)
@@ -173,6 +178,7 @@ def finalize(
             case _:
                 raise ValueError(f"Unsupported backend: {backend}")
 
+        self._unmark_model_as_prepared(model)
         if update_internal:
             self._model = model
         return model

src/coreai_opt/quantization/_eager/quantizer.py

@@ -27,7 +27,11 @@
     TorchFunctionEagerHandler,
 )
 from coreai_opt._utils.spec_utils import PartialConstructor
-from coreai_opt._utils.torch_utils import move_model_to_eval, move_model_to_train
+from coreai_opt._utils.torch_utils import (
+    fold_compression_parametrizations,
+    move_model_to_eval,
+    move_model_to_train,
+)
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
@@ -237,13 +241,13 @@ def finalize(
         """Convert quantized model to backend-specific representations.
 
         Converts fake quantization modules into backend-specific quantization ops.
-        Only call ``finalize`` when exporting to a target backend. For torch-based evaluation, use
-        the model returned by ``prepare()`` directly rather than calling ``finalize``.
+        For ``backend=ExportBackend.Torch``, weight fake quantization is folded into
+        dense weights and activation fake quantization is kept.
 
         Args:
             model: Model to finalize (uses internal model if None).
-            backend: Target export backend for the quantized model. Supports
-                CoreAI (default) and CoreML.
+            backend: Target backend for the quantized model. Supports
+                CoreAI (default), CoreML and Torch.
             mmap_dir (str | None): If provided, serialize finalized quantized
                 weights to safetensors files under this directory and re-load
                 them via mmap. Only supported with the CoreAI backend. The
@@ -272,7 +276,10 @@ def finalize(
 
         # Backend specific logic
         match backend:
-            case ExportBackend._TORCH:
+            case ExportBackend.Torch:
+                # Activation fake quantization lives in standalone modules, not
+                # parametrizations, so only weight fake quantization is folded.
+                fold_compression_parametrizations(model, FakeQuantizeImplBase)
                 finalized_model = model
             case ExportBackend.CoreAI:
                 finalized_model = prepare_for_mlir_export(model, mmap_dir=mmap_dir)
@@ -282,6 +289,7 @@ def finalize(
                 msg = f"Unsupported backend {backend}"
                 raise ValueError(msg)
 
+        self._unmark_model_as_prepared(finalized_model)
         return finalized_model
 
     @contextmanager

src/coreai_opt/quantization/_export_utils.py

@@ -418,7 +418,7 @@ def validate_activation_export_supported(fake_quant_mod: FakeQuantizeImplBase) -
             "Core AI export does not support PerBlockGranularity on activations: a "
             "per-block scale cannot be canonicalized to 0-D or 1-D. Use "
             "PerTensorGranularity or PerChannelGranularity, export with "
-            "backend=ExportBackend._TORCH, or install and activate an extension "
+            "backend=ExportBackend.Torch, or install and activate an extension "
             "package that registers an activation export handler for this granularity."
         )
 

src/coreai_opt/quantization/_graph/_prepare_for_export.py

@@ -6,7 +6,7 @@
 """Export preparation utilities for quantized models.
 
 This module provides functions to prepare PT2E-quantized models for different
-export backends (CoreML, CoreAI) by converting fake quantization operations to
+backends (CoreML, CoreAI, Torch) by converting fake quantization operations to
 backend-specific representations.
 """
 
@@ -653,6 +653,52 @@ def prepare_for_mil_export(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
     return model
 
 
+def prepare_for_torch_inference(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    """Fold weight fake quantization into dense weights for fast torch inference.
+
+    Each weight read through a fake quantization node is replaced by its
+    fake-quantized value, and the node and its module are removed. Activation
+    fake quantization is kept.
+
+    Args:
+        model (torch.fx.GraphModule): The quantized GraphModule. Modified in place.
+
+    Returns:
+        torch.fx.GraphModule: The same GraphModule with weight fake quantization folded.
+
+    Raises:
+        TypeError: If model is not a torch.fx.GraphModule.
+    """
+    if not isinstance(model, torch.fx.GraphModule):
+        raise TypeError("Model must be a torch.fx.GraphModule")
+
+    for node, fake_quant_mod in _collect_fake_quantize_nodes(model):
+        if not _is_weight_fake_quant(node, fake_quant_mod):
+            continue
+
+        weight_node = _get_fake_quant_input(node)
+        if len(weight_node.users) > 1:
+            logger.warning(
+                "Weight %r is also read directly by the graph; "
+                "keeping its fake quantization instead of folding it.",
+                weight_node.target,
+            )
+            continue
+
+        dense_weight = resolve_attr(model, str(weight_node.target))
+        with torch.no_grad():
+            dense_weight.data = fake_quant_mod(dense_weight.data)
+
+        node.replace_all_uses_with(weight_node)
+        model.graph.erase_node(node)
+        remove_fake_quant_module(model, node)
+
+    model.graph.eliminate_dead_code()
+    model.recompile()
+
+    return model
+
+
 def _release_unused_dense_weights(
     model: torch.fx.GraphModule,
     dense_weight_targets: Iterable[str],