fix naming

fix naming

fix naming

fix

Author: flybird11111

colossalai/booster/plugin/low_level_zero_plugin.py

@@ -1,5 +1,6 @@
 import logging
 import warnings
+import enum
 import os
 from functools import partial
 from pathlib import Path
@@ -43,6 +44,11 @@ def _convert_floating_point(x, dtype: torch.dtype = torch.float16):
 
 SUPPORTED_PRECISION = ["fp16", "bf16", "fp32"]
 
+class OptimizerParamCheckState(enum.Enum):
+    ORIGIN_PARAM_FINDED = 0
+    ORIGIN_PARAM_NOT_FIND = -1
+    LORA_PARM_EXISTED = -2
+
 
 class LowLevelZeroModel(ModelWrapper, AMPModelMixin):
     def __init__(self, module: nn.Module, precision: str) -> None:
@@ -354,15 +360,18 @@ def get_param_group_id(self, optimizer: Optimizer, origin_param: Parameter):
     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
+        target_group_id = None
         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
+                    return target_group_id, OptimizerParamCheckState.LORA_PARM_EXISTED
                 if id(p) == origin_param_id:
                     target_group_id = group_id
-        return target_group_id
+        if target_group_id is not None:
+            return target_group_id, OptimizerParamCheckState.ORIGIN_PARAM_FINDED
+        else:
+            return target_group_id, OptimizerParamCheckState.ORIGIN_PARAM_NOT_FIND
 
     def add_lora_params_to_optimizer(self, model, optimizer):
         """ add lora parameters to optimizer """
@@ -374,12 +383,12 @@ def add_lora_params_to_optimizer(self, model, optimizer):
             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_key = origin_key.replace(f"{model.active_adapter}", "base_layer")
                 origin_param = name2param[origin_key]
-                group_id = self.get_param_group_id(optimizer, origin_param, param)
-                if group_id == -1:
+                group_id, check_state = self.get_param_group_id(optimizer, origin_param, param)
+                if check_state == OptimizerParamCheckState.ORIGIN_PARAM_NOT_FIND:
                     warnings.warn("Origin parameter {origin_key} related to {name} doesn't exist in optimizer param_groups.")
-                elif group_id >= 0:
+                elif check_state == OptimizerParamCheckState.ORIGIN_PARAM_FINDED and group_id is not None and group_id >= 0:
                     optimizer.param_groups[group_id]['params'].append(param)
 
     def configure(

colossalai/pipeline/p2p.py

@@ -44,6 +44,20 @@ def _cuda_safe_tensor_to_object(tensor: torch.Tensor, tensor_size: torch.Size) -
 
     return unpickle
 
+def check_for_nccl_backend(group):
+
+    pg = group or c10d._get_default_group()
+    # Gate PG wrapper check on Gloo availability.
+    if c10d._GLOO_AVAILABLE:
+        # It is not expected for PG to be wrapped many times, but support it just
+        # in case
+        while isinstance(pg, c10d._ProcessGroupWrapper):
+            pg = pg.wrapped_pg
+
+    return (
+        c10d.is_nccl_available() and
+        pg.name() == c10d.Backend.NCCL
+    )
 
 def _broadcast_object_list(
     object_list: List[Any], src: int, group: ProcessGroup, device: Optional[Union[torch.device, str, int]] = None
@@ -65,7 +79,7 @@ def _broadcast_object_list(
         c10d._warn_not_in_group("broadcast_object_list")
         return
 
-    is_nccl_backend = c10d._check_for_nccl_backend(group)
+    is_nccl_backend = check_for_nccl_backend(group)
     current_device = None
 
     if device is not None:

requirements/requirements-test.txt

@@ -18,5 +18,5 @@ SentencePiece
 ninja
 flash_attn==2.0.5
 datasets
-peft
+peft>=0.7.1
 #auto-gptq now not support torch1.12

requirements/requirements.txt

@@ -14,3 +14,4 @@ einops
 sentencepiece
 google
 protobuf
+peft>=0.7.1