format

Author: lxd-cumt

flagscale/train/peft/lora.py

@@ -1,19 +1,22 @@
-import torch
-import torch.nn as nn
-from typing import Any, Callable, List, Optional, Tuple, Type, Dict
-from dataclasses import dataclass, field
+import logging
+
 from collections import namedtuple
+from dataclasses import dataclass, field
 from functools import partial
-import logging
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+
+import torch
+import torch.nn as nn
+
+from megatron.training.utils import unwrap_model
 
 from flagscale.train.peft.peft import PEFT, AdapterWrapper
 from flagscale.train.peft.utils import (
+    ParallelLinearAdapter,
     get_adapter_attributes_from_linear,
     is_expert_linear,
-    ParallelLinearAdapter,
     match_module,
 )
-from megatron.training.utils import unwrap_model
 
 
 class LoRALinear(AdapterWrapper):
@@ -25,10 +28,7 @@ class to provide a specific implementation of the forward method.
     """
 
     def forward(
-        self,
-        x: torch.Tensor,
-        *args,
-        **kwargs,
+        self, x: torch.Tensor, *args, **kwargs
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         # pylint: disable=C0115,C0116
         linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs)
@@ -44,10 +44,8 @@ class LoRA(PEFT, peft_type='lora'):
     LoRA uses a low-rank projection to adapt the weights of a pre-trained model to a new downstream task.
     This class facilitates the application of LoRA to specific modules within the model architecture.
     """
-    def __init__(
-        self,
-        config,
-    ):
+
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.target_modules = config.lora_target_modules
@@ -72,9 +70,13 @@ def transform(self, m: nn.Module, name=None, prefix=None):
         if (ans := match_module(m, name, prefix, self.target_modules)) is not None:
             (match, full_name) = ans
 
-            input_is_parallel, in_features, out_features, disable_sp_comm, base_linear_is_parallel = (
-                get_adapter_attributes_from_linear(m)
-            )
+            (
+                input_is_parallel,
+                in_features,
+                out_features,
+                disable_sp_comm,
+                base_linear_is_parallel,
+            ) = get_adapter_attributes_from_linear(m)
 
             adapter = ParallelLinearAdapter(
                 in_features=in_features,
@@ -94,12 +96,11 @@ def transform(self, m: nn.Module, name=None, prefix=None):
 
             lora_linear = LoRALinear(m, adapter)
             return lora_linear
-            
+
         return m
-    
+
     def freeze_model(self, model: nn.Module):
-        """Freeze main model for lora
-        """
+        """Freeze main model for lora"""
         for name, param in model.named_parameters():
             param.requires_grad = False
         for name, param in model.named_parameters():
@@ -115,7 +116,7 @@ def load_state_dict_hook_remap_main_model_params(
             strict: bool,
             missing_keys: List[str],
             unexpected_keys: List[str],
-            errors: List[Any]  
+            errors: List[Any],
         ):
             old_keys = [prefix + "weight", prefix + "bias"]
             new_keys = [prefix + "to_wrap.weight", prefix + "to_wrap.bias"]
@@ -125,7 +126,7 @@ def load_state_dict_hook_remap_main_model_params(
                         state_dict[new_key] = state_dict.pop(old_key)
                     else:
                         state_dict.pop(old_key)
-            
+
         for name, module in model.named_modules():
             if isinstance(module, LoRALinear):
                 module.register_load_state_dict_pre_hook(
@@ -142,13 +143,11 @@ def load_state_dict_hook_ignore_param_names(
                         f"{param_name} being removed from incompatible_keys.missing_keys in this model"
                     )
                     incompatible_keys.missing_keys.remove(param_name)
+
         lora_param_names = []
         for name in model.state_dict().keys():
             if "adapter" in name:
                 lora_param_names.append(name)
         model.register_load_state_dict_post_hook(
-            partial(
-                load_state_dict_hook_ignore_param_names,
-                lora_param_names,
-            )
+            partial(load_state_dict_hook_ignore_param_names, lora_param_names)
         )

flagscale/train/peft/peft.py

@@ -1,7 +1,8 @@
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+
 import torch
 import torch.nn as nn
-from typing import Any, Callable, List, Optional, Tuple, Type, Dict
-from abc import ABC, abstractmethod
 
 
 class PEFT(ABC):
@@ -11,6 +12,7 @@ class PEFT(ABC):
     large language models efficiently by modifying only a small subset of the model's
     parameters.
     """
+
     _REGISTRY: Dict[str, Type["PEFT"]] = {}
 
     def __init_subclass__(cls, *, peft_type: str, **kwargs):
@@ -19,14 +21,14 @@ def __init_subclass__(cls, *, peft_type: str, **kwargs):
             raise ValueError(f"peft_type={peft_type} registered already!")
         cls._REGISTRY[peft_type] = cls
         cls._peft_type = peft_type
-    
+
     @classmethod
     def from_config(cls, config) -> "PEFT":
         peft_type = config.peft_type
         if peft_type not in cls._REGISTRY:
             raise KeyError(f"Unsupported peft_type: {peft_type}, registered: {list(cls._REGISTRY)}")
         sub_cls = cls._REGISTRY[peft_type]
-        return sub_cls(config)    
+        return sub_cls(config)
 
     @abstractmethod
     def transform(self, module, name=None, prefix=None):
@@ -60,10 +62,9 @@ def apply_transform(self, model: nn.Module):
             if replaced_module == module:
                 continue
             model.set_submodule(full_name, replaced_module)
-    
+
     def freeze_model(self, model: nn.Module):
-        """Apply a default freeze method to the model.
-        """
+        """Apply a default freeze method to the model."""
         pass
 
     def load_state_dict_pre_hooks(self, model: nn.Module):
@@ -144,5 +145,7 @@ def state_dict(self, destination=None, prefix='', keep_vars=False):
         # Get state dict of the main module
         self.to_wrap.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
         # Store adapter state dict under the "adapter" prefix in the destination dict
-        self.adapter.state_dict(destination=destination, prefix=f'{prefix}adapter.', keep_vars=keep_vars)
+        self.adapter.state_dict(
+            destination=destination, prefix=f'{prefix}adapter.', keep_vars=keep_vars
+        )
         return destination

flagscale/train/peft/utils.py

@@ -1,10 +1,12 @@
 import math
 import re
+
 from importlib.metadata import version
 from typing import Any, Callable, List, Optional, Tuple, Type
 
 import packaging
 import torch
+
 from torch import nn
 
 from megatron.core import ModelParallelConfig, parallel_state
@@ -18,13 +20,14 @@
 
 try:
     from megatron.core.extensions.transformer_engine import (
+        TEColumnParallelGroupedLinear,
         TEColumnParallelLinear,
         TELayerNormColumnParallelLinear,
-        TEColumnParallelGroupedLinear,
-        TERowParallelLinear,
-        TERowParallelGroupedLinear,
         TELinear,
+        TERowParallelGroupedLinear,
+        TERowParallelLinear,
     )
+
     HAVE_TE = True
 except ImportError:
     HAVE_TE = False
@@ -34,8 +37,7 @@
 
 
 def match_module(m, name, prefix, target_modules):
-    """
-    """
+    """ """
     full_name = f"{prefix}.{name}" if prefix else name
     for pattern in target_modules:
         if name == pattern or wildcard_match(pattern, full_name):
@@ -117,7 +119,13 @@ def get_adapter_attributes_from_linear(m: nn.Module):
     else:
         raise NotImplementedError(f"Layer type is unrecognized for LoRA: {type(m)}")
 
-    return input_is_parallel, in_features, out_features, disable_sequence_parallel_comm, base_linear_is_parallel
+    return (
+        input_is_parallel,
+        in_features,
+        out_features,
+        disable_sequence_parallel_comm,
+        base_linear_is_parallel,
+    )
 
 
 def is_expert_linear(fqn):
@@ -188,7 +196,9 @@ def __init__(
         if model_parallel_config is None:
             model_parallel_config = ModelParallelConfig()
         _sequence_parallel = model_parallel_config.sequence_parallel
-        model_parallel_config.sequence_parallel = False  # SP is irrelevant for the lora linear layer
+        model_parallel_config.sequence_parallel = (
+            False  # SP is irrelevant for the lora linear layer
+        )
         self.config = model_parallel_config
 
         if input_is_parallel:
@@ -215,7 +225,7 @@ def __init__(
                 symmetric_ar_type=model_parallel_config.symmetric_ar_type,
                 skip_weight_param_allocation=False,
             )
-        
+
         if input_is_parallel:
             self.linear_out = TELinear(
                 input_size=dim,
@@ -246,7 +256,7 @@ def __init__(
             self.dropout = torch.nn.Dropout(p=dropout)
         else:
             self.dropout = None
-        
+
         # revert config change in case it is read elsewhere
         model_parallel_config.sequence_parallel = _sequence_parallel
         self.disable_sequence_parallel_comm = disable_sequence_parallel_comm
@@ -271,26 +281,34 @@ def forward(self, x):
         if self.dropout is not None and self.dropout_position == 'pre':
             x = self.dropout(x)
 
-        if not self.disable_sequence_parallel_comm and not self.input_is_parallel and not self.is_expert:
+        if (
+            not self.disable_sequence_parallel_comm
+            and not self.input_is_parallel
+            and not self.is_expert
+        ):
             # for attention_qkv and linear_fc1
             # layernorm before lora is impacted by sequence parallel,
             # hence seq dim need to be gathered right before lora linear layers
             # this function also handles the backward pass correctly
             x = gather_from_sequence_parallel_region(x)
-        
-        x, _= self.linear_in(x)
-        x, _= self.linear_out(x)
 
-        if not self.disable_sequence_parallel_comm and self.input_is_parallel and not self.is_expert:
+        x, _ = self.linear_in(x)
+        x, _ = self.linear_out(x)
+
+        if (
+            not self.disable_sequence_parallel_comm
+            and self.input_is_parallel
+            and not self.is_expert
+        ):
             # for attention_dense and linear_fc2
             # layernorm after lora is impacted by sequence parallel,
             # hence seq dim need to be scattered right after lora linear layers
             # this function also handles the backward pass correctly
             x = scatter_to_sequence_parallel_region(x)
-        
+
         if self.dropout is not None and self.dropout_position == 'post':
             x = self.dropout(x)
-        
+
         x = x * (self.alpha / self.dim)
 
         return x
@@ -303,17 +321,21 @@ def sharded_state_dict(
         since TP is sharded separately for the two logical matrices (gate and up)
         """
         sharded_state_dict = {}
-        linear_in_sd = self.linear_in.sharded_state_dict(f"{prefix}linear_in.", sharded_offsets, metadata)
-        linear_out_sd = self.linear_out.sharded_state_dict(f"{prefix}linear_out.", sharded_offsets, metadata)
+        linear_in_sd = self.linear_in.sharded_state_dict(
+            f"{prefix}linear_in.", sharded_offsets, metadata
+        )
+        linear_out_sd = self.linear_out.sharded_state_dict(
+            f"{prefix}linear_out.", sharded_offsets, metadata
+        )
 
         sharded_state_dict.update(linear_in_sd)
         sharded_state_dict.update(linear_out_sd)
         return sharded_state_dict
-        
+
     def __repr__(self):
         return (
             f"{type(self).__name__}(linear_in: in_features={self.linear_in.in_features}, "
             f"out_features={self.linear_in.out_features}, bias={self.linear_in.use_bias}, TP={self.linear_in.tp_size}), "
             f"(linear_out: in_features={self.linear_out.in_features}, "
             f"out_features={self.linear_out.out_features}, bias={self.linear_out.use_bias}, TP={self.linear_out.tp_size})"
-        )
+        )