protect fused modules with multiple hooks from dict mutate exception (#872)

Author: bfineran

src/sparseml/pytorch/sparsification/quantization/helpers.py

@@ -17,7 +17,7 @@
 """
 
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from torch.nn import BatchNorm2d, Conv2d, Embedding, Module, ReLU
@@ -33,6 +33,7 @@
 from dataclasses import dataclass, field
 
 from sparseml.pytorch.nn import ReLU as ReLU_nm
+from sparseml.pytorch.utils import get_layer
 
 
 __all__ = [
@@ -91,6 +92,25 @@
     else None
 )
 
+_FUSED_MODULE_TYPES = (
+    (
+        # Conv based layers
+        nni.ConvBn1d,
+        nni.ConvBn2d,
+        nni.ConvBn3d,
+        nni.ConvReLU1d,
+        nni.ConvReLU2d,
+        nni.ConvReLU3d,
+        nni.ConvBnReLU1d,
+        nni.ConvBnReLU2d,
+        nni.ConvBnReLU3d,
+        # Linear Layers
+        nni.LinearReLU,
+    )
+    if nni  # nni will always import if torch.quantization is available
+    else tuple()
+)
+
 
 @dataclass
 class QConfigProperties:
@@ -675,7 +695,16 @@ def fuse_module_conv_bn_relus(
     if len(current_block) > 1:
         conv_blocks.append(current_block)
     if conv_blocks:
+        # manually save and move hooks surrounding fused blocks into new fused modules
+        # due to torch.quantization error when a module has more than one hook
+        block_hooks = _delete_get_block_hooks(module, conv_blocks)
+
+        # run torch fusion
         torch_quantization.fuse_modules(module, conv_blocks, inplace=True)
+
+        # add hooks back
+        _add_fused_block_hooks(module, block_hooks)
+
     return module
 
 
@@ -701,6 +730,52 @@ def prepare_embeddings_qat(
             _prepare_qat_embedding(submodule, qconfig)
 
 
+def _delete_get_block_hooks(
+    module: Module,
+    fuse_blocks: List[str],
+) -> List[Tuple[Any, Any]]:
+    block_hooks = []
+    for block in fuse_blocks:
+        pre_hooks = []
+        post_hooks = []
+
+        # get first and last Module objects in block by their names
+        block_head = get_layer(block[0], module)
+        block_tail = get_layer(block[-1], module)
+
+        for handle_id, pre_hook_fn in list(block_head._forward_pre_hooks.items()):
+            pre_hooks.append(pre_hook_fn)
+            del block_head._forward_pre_hooks[handle_id]
+
+        for handle_id, hook_fn in list(block_tail._forward_hooks.items()):
+            post_hooks.append(hook_fn)
+            del block_tail._forward_hooks[handle_id]
+
+        block_hooks.append((pre_hooks, post_hooks))
+
+    return block_hooks
+
+
+def _add_fused_block_hooks(module: Module, block_hooks: List[Tuple[Any, Any]]):
+    fused_modules = [
+        mod for mod in module.modules() if isinstance(mod, _FUSED_MODULE_TYPES)
+    ]
+
+    if len(fused_modules) != len(block_hooks):
+        raise RuntimeError(
+            f"Number of fused modules ({len(fused_modules)}) after layer fusion in "
+            f"module {module.__class__.__name__}. does not match expected "
+            f"({len(block_hooks)}). Module may have already been fused or block "
+            "skipped during torch.quantization.fuse_modules"
+        )
+
+    for fused_module, (pre_hooks, post_hooks) in zip(fused_modules, block_hooks):
+        for pre_hook in pre_hooks:
+            fused_module.register_forward_pre_hook(pre_hook)
+        for post_hook in post_hooks:
+            fused_module.register_forward_hook(post_hook)
+
+
 def _prepare_qat_embedding(embedding: Module, qconfig: "torch.quantization.QConfig"):
     embedding.weight_fake_quant = qconfig.weight()