from pytorch_lightning import LightningModule from pytorch_lightning.strategies import DeepSpeedStrategy from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam from transformers.optimization import AdamW, TYPE_TO_SCHEDULER_FUNCTION from torch.optim import Optimizer from torch.optim.lr_scheduler import _LRScheduler from transformers.trainer_utils import SchedulerType from typing import Optional, Union import warnings import types def add_module_args(parent_args): parser = parent_args.add_argument_group('Basic Module') parser.add_argument('--learning_rate', default=5e-5, type=float) parser.add_argument('--min_learning_rate', default=1e-7, type=float) parser.add_argument('--lr_decay_steps', default=0, type=int) # lr decay的时候会依赖total_steps，这里设置的是total_steps的比例，比如我只需要前50%步做decay，ratio设置为0.5 parser.add_argument('--lr_decay_ratio', default=1.0, type=float) parser.add_argument('--warmup_steps', default=0, type=int) parser.add_argument('--warmup_ratio', default=0.1, type=float) parser.add_argument('--weight_decay', default=1e-1, type=float) parser.add_argument('--adam_beta1', default=0.9, type=float) parser.add_argument('--adam_beta2', default=0.999, type=float) parser.add_argument('--adam_epsilon', default=1e-8, type=float) parser.add_argument('--model_path', default=None, type=str) parser.add_argument('--scheduler_type', default='polynomial', type=str) return parent_args def add_inverse_square_args(parent_args): parser = parent_args.add_argument_group('Basic Module') parser.add_argument('--warmup_min_lr', default=1e-9, type=float) parser.add_argument('--warmup_max_lr', default=1e-4, type=float) return parent_args def get_default_update_params(pl_model: LightningModule): no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'layer_norm.', 'layernorm.'] optimizer_grouped_params = [ {'params': [p for n, p in pl_model.named_parameters() if not any( nd in n for nd in no_decay) and p.requires_grad], 'weight_decay': pl_model.hparams.weight_decay}, {'params': [p for n, p in pl_model.named_parameters() if any( nd in n for nd in no_decay) and p.requires_grad], 'weight_decay': 0.0} ] return optimizer_grouped_params def configure_optimizers(pl_model: LightningModule, model_params=None): ''' Args: pl_model： lightning module model_params: 需要优化的模型参数 ''' # get params that optimizer need if model_params is None: optimizer_grouped_params = get_default_update_params(pl_model) else: optimizer_grouped_params = model_params # Configure optimizer. if isinstance(pl_model.trainer.strategy, DeepSpeedStrategy): if 'offload_optimizer' in pl_model.trainer.strategy.config['zero_optimization']: optimizer = DeepSpeedCPUAdam( optimizer_grouped_params, adamw_mode=True, lr=pl_model.hparams.learning_rate, betas=(pl_model.hparams.adam_beta1, pl_model.hparams.adam_beta2), eps=pl_model.hparams.adam_epsilon) else: optimizer = FusedAdam( optimizer_grouped_params, adam_w_mode=True, lr=pl_model.hparams.learning_rate, betas=(pl_model.hparams.adam_beta1, pl_model.hparams.adam_beta2), eps=pl_model.hparams.adam_epsilon) # elif isinstance(pl_model.trainer.strategy, ColossalAIStrategy): # from colossalai.nn.optimizer import HybridAdam # optimizer = HybridAdam( # optimizer_grouped_params, # lr=pl_model.hparams.learning_rate, # betas=(pl_model.hparams.adam_beta1, pl_model.hparams.adam_beta2), # eps=pl_model.hparams.adam_epsilon) else: optimizer = AdamW(optimizer_grouped_params, lr=pl_model.hparams.learning_rate, betas=(pl_model.hparams.adam_beta1, pl_model.hparams.adam_beta2), eps=pl_model.hparams.adam_epsilon) # Configure learning rate scheduler. warmup_steps = pl_model.hparams.warmup_ratio * \ pl_model.total_steps if pl_model.hparams.warmup_steps == 0 else pl_model.hparams.warmup_steps if pl_model.hparams.scheduler_type == "inverse_sqrt": scheduler = inverse_square_root_schedule(optimizer=optimizer, num_warmup_steps=warmup_steps, lr_min=pl_model.hparams.warmup_min_lr, lr_max=pl_model.hparams.warmup_max_lr) else: total_steps = pl_model.hparams.lr_decay_ratio * \ pl_model.total_steps if pl_model.hparams.lr_decay_steps == 0 else pl_model.hparams.lr_decay_steps scheduler = get_scheduler(name=pl_model.hparams.scheduler_type, optimizer=optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, lr_end=pl_model.hparams.min_learning_rate) scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1} return [optimizer], [scheduler] def inverse_square_root_schedule( optimizer: Optimizer, num_warmup_steps: int = 4000, lr_min=1e-9, lr_max=1e-4, power=0.5, last_epoch: int = -1): lr_init = optimizer.defaults["lr"] if (lr_min > lr_max): raise ValueError(f"lr_min ({lr_min}) must be be smaller than lr_max ({lr_max})") lr_step = (lr_max - lr_init) / num_warmup_steps decay_factor = lr_max * num_warmup_steps**power def lr_lambda(current_step: int): # 自定义函数 if current_step < num_warmup_steps: return lr_step * current_step return decay_factor * current_step ** (-power) return Direct_LR(optimizer, lr_lambda, last_epoch, True) class Direct_LR(_LRScheduler): """ Modified from LambdaLR """ def __init__(self, optimizer, lr_lambda, last_epoch=-1, warmup_steps=4000, verbose=False): self.optimizer = optimizer self.warmup_steps = warmup_steps if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple): self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups) else: if len(lr_lambda) != len(optimizer.param_groups): raise ValueError("Expected {} lr_lambdas, but got {}".format( len(optimizer.param_groups), len(lr_lambda))) self.lr_lambdas = list(lr_lambda) super(Direct_LR, self).__init__(optimizer, last_epoch, verbose) def state_dict(self): """Returns the state of the scheduler as a :class:`dict`. It contains an entry for every variable in self.__dict__ which is not the optimizer. The learning rate lambda functions will only be saved if they are callable objects and not if they are functions or lambdas. When saving or loading the scheduler, please make sure to also save or load the state of the optimizer. """ state_dict = {key: value for key, value in self.__dict__.items() if key not in ('optimizer', 'lr_lambdas')} state_dict['lr_lambdas'] = [None] * len(self.lr_lambdas) for idx, fn in enumerate(self.lr_lambdas): if not isinstance(fn, types.FunctionType): state_dict['lr_lambdas'][idx] = fn.__dict__.copy() return state_dict def load_state_dict(self, state_dict): """Loads the schedulers state. When saving or loading the scheduler, please make sure to also save or load the state of the optimizer. Args: state_dict (dict): scheduler state. Should be an object returned from a call to :meth:`state_dict`. """ lr_lambdas = state_dict.pop('lr_lambdas') self.__dict__.update(state_dict) # Restore state_dict keys in order to prevent side effects # https://github.com/pytorch/pytorch/issues/32756 state_dict['lr_lambdas'] = lr_lambdas for idx, fn in enumerate(lr_lambdas): if fn is not None: self.lr_lambdas[idx].__dict__.update(fn) def get_lr(self): if not self._get_lr_called_within_step: warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.") if self._step_count < self.warmup_steps: return [base_lr + lmbda(self.last_epoch) for lmbda, base_lr in zip(self.lr_lambdas, self.base_lrs)] return [lmbda(self.last_epoch) for lmbda in self.lr_lambdas] def get_total_steps(trainer, hparams): train_loader = trainer._data_connector._train_dataloader_source.dataloader() # Calculate total steps if trainer.max_epochs > 0: world_size = trainer.world_size tb_size = hparams.train_batchsize * max(1, world_size) ab_size = trainer.accumulate_grad_batches total_steps = (len(train_loader.dataset) * trainer.max_epochs // tb_size) // ab_size else: total_steps = trainer.max_steps return total_steps def get_scheduler( name: Union[str, SchedulerType], optimizer: Optimizer, num_warmup_steps: Optional[int] = None, num_training_steps: Optional[int] = None, lr_end: Optional[float] = None ): """ Unified API to get any scheduler from its name. Args: name (`str` or `SchedulerType`): The name of the scheduler to use. optimizer (`torch.optim.Optimizer`): The optimizer that will be used during training. num_warmup_steps (`int`, *optional*): The number of warmup steps to do. This is not required by all schedulers (hence the argument being optional), the function will raise an error if it's unset and the scheduler type requires it. num_training_steps (`int``, *optional*): The number of training steps to do. This is not required by all schedulers (hence the argument being optional), the function will raise an error if it's unset and the scheduler type requires it. """ name = SchedulerType(name) schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] if name == SchedulerType.CONSTANT: return schedule_func(optimizer) # All other schedulers require `num_warmup_steps` if num_warmup_steps is None: raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") if name == SchedulerType.CONSTANT_WITH_WARMUP: return schedule_func(optimizer, num_warmup_steps=num_warmup_steps) # All other schedulers require `num_training_steps` if num_training_steps is None: raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") if name == SchedulerType.POLYNOMIAL: return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, lr_end=lr_end) return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)