Refactored Model Patcher Class

plugins/accelerated-peft/src/fms_acceleration_peft/autogptq_utils.py

@@ -52,7 +52,7 @@ def build_patch_to_view_tensor_to_parameter_for_fsdp_gptq(
         torch_dtype=torch.int32,  # patch it back to
     )
 
-def load_fsdp_gptq_patch(target_module, torch_dtype):
+def register_tensors_as_parameters_patch_rule(target_module, torch_dtype):
     # Register patch
     ModelPatcher.register(
         ModelPatcherRule(

plugins/accelerated-peft/src/fms_acceleration_peft/framework_plugin_autogptq.py

@@ -195,8 +195,11 @@ def model_loader(self, model_name: str, **kwargs):
             and os.environ.get("ACCELERATE_USE_FSDP", "false").lower() == "true"
         ):
             # register FSDP patch
-            from .autogptq_utils import load_fsdp_gptq_patch
-            load_fsdp_gptq_patch(target_module = QuantLinear, torch_dtype = torch_dtype)
+            from .autogptq_utils import register_tensors_as_parameters_patch_rule
+            register_tensors_as_parameters_patch_rule(
+                target_module=QuantLinear, 
+                torch_dtype=torch_dtype,
+            )
 
             # replace
             AutoModelForCausalLM.from_config = _old_from_config

plugins/framework/src/fms_acceleration/model_patcher.py

@@ -161,11 +161,11 @@ class ModelPatcherRule:
     ] = None
 
     def __post_init__(self):
-        if (
-            self.forward is not None
-            and self.forward_builder is not None
-            and self.import_and_maybe_reload is not None
-        ):
+        if sum([
+            self.forward is not None,
+            self.forward_builder is not None,
+            self.import_and_maybe_reload is not None,
+        ])>1:
             raise ValueError(
                 f"Rule '{self.rule_id}' must only have only one of forward, "
                 "foward builder, or import_and_maybe_reload, specified."
@@ -305,7 +305,7 @@ def _import_and_reload(model: torch.nn.Module):
                 elif _target.startswith(module_path):
                     _no_reload.append(rule)
 
-        assert len(_with_reload) <= 1, "cannot have have at most one rule with reload"
+        assert len(_with_reload) <= 1, "can only have at most one rule with reload"
 
         # handle those with reload first
         for rule in _with_reload + _no_reload:

plugins/framework/src/fms_acceleration/utils/test_utils.py

@@ -21,7 +21,6 @@
 from typing import Any, Callable, Dict, List, Set, Tuple, Type
 
 # Third Party
-from torch.nn import CrossEntropyLoss
 import torch
 import yaml
 
@@ -182,12 +181,12 @@ def dummy_custom_loader(self, model_name, **kwargs):
     "dummy custom loader returning dummy model"
     return create_noop_model_with_archs(archs=["DummyModel"])  #
 
-
-class DummyModule(torch.nn.Module):
-    def __init__(self, hidden_size, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self.linear = torch.nn.Linear(hidden_size, hidden_size)
-        self.loss_fn = CrossEntropyLoss()
-
-    def forward(self, X):
-        return self.linear(X)
+@contextmanager
+def instantiate_model_patcher():
+    from fms_acceleration.model_patcher import ModelPatcher
+    old_registrations = ModelPatcher.rules
+    ModelPatcher.rules = {}
+    try:
+        yield 
+    finally:
+        ModelPatcher.rules = old_registrations

plugins/framework/tests/model_patcher_test_utils.py

@@ -0,0 +1,18 @@
+import torch
+
+UNPATCHED_RESPONSE = 0
+PATCHED_RESPONSE = 1
+
+class DummyAttribute(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        return UNPATCHED_RESPONSE
+
+class PatchedAttribute(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        return PATCHED_RESPONSE