wehub-resource-sync
75 lines
3.0 KiB
Python
75 lines
3.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch.nn as nn
|
|
from peft import PeftModel
|
|
from transformers import Trainer as HfTrainer
|
|
from typing import List, Optional, Tuple
|
|
|
|
from swift.utils import get_logger
|
|
from .base import OptimizerCallback
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def get_param_startswith(model,
|
|
chosen_prefix: List[str],
|
|
rejected_prefix: Optional[List[str]] = None) -> List[Tuple[str, nn.Parameter]]:
|
|
chosen_prefix = chosen_prefix or []
|
|
rejected_prefix = rejected_prefix or []
|
|
res = []
|
|
if not chosen_prefix:
|
|
return res
|
|
is_peft_model = isinstance(model, PeftModel)
|
|
if is_peft_model:
|
|
model = model.model
|
|
for n, p in model.named_parameters():
|
|
if not p.requires_grad:
|
|
continue
|
|
is_rejected = False
|
|
for prefix in rejected_prefix:
|
|
if n.startswith(prefix):
|
|
is_rejected = True
|
|
break
|
|
if is_rejected:
|
|
continue
|
|
for prefix in chosen_prefix:
|
|
if n.startswith(prefix):
|
|
if is_peft_model:
|
|
n = f'base_model.model.{n}'
|
|
res.append((n, p))
|
|
break
|
|
return res
|
|
|
|
|
|
class MultimodalOptimizerCallback(OptimizerCallback):
|
|
|
|
def create_optimizer(self, model=None):
|
|
"""ViT/Aligner/LLM use different learning rates."""
|
|
args = self.args
|
|
if model is None:
|
|
model = self.trainer.model
|
|
decay_parameters = set(HfTrainer.get_decay_parameter_names(None, model))
|
|
model_arch = model.model_meta.model_arch
|
|
vit_parameters = get_param_startswith(model, model_arch.vision_tower, model_arch.aligner)
|
|
aligner_parameters = get_param_startswith(model, model_arch.aligner)
|
|
llm_parameters = get_param_startswith(model, model_arch.language_model)
|
|
optimizer_grouped_parameters = []
|
|
vit_lr = args.vit_lr if args.vit_lr is not None else args.learning_rate
|
|
aligner_lr = args.aligner_lr if args.aligner_lr is not None else args.learning_rate
|
|
logger.info(f'vit_lr: {vit_lr}, aligner_lr: {aligner_lr}, llm_lr: {args.learning_rate}')
|
|
for lr, parameters in zip([vit_lr, aligner_lr, args.learning_rate],
|
|
[vit_parameters, aligner_parameters, llm_parameters]):
|
|
for use_wd, wd in zip([False, True], [0., args.weight_decay]):
|
|
if use_wd:
|
|
params = [p for n, p in parameters if n in decay_parameters]
|
|
else:
|
|
params = [p for n, p in parameters if n not in decay_parameters]
|
|
if not params:
|
|
continue
|
|
optimizer_grouped_parameters.append({
|
|
'params': params,
|
|
'weight_decay': wd,
|
|
'lr': lr,
|
|
})
|
|
optimizer_cls, optimizer_kwargs = HfTrainer.get_optimizer_cls_and_kwargs(args, model)
|
|
return optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
|