Update sharded axis

Signed-off-by: Jingyu Xin <[email protected]>

Author: jingyu-ml

modelopt/torch/peft/convert.py

@@ -44,11 +44,8 @@ def update_model(
     Returns:
         The updated model with LoRA adapters
     """
-    # Validate config by converting to PEFTConfig if needed
-
     # Check if model is already in PEFT mode by looking for LoRA modules
     if not is_peft_model(model):
-        # First time - need to convert to PEFT mode
         apply_mode(model, mode=[("peft", config)], registry=PEFTModeRegistry)
     else:
         if not isinstance(config, PEFTConfig):

modelopt/torch/peft/lora/layer.py

@@ -64,6 +64,8 @@ def _register_adapter(
         self.add_module(f"lora_b_{adapter_name}", lora_b)
 
         # Store in adapter dictionary with explicit rank
+        if adapter_name in self._lora_adapters:
+            raise ValueError(f"adapter_name: {adapter_name} is already exist..")
         self._lora_adapters[adapter_name] = {
             "lora_a": lora_a,
             "lora_b": lora_b,

modelopt/torch/peft/lora/tp_layer.py

@@ -6,6 +6,7 @@
 import torch.nn as nn
 import torch.nn.init as init
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
 
 from ..config import PEFTAttributeConfig
 from .layer import LoRAModule, LoRAModuleRegistry
@@ -129,6 +130,40 @@ def update_layer_lora(
             adapter_name, lora_a, lora_b, attr_config.rank, attr_config.scale, attr_config.enable
         )
 
+    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+        """Sharding along axis 0 for ColumnParallelLinear, bias not sharded.
+
+        For ColumnParallelLinear:
+        - lora_a weight: sharded at dim 0
+        - lora_b weight: sharded at dim 0
+        """
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+
+        if hasattr(self, "_lora_adapters"):
+            lora_state_dict = {}
+            state_dict = self.state_dict(prefix="", keep_vars=True)
+
+            for adapter_name in self._lora_adapters:
+                lora_a_key = f"lora_a_{adapter_name}.weight"
+                lora_b_key = f"lora_b_{adapter_name}.weight"
+
+                if lora_a_key in state_dict:
+                    lora_state_dict[lora_a_key] = state_dict[lora_a_key]
+                if lora_b_key in state_dict:
+                    lora_state_dict[lora_b_key] = state_dict[lora_b_key]
+
+            lora_sharding_dims = {}
+            for key in lora_state_dict:
+                lora_sharding_dims[key] = 0
+
+            if lora_state_dict:
+                lora_sharded = make_sharded_tensors_for_checkpoint(
+                    lora_state_dict, prefix, lora_sharding_dims, sharded_offsets
+                )
+                sharded_state_dict.update(lora_sharded)
+
+        return sharded_state_dict
+
 
 @LoRAModuleRegistry.register({RowParallelLinear: "megatron_RowParallelLinear"})
 class _LoRAMegatronRowParallelLinear(_MegatronParallelLoRABase):
@@ -172,6 +207,43 @@ def update_layer_lora(
             adapter_name, lora_a, lora_b, attr_config.rank, attr_config.scale, attr_config.enable
         )
 
+    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+        """Sharding along axis 1 for RowParallelLinear, bias not sharded.
+
+        For RowParallelLinear:
+        - lora_a weight: sharded at dim 1 (RowParallelLinear)
+        - lora_b weight: sharded at dim 0 (ColumnParallelLinear)
+        """
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+
+        if hasattr(self, "_lora_adapters"):
+            lora_state_dict = {}
+            state_dict = self.state_dict()
+
+            for adapter_name in self._lora_adapters:
+                lora_a_key = f"lora_a_{adapter_name}.weight"
+                lora_b_key = f"lora_b_{adapter_name}.weight"
+
+                if lora_a_key in state_dict:
+                    lora_state_dict[lora_a_key] = state_dict[lora_a_key]
+                if lora_b_key in state_dict:
+                    lora_state_dict[lora_b_key] = state_dict[lora_b_key]
+
+            lora_sharding_dims = {}
+            for key in lora_state_dict:
+                if "lora_a_" in key:
+                    lora_sharding_dims[key] = 1
+                elif "lora_b_" in key:
+                    lora_sharding_dims[key] = 0
+
+            if lora_state_dict:
+                lora_sharded = make_sharded_tensors_for_checkpoint(
+                    lora_state_dict, prefix, lora_sharding_dims, sharded_offsets
+                )
+                sharded_state_dict.update(lora_sharded)
+
+        return sharded_state_dict
+
 
 # Register quantized versions if available
 if QUANT_MODULES_AVAILABLE: