fix(lisa): select language layers in multimodal models

swift/callbacks/lisa.py

@@ -1,6 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import numpy as np
 import torch
+from accelerate.utils import extract_model_from_parallel
 from typing import TYPE_CHECKING
 
 from .base import TrainerCallback
@@ -19,10 +20,21 @@ def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
 
     def on_train_begin(self, args, state, control, **kwargs):
         # Wait until the model exists and the optimizer has registered all layers.
-        self.model = kwargs['model']
+        self.model = extract_model_from_parallel(kwargs['model'], keep_torch_compile=False)
+        model_arch = getattr(getattr(self.model, 'model_meta', None), 'model_arch', None)
+        layer_prefixes = getattr(model_arch, 'language_model', None)
+        if getattr(model_arch, 'module_list', None):
+            layer_prefixes = [model_arch.module_list]
+        excluded_prefixes = [
+            prefix for key in ('vision_tower', 'aligner', 'generator') for prefix in getattr(model_arch, key, [])
+        ]
         layers_name = None
         layers = None
         for name, module in self.model.named_modules():
+            if layer_prefixes and not any(name == prefix or name.startswith(f'{prefix}.') for prefix in layer_prefixes):
+                continue
+            if any(name == prefix or name.startswith(f'{prefix}.') for prefix in excluded_prefixes):
+                continue
             if isinstance(module, torch.nn.ModuleList):
                 layers_name = name
                 layers = module