Feat (utils): replace weights with quantized ones

src/brevitas/nn/utils.py

@@ -1,9 +1,18 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Tuple
+
 import torch
 from torch.nn import Parameter
+from torch.utils.hooks import RemovableHandle
 
+from brevitas import config
+from brevitas.inject.enum import FloatToIntImplType
+from brevitas.inject.enum import ScalingImplType
 from brevitas.utils.torch_utils import compute_channel_view_shape
 
 
@@ -60,6 +69,107 @@ def rename_state_dict_by_postfix(old_postfix, new_postfix, state_dict):
         state_dict[keys_map[old_key]] = state_dict.pop(old_key)
 
 
+def merge_quant_weights(
+        model: torch.nn.Module,
+        example_input: torch.Tensor,
+        preserve_original_weights: bool = False) -> None:
+    """Merge quantized weights into model weights.
+
+    This could be useful for example with Learned Round.
+    After learned round training, the rounding decision for each weight element is
+    deterministic. This function uses forward hooks to discover the association
+    between weight tensors and its quantized counterparts, and update the module's weights.
+    A single forward pass is performed using ``example_input``.
+
+    Usage::
+
+        model.eval()
+        merge_quant_weights(model, sample_input)
+        # Weights are now merged and rounding mode is ROUND.
+
+    Args:
+        model: A model containing quantised layers with learned round quantisers.
+        example_input: A single example input tensor used to run the forward pass.
+        preserve_original_weights: If ``True``, the original weights are saved to
+            ``m.weight_orig`` before overwriting (default ``False``).
+    """
+    # Imported here to avoid a circular import
+    from brevitas.proxy.parameter_quant import WeightQuantProxyFromInjectorBase
+
+    hooks: List[RemovableHandle] = []
+    proxy_list: List[WeightQuantProxyFromInjectorBase] = []
+    module_tensor_id_mapping: Dict = {}
+
+    def hook(module: WeightQuantProxyFromInjectorBase, args: Tuple[Any, ...], output: Any) -> None:
+        input_tensor = args[0]
+        with torch.no_grad():
+            for m in module.tracked_module_list:
+                # We match the module based on its weights and the ID of the tensor to quantize
+                if id(m.weight.data) == id(input_tensor.data):
+                    m.weight.data = output.value.data
+                    # We track how many modules have been converted
+                    if module not in module_tensor_id_mapping:
+                        module_tensor_id_mapping[module] = 1
+                    else:
+                        module_tensor_id_mapping[module] += 1
+            proxy_list.append(module)
+
+    # Register Proxy hooks
+    for module in model.modules():
+        if not isinstance(module, WeightQuantProxyFromInjectorBase):
+            continue
+        _change_scale_impl_type(module)
+        handle = module.register_forward_hook(hook)
+        hooks.append(handle)
+
+    # Run a single forward pass to trigger the hooks
+    try:
+        model(example_input)
+    finally:
+        # Remove all hooks
+        for h in hooks:
+            h.remove()
+        hooks.clear()
+
+    # Reset quantizers from LEARNED_ROUND to ROUND
+    with torch.no_grad():
+        for module in proxy_list:
+            if module_tensor_id_mapping[module] < len(module.tracked_module_list):
+                raise RuntimeError("Not all weights associated to this quantizer were replaced")
+            _reset_quantizer(module)
+
+
+def _change_scale_impl_type(proxy) -> None:
+    """Change the scaling implementation type to PARAMETER_FROM_STATS."""
+    reinit_on_state_dict = config.REINIT_ON_STATE_DICT_LOAD
+    ignore_missing_key = config.IGNORE_MISSING_KEYS
+    config.REINIT_ON_STATE_DICT_LOAD = False
+    config.IGNORE_MISSING_KEYS = True
+    state_dict = proxy.state_dict()
+    proxy.quant_injector = proxy.quant_injector.let(
+        scaling_impl_type=ScalingImplType.PARAMETER_FROM_STATS)
+    proxy.init_tensor_quant()
+    proxy.load_state_dict(state_dict, strict=False)
+    config.IGNORE_MISSING_KEYS = ignore_missing_key
+    config.REINIT_ON_STATE_DICT_LOAD = reinit_on_state_dict
+
+
+def _reset_quantizer(proxy) -> None:
+    """Switch a weight quant proxy from LearnedRound back to standard Round."""
+    reinit_on_state_dict = config.REINIT_ON_STATE_DICT_LOAD
+    ignore_missing_key = config.IGNORE_MISSING_KEYS
+    config.REINIT_ON_STATE_DICT_LOAD = False
+    config.IGNORE_MISSING_KEYS = True
+    state_dict = {
+        k: v for k, v in proxy.state_dict().items() if not k.endswith('float_to_int_impl.value')}
+
+    proxy.quant_injector = proxy.quant_injector.let(float_to_int_impl_type=FloatToIntImplType.ROUND)
+    proxy.init_tensor_quant()
+    proxy.load_state_dict(state_dict, strict=False)
+    config.IGNORE_MISSING_KEYS = ignore_missing_key
+    config.REINIT_ON_STATE_DICT_LOAD = reinit_on_state_dict
+
+
 def check_tensors_same_ptr(tensor_list):
     pointers = []
     for t in tensor_list:

tests/brevitas/nn/test_merge_quant_weights.py

@@ -0,0 +1,108 @@
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import pytest
+import torch
+
+from brevitas.core.function_wrapper.learned_round import LearnedRoundSte
+from brevitas.core.scaling import ParameterFromStatsFromParameterScaling
+from brevitas.inject.enum import FloatToIntImplType
+from brevitas.inject.enum import LearnedRoundImplType
+from brevitas.nn import QuantLinear
+from brevitas.nn.utils import merge_quant_weights
+from brevitas.quant_tensor import QuantTensor
+from brevitas_examples.common.learned_round.learned_round_method import \
+    insert_learned_round_quantizers
+from tests.conftest import SEED
+
+IN_FEATURES = 8
+OUT_FEATURES = 16
+
+LEARNED_ROUND_OPTIONS = [
+    LearnedRoundImplType.HARD_SIGMOID, LearnedRoundImplType.SIGMOID, LearnedRoundImplType.IDENTITY]
+
+
+def _get_quant_weights(model):
+    """Get the quantised weight outputs for all QuantLinear layers in the model."""
+    results = {}
+    for name, module in model.named_modules():
+        if isinstance(module, QuantLinear):
+            quant_weight = module.quant_weight()
+            if isinstance(quant_weight, QuantTensor):
+                quant_weight = quant_weight.value
+            results[name] = quant_weight.detach().clone()
+    return results
+
+
+def _randomise_learned_round(model):
+    """Randomise learned round values to simulate training."""
+    for module in model.modules():
+        if isinstance(module, LearnedRoundSte):
+            module.value.data = torch.randn_like(module.value.data)
+
+
+@pytest.mark.parametrize("learned_round_param", LEARNED_ROUND_OPTIONS)
+def test_merge_quant_weights_preserves_quantised_weights(learned_round_param):
+    """After merging, standard round should preserve the quantised weights, remove learned
+    round and its forward hooks, and reset the rounding mode to ROUND."""
+    torch.manual_seed(SEED)
+    model = QuantLinear(in_features=IN_FEATURES, out_features=OUT_FEATURES, bias=False)
+    model.eval()
+    insert_learned_round_quantizers(model, learned_round_param)
+    assert model.weight_quant.rounding_mode == "LEARNED_ROUND"
+
+    _randomise_learned_round(model)
+    model.eval()
+
+    # Get quantised weights with learned round active
+    quant_before = _get_quant_weights(model)
+    hooks_before = len(model._forward_hooks)
+
+    # Merge learned round into weights
+    x = torch.randn(4, IN_FEATURES)
+    merge_quant_weights(model, x)
+
+    # Verify that learned round has been removed
+    for module in model.modules():
+        assert not isinstance(module, LearnedRoundSte), \
+            "LearnedRoundSte should be removed after merge"
+
+    # Verify that the merge's forward hooks were cleaned up
+    hooks_after = len(model._forward_hooks)
+    assert hooks_after == hooks_before, "Forward hooks were not cleaned up after merge"
+
+    # Verify that the rounding mode has been reset to standard round
+    assert isinstance(
+        model.weight_quant.tensor_quant.scaling_impl, ParameterFromStatsFromParameterScaling)
+    assert model.weight_quant.rounding_mode == "ROUND"
+
+    # The quantised outputs should match
+    quant_after = _get_quant_weights(model)
+    for name in quant_before:
+        assert torch.allclose(quant_before[name], quant_after[name], atol=1e-6), \
+            f"Quantised weights differ for {name} after merge"
+
+
+@pytest.mark.parametrize("learned_round_param", LEARNED_ROUND_OPTIONS)
+def test_merge_quant_weights_forward_equivalence(learned_round_param):
+    """The model forward output should be identical before and after merging."""
+    torch.manual_seed(SEED)
+    model = QuantLinear(in_features=IN_FEATURES, out_features=OUT_FEATURES, bias=True)
+    model.eval()
+
+    insert_learned_round_quantizers(model, learned_round_param)
+    _randomise_learned_round(model)
+
+    model.eval()
+    x = torch.randn(4, IN_FEATURES)
+
+    with torch.no_grad():
+        out_before = model(x).clone()
+
+    merge_quant_weights(model, x)
+
+    with torch.no_grad():
+        out_after = model(x)
+
+    assert torch.allclose(out_before, out_after, atol=1e-5), \
+        "Model outputs differ after merge"