fix

fix

fix

fix

fix

Author: flybird11111

colossalai/booster/plugin/low_level_zero_plugin.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 import os
 from functools import partial
 from pathlib import Path
@@ -9,6 +10,7 @@
 
 import torch
 import torch.nn as nn
+from torch.nn import Parameter
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
 from torch.utils._pytree import tree_map
@@ -335,13 +337,27 @@ def enable_lora(
         from peft import PeftModel, get_peft_model
         assert not isinstance(model, LowLevelZeroModel), "Lora should be enabled before boosting the model."
         self.lora_enabled = True
+        warnings.warn("You have enabled LoRa training. Please check the hyperparameter such as lr")
 
         if pretrained_dir is None:
             peft_model = get_peft_model(model, lora_config)
         else:
             peft_model = PeftModel.from_pretrained(model, pretrained_dir, is_trainable=True)
         return peft_model
 
+    def get_param_group_id(self, optimizer: Optimizer, origin_param: Parameter, add_param: Parameter):
+        origin_param_id = id(origin_param)
+        add_param_id = id(add_param)
+        group_id = -1
+        for pg_id, param_group in enumerate(optimizer.param_groups):
+            for p in param_group['params']:
+                if id(p) == add_param_id:
+                    return -2
+                if id(p) == origin_param_id:
+                    group_id = pg_id
+        return group_id
+
+    
     def configure(
         self,
         model: nn.Module,
@@ -353,12 +369,21 @@ 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"
-            
-            optim_params_nums = 0
-            for param_group in optimizer.param_groups:
-                optim_params_nums += len(param_group['params'])
-            model_params_nums = len(list(model.named_parameters()))
-            assert optim_params_nums == model_params_nums, "Optimizer should be initialized after enabling lora."
+
+            # add lora parameters to optimizer
+            name2param= {}
+            for name, param in model.named_parameters():
+                name2param[name] = param
+            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)
+                    assert group_id != -1, "Parameter error, origin parameter does't exists."
+                    if group_id >= 0:
+                        optimizer.param_groups[group_id]['params'].append(param)
 
         if not isinstance(model, ModelWrapper):
             model = LowLevelZeroModel(model, self.precision)

tests/test_booster/test_plugin/test_low_level_zero_plugin.py

@@ -48,6 +48,7 @@ def run_fn(stage, model_fn, data_gen_fn, output_transform_fn, lora_config=None)
 
     except Exception as e:
         return repr(e)
+        # raise e
 
 
 

tests/test_checkpoint_io/test_low_level_zero_checkpoint_io.py

@@ -168,7 +168,7 @@ def check_low_level_zero_lora_checkpointIO(stage: int, shard: bool, offload: boo
 
 def run_dist(rank, world_size, port):
     colossalai.launch(config=(dict()), rank=rank, world_size=world_size, port=port, host="localhost")
-    # check_low_level_zero_checkpointIO()
+    check_low_level_zero_checkpointIO()
     check_low_level_zero_lora_checkpointIO()
     torch.cuda.empty_cache()