diff --git a/swift/callbacks/lisa.py b/swift/callbacks/lisa.py index e44896336a..e16f8a80e4 100644 --- a/swift/callbacks/lisa.py +++ b/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