feat: lora for accelerated moe

plugins/accelerated-moe/src/fms_acceleration_moe/framework_plugin_scattermoe.py

@@ -59,6 +59,7 @@ def augmentation(
         modifiable_args: Tuple[LoraConfig],
     ):
         rank, world_size = 0, 1
+        (peft_config,) = modifiable_args
         if torch.distributed.is_initialized():
             world_size = torch.distributed.get_world_size()
             # we do not need to use the fallback as this is wrapped in an `is_initialized` block
@@ -78,6 +79,7 @@ def augmentation(
             world_size=world_size,
             ep_degree=self._ep_degree,
             mixed_precision=False,  # Currently this is hardcoded to OFF
+            lora_config=peft_config,
         )
         return model, modifiable_args
 

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

@@ -53,6 +53,10 @@
 KEY_MODEL = "model"
 KEY_OPTIMIZER = "optimizer"
 
+ADAPTER_CONFIG_NAME = "adapter_config.json"
+ADAPTER_WEIGHTS_NAME = "adapter_model.bin"
+ADAPTER_SAFE_WEIGHTS_NAME = "adapter_model.safetensors"
+
 # Below are rewrite of HF FSDP model saving functions to be able to handle
 # that the parameters are now a mixture of regular and Dtensors.
 # - these functions are found in accelerate.utils.fsdp_utils.py
@@ -110,15 +114,30 @@ def save_fsdp_optimizer(
     # get the state dicts for model and optimize
     (model_state_dict, optimizer_state_dict) = get_state_dict(model, optimizer)
 
-    # - save model
-    ckpt_model = os.path.join(output_dir, f"{FSDP_MODEL_NAME}_{MODEL_INDEX}")
-    os.makedirs(ckpt_model, exist_ok=True)
-    logger.info(f"Saving model to {ckpt_model}")
-    dcp.save(
-        state_dict={KEY_MODEL: model_state_dict},
-        storage_writer=dcp.FileSystemWriter(ckpt_model),
-        planner=DefaultSavePlanner(),
-    )
+    # filter out lora state dict
+    lora_state_dict = {
+        k: v for k, v in model_state_dict.items() if "lora_A" in k or "lora_B" in k
+    }
+
+    # - save mode
+    if lora_state_dict:
+        ckpt_model = os.path.join(output_dir, f"{FSDP_MODEL_NAME}_{MODEL_INDEX}")
+        os.makedirs(ckpt_model, exist_ok=True)
+        logger.info(f"Saving lora model to {ckpt_model}")
+        dcp.save(
+            state_dict={KEY_MODEL: lora_state_dict},
+            storage_writer=dcp.FileSystemWriter(ckpt_model),
+            planner=DefaultSavePlanner(),
+        )
+    else: 
+        ckpt_model = os.path.join(output_dir, f"{FSDP_MODEL_NAME}_{MODEL_INDEX}")
+        os.makedirs(ckpt_model, exist_ok=True)
+        logger.info(f"Saving ft model to {ckpt_model}")
+        dcp.save(
+            state_dict={KEY_MODEL: model_state_dict},
+            storage_writer=dcp.FileSystemWriter(ckpt_model),
+            planner=DefaultSavePlanner(),
+        )
     logger.info(f"Model saved to {ckpt_model}")
 
     # - save optimizer
@@ -344,7 +363,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|lora_B)")
         found = {}
 
         for k in sd_keys:
@@ -398,6 +417,7 @@ def _infer_prefixes_and_module_names(
         #   model param and they need to be cat
         for scatter_key, list_of_params in checkpoint_metadata.items():
             scatter_key_fqdn = ".".join([prefix, module_name, scatter_key])
+
             scatter_param = sd[scatter_key_fqdn]
 
             # remove from state dict
@@ -439,12 +459,49 @@ def _infer_prefixes_and_module_names(
             #  it will go by order of scatter keys
             scatter_keys = sorted(scatter_params.keys())
 
+            scatter_keys_fqdn = [".".join([prefix, module_name, scatter_key]) for scatter_key in scatter_keys]
+
             assert (
                 len(scatter_keys) > 0
             ), f"Obtained zero scatter keys for model_key '{model_key}'"
 
-            if len(scatter_keys) == 1:
+            if any("lora_A" in k for k in scatter_keys) and any("lora_B" in k for k in scatter_keys):
+                # If lora, split input linear and output linear into lora layers
+                def transform_model_key(model_key, lora_key):
+                    lora_parts = lora_key.split(".")
+                    model_parts = model_key.split(".")
+
+                    try:
+                        lora_index = lora_parts.index("block_sparse_moe")
+                        model_index = model_parts.index("block_sparse_moe")
+                    except ValueError:
+                        raise ValueError("Both keys must contain 'block_sparse_moe'")
+
+                    # Replace the component after 'block_sparse_moe' in lora_key
+                    updated_lora_parts = lora_parts[:]
+                    updated_lora_parts[lora_index + 1] = model_parts[model_index + 1]
+
+                    # Return the updated lora parts as the model key
+                    return ".".join(updated_lora_parts)
+                for i, lora_key in enumerate(scatter_keys):
+                    new_model_key = transform_model_key(model_key, scatter_keys_fqdn[i])
+                    if len(scatter_keys) == 2:
+                        sd[new_model_key] = scatter_params[lora_key]
+                    else:
+                        if "lora_A" in new_model_key:
+                            filtered_keys = [k for k in scatter_keys if "lora_A" in k]
+                        elif "lora_B" in new_model_key:
+                            filtered_keys = [k for k in scatter_keys if "lora_B" in k]
+                        else:
+                            raise ValueError(f"Unexpected LoRA key type in {new_model_key}")
+
+                        sd[new_model_key] = torch.cat(
+                            [scatter_params[k] for k in filtered_keys], dim=1
+                        )
+
+            elif len(scatter_keys) == 1:
                 sd[model_key] = scatter_params[scatter_keys[0]]
+
             else:
                 # unfortunately, there this is a in
                 # scattermoe_state_dict._maybe_reshape_scattermoe_expert_weights
@@ -468,29 +525,50 @@ def save_sharded_safetensors(
     metadata: Dict,
     max_shard_size: Union[int, str] = "5GB",
 ):
-    filename_pattern = SAFE_WEIGHTS_NAME.replace(".bin", "{suffix}.bin").replace(
-        ".safetensors", "{suffix}.safetensors"
-    )
-    state_dict_split = split_torch_state_dict_into_shards(
-        input_state_dict,
-        filename_pattern=filename_pattern,
-        max_shard_size=max_shard_size,
-    )
-    index = {
-        "metadata": state_dict_split.metadata,
-        "weight_map": state_dict_split.tensor_to_filename,
-    }
-    # Save the index
-    with open(
-        os.path.join(save_directory, SAFE_WEIGHTS_INDEX_NAME), "w", encoding="utf-8"
-    ) as f:
-        content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-        f.write(content)
-
-    filename_to_tensors = state_dict_split.filename_to_tensors.items()
-    for shard_file, tensors in filename_to_tensors:
-        shard = {tensor: input_state_dict[tensor].contiguous() for tensor in tensors}
-        save_file(shard, os.path.join(save_directory, shard_file), metadata=metadata)
+    lora = False
+    for name, _ in input_state_dict.items():
+        if "lora_A" or "lora_B" in name:
+            lora = True
+            break
+
+    if not lora:
+        filename_pattern = SAFE_WEIGHTS_NAME.replace(".bin", "{suffix}.bin").replace(
+            ".safetensors", "{suffix}.safetensors"
+        )
+        state_dict_split = split_torch_state_dict_into_shards(
+            input_state_dict,
+            filename_pattern=filename_pattern,
+            max_shard_size=max_shard_size,
+        )
+
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        # Save the index
+        with open(
+            os.path.join(save_directory, SAFE_WEIGHTS_INDEX_NAME), "w", encoding="utf-8"
+        ) as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+        filename_to_tensors = state_dict_split.filename_to_tensors.items()
+        for shard_file, tensors in filename_to_tensors:
+            shard = {tensor: input_state_dict[tensor].contiguous() for tensor in tensors}
+            save_file(shard, os.path.join(save_directory, shard_file), metadata=metadata)
+    else:
+        filename_pattern = ADAPTER_SAFE_WEIGHTS_NAME.replace(".bin", "{suffix}.bin").replace(
+            ".safetensors", "{suffix}.safetensors"
+        )
+        state_dict_split = split_torch_state_dict_into_shards(
+            input_state_dict,
+            filename_pattern=filename_pattern,
+            max_shard_size=max_shard_size,
+        )
+        filename_to_tensors = state_dict_split.filename_to_tensors.items()
+        for shard_file, tensors in filename_to_tensors:
+            shard = {tensor: input_state_dict[tensor].contiguous() for tensor in tensors}
+            save_file(shard, os.path.join(save_directory, shard_file), metadata=metadata)
 
 
 # --------------------------- SCRIPT -------------------------

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:
@@ -91,8 +93,9 @@ def _hook(grad):
         )
 
         # install gradient scaling hook
-        if KEY_SCATTERMOE_ROUTER not in weight_name:
-            param.register_hook(_hook)
+        if KEY_SCATTERMOE_ROUTER not in weight_name and KEY_SCATTERMOE_LORA_A_ROUTER not in weight_name and KEY_SCATTERMOE_LORA_B_ROUTER not in weight_name:
+            if param.requires_grad:
+                param.register_hook(_hook) 
 
         # register the sharded parameter onto the megablocks.dmoe
         mod.register_parameter(name, param)

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

@@ -27,6 +27,8 @@
 from .scattermoe_constants import (
     DIM_EXPERT,
     KEY_SCATTERMOE_ROUTER,
+    KEY_SCATTERMOE_LORA_A_ROUTER,
+    KEY_SCATTERMOE_LORA_B_ROUTER,
     PARAM_NAME_WEIGHT_SCATTERMOE,
 )
 
@@ -121,6 +123,8 @@ def _insert(L: List, i: int, v):
             n -= 1
         L[i] = v
 
+    lora = False
+
     # if expert_name = input_linear|output_linear|input_linear
     # - in this case will map
     # - input_linear: [w1, w3], output_linear: {w2}
@@ -147,6 +151,12 @@ def _insert(L: List, i: int, v):
     # `w1.weight`: [...]
     _map = defaultdict(list)
     prefix = f"{prefix}.{instance_name}."
+    # Lora case where it prefix looks like base_model.model.model...
+    # instead of model...
+    if not prefix.startswith("model."):
+        prefix=prefix.replace("base_model.model.", "", 1)
+        lora=True
+
     for k, stfile in weight_map.items():
         if not k.startswith(prefix):
             continue
@@ -163,13 +173,22 @@ def _insert(L: List, i: int, v):
                 f"'{router_name}' or expert_name '{expert_name}'"
             )
         if m.group(1) == router_name:
-            _map[KEY_SCATTERMOE_ROUTER].append((k, stfile))
+            if lora:
+                _map["router.lora_A.default.weight"].append((k, stfile))
+                _map["router.lora_B.default.weight"].append((k, stfile))
+            else:
+                _map[KEY_SCATTERMOE_ROUTER].append((k, stfile))
         elif m.group(1) in expert_name:
             index = m.group(2)
             index = 0 if index is None else int(index)
             mod = None
-            for mod in expert_map.get(m.group(1), expert_map.get(m.group(3))):
-                _insert(_map[f"{mod}.weight"], index, (k, stfile))
+            if not lora:
+                for mod in expert_map.get(m.group(1), expert_map.get(m.group(3))):
+                    _insert(_map[f"{mod}.weight"], index, (k, stfile))
+            else:
+                for mod in expert_map.get(m.group(1), expert_map.get(m.group(3))):
+                    _insert(_map[f"{mod}.lora_A"], index, (k, stfile))
+                    _insert(_map[f"{mod}.lora_B"], index, (k, stfile))
 
             assert mod is not None, f"cannot map '{rel_k}'"
 
@@ -295,7 +314,7 @@ def get_state_dict_from_checkpoint_metadata(
         # go by one weight at a time.
         for scatter_key, vs in checkpoint_metadata.items():
 
-            if KEY_SCATTERMOE_ROUTER in scatter_key:
+            if KEY_SCATTERMOE_ROUTER in scatter_key or KEY_SCATTERMOE_LORA_A_ROUTER in scatter_key or KEY_SCATTERMOE_LORA_B_ROUTER in scatter_key:
                 k, fi = vs[0]  # only one item
                 param = files[fi].get_tensor(k)