qbmm context manager auto check

Signed-off-by: cliu-us <cliu@us.ibm.com>

Author: chichun-charlie-liu

fms_mo/fx/dynamo_utils.py

@@ -1242,6 +1242,32 @@ def call_seq_hook(mod, *_args, **_kwargs):
                 )
                 setattr(mod_bmm_happened, f"QBmm{ln}", newQBmm)
 
+        # add auto QBmm check to last layer if any QBmms in model (only for transformers)
+        def qbmm_auto_check(_mod, *_args, **_kwargs):
+            """Automatic QBmm check. This hook will be attached to the last module and check once
+            only at the end of first forward() call. Throw a "warning" if a model has QBmm attached
+            but not called (as it could be intentional.)
+            """
+            num_called_qbmms = []
+            for lay, line_nums in qcfg["bmm_prep"]["layers_with_bmm"].items():
+                for ln in line_nums:
+                    qbmm_i = model.get_submodule(f"{lay}.QBmm{ln}")
+                    num_called_qbmms.append(qbmm_i.num_module_called == 1)
+
+            if not all(num_called_qbmms):
+                err_msg = (
+                    "QBmms were attached but not called during forward()."
+                    "Possibly patch_torch_bmm() context manager is missing."
+                )
+                if qcfg["force_stop_if_qbmm_auto_check_failed"]:
+                    raise RuntimeError(err_msg)
+                logger.warning(err_msg)
+
+            qcfg["hook_qbmm_auto_check"].remove()
+
+        last_mod = model.get_submodule(qcfg["mod_call_seq"][-1])
+        qcfg["hook_qbmm_auto_check"] = last_mod.register_forward_hook(qbmm_auto_check)
+
     # c) identify RPN/FPN
     # TODO this hack only works for torchvision models. will use find_rpn_fpn_gm()
 

fms_mo/utils/qconfig_utils.py

@@ -198,6 +198,7 @@ def qconfig_init(recipe: str = None, args: Any = None):
     qcfg["which2patch_contextmanager"] = (
         None  # an internal var that should not be set by user
     )
+    qcfg["force_stop_if_qbmm_auto_check_failed"] = False
 
     # LSTM related, if any of these is not None, then last layer (FC) will not be skipped.
     qcfg["nbits_w_lstm"] = None

tests/models/test_qmodelprep.py

@@ -300,6 +300,8 @@ def test_bert_dynamo_wi_qbmm(
             other_qmodules.append(m)
 
     # check 2: model call without our "patch" context manager, will not reach QBmm
+    #           we have an auto check in place, but it will only log warning, unless this flag
+    #           qcfg["force_stop_if_qbmm_auto_check_failed"] = True
     with torch.no_grad():
         model_bert_eager(**input_bert)
     assert all(