git # This is a combination of 3 commits.

Update low_level_zero_plugin.py

Update low_level_zero_plugin.py

fix

fix

fix

Author: flybird11111

.github/workflows/build_on_pr.yml

@@ -208,7 +208,7 @@ jobs:
 
       - name: Execute Unit Testing
         run: |
-          CURL_CA_BUNDLE="" PYTHONPATH=$PWD pytest -m "not largedist" --testmon --testmon-forceselect --testmon-cov=. --durations=10 tests/test_booster/test_plugin/test_3d_plugin.py 
+          CURL_CA_BUNDLE="" PYTHONPATH=$PWD pytest -m "not largedist" --testmon --testmon-forceselect --testmon-cov=. --durations=10 tests/
         env:
           DATA: /data/scratch/cifar-10
           NCCL_SHM_DISABLE: 1

colossalai/booster/plugin/low_level_zero_plugin.py

@@ -351,26 +351,35 @@ def get_param_group_id(self, optimizer: Optimizer, origin_param: Parameter):
                     return group_id
         return -1
 
+    def get_param_group_id(self, optimizer: Optimizer, origin_param: Parameter, lora_param: Parameter):
+        origin_param_id = id(origin_param)
+        lora_param_id = id(lora_param)
+        target_group_id = -1
+        for group_id, param_group in enumerate(optimizer.param_groups):
+            for p in param_group['params']:
+                if id(p) == lora_param_id:
+                    # check if the lora parameter exists.
+                    return -2
+                if id(p) == origin_param_id:
+                    target_group_id = group_id
+        return target_group_id
+    
     def add_lora_params_to_optimizer(self, model, optimizer):
         """ add lora parameters to optimizer """
         name2param= {}
         for name, param in model.named_parameters():
             name2param[name] = param
 
-        optimizer_param_nums = 0
-        for param_group in optimizer.param_groups:
-            optimizer_param_nums += len(param_group['params'])
-
-        # Check if the optimizer is created after the model is transformed into a LoRa model.
-        if len(name2param) != optimizer_param_nums:
-            for name, param in name2param.items():
-                if 'lora_A' in name or 'lora_B' in name:
-                    origin_key = name.replace("lora_A.", "")
-                    origin_key = origin_key.replace("lora_B.", "")
-                    origin_key = origin_key.replace(f"{model.active_adapter}.", "")
-                    origin_param = name2param[origin_key]
-                    group_id = self.get_param_group_id(optimizer, origin_param)
-                    assert group_id != -1, "Parameter error, origin parameter does't exists."
+        for name, param in name2param.items():
+            if 'lora_A' in name or 'lora_B' in name:
+                origin_key = name.replace("lora_A.", "")
+                origin_key = origin_key.replace("lora_B.", "")
+                origin_key = origin_key.replace(f"{model.active_adapter}.", "")
+                origin_param = name2param[origin_key]
+                group_id = self.get_param_group_id(optimizer, origin_param, param)
+                if group_id == -1:
+                    warnings.warn("Origin parameter {origin_key} related to {name} doesn't exist in optimizer param_groups.")
+                elif group_id >= 0:
                     optimizer.param_groups[group_id]['params'].append(param)
 
     def configure(
@@ -384,7 +393,8 @@ def configure(
         if self.lora_enabled:
             from peft import PeftModel
             assert isinstance(model, PeftModel), "The model should have been wrapped as a PeftModel when self.lora_enabled is True"
-            self.add_lora_params_to_optimizer(model, optimizer)
+            if optimizer is not None:
+                self.add_lora_params_to_optimizer(model, optimizer)
 
 
         if not isinstance(model, ModelWrapper):

tests/test_booster/test_plugin/test_dp_plugin_base.py

@@ -1,4 +1,4 @@
-from typing import Callable, Iterator, List, Tuple, Union
+from typing import Callable, Iterator, List, Tuple, Union, Dict
 
 import torch
 import torch.distributed as dist
@@ -51,6 +51,12 @@ def supported_precisions(self) -> List[str]:
     def no_sync(self, model: nn.Module) -> Iterator[None]:
         pass
 
+    def enable_lora(self, model: nn.Module, pretrained_dir: str, lora_config: Dict) -> nn.Module:
+        pass
+
+    def support_lora(self) -> bool:
+        pass
+
 
 def check_dataloader_sharding():
     plugin = DPPluginWrapper()