fix

fix

fix

fix

fix

fix

fix

Author: flybird11111

colossalai/booster/plugin/low_level_zero_plugin.py

@@ -1,14 +1,14 @@
 import logging
+import warnings
 import os
 from functools import partial
 from pathlib import Path
 from types import MethodType
 from typing import Callable, Dict, Iterator, List, Optional, Tuple, Dict
 
-from peft import LoraConfig, TaskType, get_peft_model
-
 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 +335,44 @@ 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 hyperparameters 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):
+        origin_param_id = id(origin_param)
+        for group_id, param_group in enumerate(optimizer.param_groups):
+            for p in param_group['params']:
+                if id(p) == origin_param_id:
+                    return group_id
+        return -1
+    
+    def add_lora_para_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."
+                    optimizer.param_groups[group_id]['params'].append(param)
+    
     def configure(
         self,
         model: nn.Module,
@@ -353,12 +384,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"
-            
-            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."
+            self.add_lora_para_to_optimizer(model, optimizer)
+
 
         if not isinstance(model, ModelWrapper):
             model = LowLevelZeroModel(model, self.precision)

colossalai/zero/low_level/bookkeeping/gradient_store.py

@@ -82,6 +82,7 @@ def get_working_grads_by_group_id(self, group_id: int) -> List:
         """
 
         grad_list = []
+        # When using LoRa and the user sets multiple param_groups, it is possible that some param_groups have no parameters with gradients.
         if group_id not in self._grads_of_params.keys():
             return grad_list
         for param_grads in self._grads_of_params[group_id].values():

tests/test_booster/test_plugin/test_low_level_zero_plugin.py

@@ -24,11 +24,11 @@ def run_fn(stage, model_fn, data_gen_fn, output_transform_fn, lora_config=None)
         plugin = LowLevelZeroPlugin(stage=stage, max_norm=1.0, initial_scale=2**5)
         booster = Booster(plugin=plugin)
         model = model_fn()
+        optimizer = HybridAdam(model.parameters(), lr=1e-3)
 
         if lora_config is not None:
             model = booster.enable_lora(model, lora_config=lora_config)
 
-        optimizer = HybridAdam(model.parameters(), lr=1e-3)
         criterion = lambda x: x.mean()
         data = data_gen_fn()
 
@@ -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

@@ -80,9 +80,10 @@ def run_fn(stage, shard, offload, model_fn, data_gen_fn, output_transform_fn, lo
         booster = Booster(plugin=plugin)
         new_booster = Booster(plugin=new_plugin)
         model = model_fn()
+        optimizer = HybridAdam(model.parameters(), lr=1e-3)
         new_model = deepcopy(model)
+        new_optimizer = HybridAdam(new_model.parameters(), lr=1e-3)
         model = booster.enable_lora(model, lora_config=lora_config)
-        optimizer = HybridAdam(model.parameters(), lr=1e-3)
         criterion = lambda x: x.mean()
         data = data_gen_fn()
 
@@ -107,7 +108,6 @@ def run_fn(stage, shard, offload, model_fn, data_gen_fn, output_transform_fn, lo
             booster.save_lora_as_pretrained(model, model_ckpt_path)
             booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=False)
             new_model = new_booster.enable_lora(new_model, pretrained_dir=model_ckpt_path, lora_config=lora_config)
-            new_optimizer = HybridAdam(new_model.parameters(), lr=1e-3)
             new_model, new_optimizer, criterion, _, _ = new_booster.boost(new_model, new_optimizer, criterion)
             check_state_dict_equal(model.state_dict(), new_model.state_dict(), False)
 
@@ -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()