logic

Signed-off-by: Will Johnson <mwjohnson728@gmail.com>

Author: willmj

plugins/accelerated-moe/src/fms_acceleration_moe/utils/checkpoint_utils.py

@@ -405,14 +405,27 @@ def _infer_prefixes_and_module_names(
         #        'model-00001-of-00002.safetensors')],
         #      'router.weight': [('model.layers.0.block_sparse_moe.router.layer.weight',
         #        'model-00001-of-00002.safetensors')]})
+        ip_op_layers = False
+        router_layer = False
+        lora_utils = None
+        if lora:
+            for namex, _ in sd.items():
+                if "w1" in namex:
+                    ip_op_layers = True
+                    break
+            for namex, _ in sd.items():
+                if "router" in namex:
+                    router_layer = True
+                    break
+            lora_utils = [router_layer, ip_op_layers]
 
         checkpoint_metadata = get_checkpoint_meta_from_sharded_safetensor(
             weight_map,
             prefix,
             module_name,
             router_name,
             expert_name,
-            lora_utils=lora,
+            lora_utils=lora_utils,
         )
 
         model2scatter = defaultdict(dict)

plugins/accelerated-moe/src/fms_acceleration_moe/utils/scattermoe_state_dict.py

@@ -89,7 +89,7 @@ def get_checkpoint_meta_from_sharded_safetensor(
     expert_name: str = "experts",  # e.g., named "experts" within block_sparse_moe
     expert_map: Dict = None,  # map -> [w1,w2,w3]
     lora_start: bool = False,  # if lora is detected in prepare_scattermoe.py
-    lora_utils: bool = False,  # if lora is detected in checkpoint_utils.py
+    lora_utils: List = None,  # if lora is detected in checkpoint_utils.py
     target_modules: Dict = None,  # target modules from prepare_scattermoe.py
 ) -> Dict[str, List[Tuple]]:
     """
@@ -171,7 +171,7 @@ def _insert(L: List, i: int, v):
                 f"'{router_name}' or expert_name '{expert_name}'"
             )
         if m.group(1) == router_name:
-            if lora_utils:
+            if lora_utils[0]:
                 _map[KEY_SCATTERMOE_LORA_A_ROUTER].append((k, stfile))
                 _map[KEY_SCATTERMOE_LORA_B_ROUTER].append((k, stfile))
             else:
@@ -184,7 +184,7 @@ def _insert(L: List, i: int, v):
                 index = m.group(2)
                 index = 0 if index is None else int(index)
                 mod = None
-                if not lora_utils:
+                if not lora_utils[1]:
                     for mod in expert_map.get(m.group(1), expert_map.get(m.group(3))):
                         _insert(_map[f"{mod}.weight"], index, (k, stfile))
                 else: