Fix galore linears (#541)

Author: tastelikefeet

swift/llm/tuner.py

@@ -179,20 +179,18 @@ def prepare_model(model, args: SftArguments):
 
     if args.use_galore:
         from swift.trainers.optimizers.galore import GaLoreConfig
-        model_type = args.model_type
-        for key in MODEL_KEYS_MAPPING.keys():
-            if key in model_type.lower():
-                model_type = key
-                break
+        if args.galore_target_modules is None:
+            args.galore_target_modules = find_all_linears(
+                model, 0, args.model_type)
+        if args.galore_with_embedding:
+            args.galore_target_modules += find_embedding(model)
         args.training_args.galore_config = GaLoreConfig(
-            model_type=model_type,
             target_modules=args.galore_target_modules,
             rank=args.galore_rank,
             update_proj_gap=args.galore_update_proj_gap,
             galore_scale=args.galore_scale,
             proj_type=args.galore_proj_type,
             optim_per_parameter=args.galore_optim_per_parameter,
-            with_embedding=args.galore_with_embedding,
         )
 
     class TrainerAdapterCallback(TrainerCallback):

swift/llm/utils/argument.py

@@ -125,7 +125,7 @@ class SftArguments:
     # galore
     use_galore: bool = False
     galore_rank: int = 128
-    galore_target_modules: Union[str, List[str]] = None
+    galore_target_modules: Optional[List[str]] = None
     galore_update_proj_gap: int = 50
     galore_scale: float = 1.0
     galore_proj_type: str = 'std'

swift/trainers/optimizers/galore/utils.py

@@ -8,7 +8,6 @@
 from torch.optim.lr_scheduler import LRScheduler
 from transformers import Trainer, TrainingArguments, get_scheduler
 
-from swift.tuners.module_mapping import MODEL_KEYS_MAPPING
 from swift.utils import get_logger
 
 logger = get_logger()
@@ -23,7 +22,6 @@ class GaLoreConfig:
     See https://arxiv.org/abs/2403.03507
 
     Args:
-        model_type (`str`): The model_type of Galore
         rank (`int`): The galore rank
         target_modules (`Union[str, List[str]]`): The target modules to use, if `None`,
             will use all attn and mlp linears
@@ -33,13 +31,11 @@ class GaLoreConfig:
         galore_scale(float): the scale of gradient
         optim_per_parameter(bool): Gives one optimizer per parameter
     """
-    model_type: str = None
     rank: int = 128
     target_modules: Union[str, List[str]] = None
     update_proj_gap: int = 50
     galore_scale: float = 1.0
     proj_type: str = 'std'
-    with_embedding: bool = False
     optim_per_parameter: bool = False
 
 
@@ -72,19 +68,6 @@ def step(self, *args, **kwargs) -> None:
 def create_optimizer_and_scheduler(model: nn.Module, args: TrainingArguments,
                                    config: GaLoreConfig, max_steps,
                                    **defaults):
-    if not config.target_modules:
-        if config.model_type in MODEL_KEYS_MAPPING:
-            target_modules_list = [
-                MODEL_KEYS_MAPPING[config.model_type].attention.split('.{}.')
-                [1], MODEL_KEYS_MAPPING[config.model_type].mlp.split('.{}.')[1]
-            ]
-            config.target_modules = target_modules_list
-            if config.with_embedding:
-                embedding = MODEL_KEYS_MAPPING[config.model_type].embedding
-                idx = embedding.rfind('.')
-                embedding = embedding[idx + 1:]
-                target_modules_list.append(embedding)
-
     galore_params = []
     for module_name, module in model.named_modules():
         if not isinstance(module, (nn.Linear, nn.Embedding)) or \