fix: lora constants

Signed-off-by: Will Johnson <[email protected]>

Author: willmj

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

@@ -344,7 +344,7 @@ def _infer_prefixes_and_module_names(
     ):
         _name = "|".join([PARAM_NAME_ROUTER_SCATTERMOE, *PARAM_NAME_WEIGHT_SCATTERMOE])
         # pylint: disable=anomalous-backslash-in-string
-        _reg = re.compile(f"(.*)\.({_name})\.weight")
+        _reg = re.compile(rf"(.*)\.({_name})\.(?:weight|lora_A\.weight|lora_B\.weight)")
         found = {}
 
         for k in sd_keys:

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

@@ -24,7 +24,9 @@
 KEY_EXPERT_PARALLEL = "expert_parallel"
 DIM_EXPERT = 0
 
-KEY_SCATTERMOE_ROUTER = PARAM_NAME_ROUTER_SCATTERMOE + ".weight"
+KEY_SCATTERMOE_ROUTER = "router.weight"
+KEY_SCATTERMOE_LORA_A_ROUTER = "router.lora_A.weight"
+KEY_SCATTERMOE_LORA_B_ROUTER = "router.lora_B.weight"
 
 # Currently out ScatterMoE drop supports an up/down proj, and
 # and optional gate_proj.

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

@@ -34,6 +34,8 @@
     KEY_EXPERT_PARALLEL,
     KEY_REPLICATE,
     KEY_SCATTERMOE_ROUTER,
+    KEY_SCATTERMOE_LORA_A_ROUTER,
+    KEY_SCATTERMOE_LORA_B_ROUTER,
     get_scattermoe_conv_spec_from_archs,
 )
 from .scattermoe_state_dict import (
@@ -66,7 +68,7 @@ def _hook(grad):
 
     for weight_name, param in state_dict.items():
 
-        if KEY_SCATTERMOE_ROUTER in weight_name:
+        if KEY_SCATTERMOE_ROUTER in weight_name or KEY_SCATTERMOE_LORA_A_ROUTER in weight_name or KEY_SCATTERMOE_LORA_B_ROUTER in weight_name:
             # if its the router, replicate
             param = distribute_tensor(param, device_mesh, reps + [Replicate()])
         elif param.shape[0] > num_experts_per_device: