Some minor updates

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

Author: jingyu-ml

modelopt/torch/peft/config.py

@@ -72,6 +72,17 @@ class PEFTAttributeConfig(ModeloptBaseConfig):
         description="Custom initialization function for LoRA B matrix. Default to zero initialization.",
     )
 
+    @field_validator("lora_a_init", "lora_b_init")
+    @classmethod
+    def validate_init_method(cls, v):
+        """Validate initialization method is supported."""
+        valid_methods = {"kaiming_init", "zero_init"}
+        if v not in valid_methods:
+            raise ValueError(
+                f"Invalid initialization method: {v}. Supported methods: {', '.join(valid_methods)}"
+            )
+        return v
+
     @field_validator("rank")
     @classmethod
     def validate_rank(cls, v):

modelopt/torch/peft/convert.py

@@ -62,7 +62,8 @@ def update_model(
     Returns:
         The updated model with LoRA adapters
     """
-    assert is_megatron_core_model(model), "We only support mcore format for the PEFT mode"
+    if not is_megatron_core_model(model):
+        raise ValueError("PEFT mode currently supports Megatron-Core models only.")
 
     # Check if model is already in PEFT mode by looking for LoRA modules
     if not is_peft_model(model):

modelopt/torch/peft/lora/layer.py

@@ -72,7 +72,7 @@ def _register_adapter(
 
         # Store in adapter dictionary with explicit rank
         if adapter_name in self._lora_adapters:
-            raise ValueError(f"adapter_name: {adapter_name} is already exist..")
+            raise ValueError(f"Adapter '{adapter_name}' already exists.")
         self._lora_adapters[adapter_name] = {
             "lora_a": lora_a,
             "lora_b": lora_b,
@@ -94,10 +94,7 @@ def update_layer_lora(
 
         Args:
             adapter_name: Name for the new adapter
-            rank: Rank of the LoRA decomposition (default: 64)
-            scale: Scale factor for the LoRA output (default: 1.0)
-            lora_a_init: Optional initialization function for LoRA A matrix
-            lora_b_init: Optional initialization function for LoRA B matrix
+            attr_config: PEFTAttributeConfig containing rank, scale, and initialization settings
         """
         raise NotImplementedError("Subclasses must implement update_layer_lora")