filter out router

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

Author: willmj

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

@@ -113,8 +113,11 @@ def save_fsdp_optimizer(
     (model_state_dict, optimizer_state_dict) = get_state_dict(model, optimizer)
 
     # filter out lora state dict
+    # TODO: Once expert layers are supported for LoRA tuning
+    # remove the "router" filtering
     lora_state_dict = {
-        k: v for k, v in model_state_dict.items() if "lora_A" in k or "lora_B" in k
+        k: v for k, v in model_state_dict.items()
+        if ("lora_A" in k or "lora_B" in k) and "router" not in k
     }
 
     # - save model