From 3b5a9d1d8887542349db15c0bfcbe83a8b9a746f Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 19 Feb 2026 23:49:32 -0500 Subject: [PATCH] update create_optimizer for updated api --- src/axolotl/core/trainers/mixins/optimizer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/axolotl/core/trainers/mixins/optimizer.py b/src/axolotl/core/trainers/mixins/optimizer.py index 850442c60..dc011d2b1 100644 --- a/src/axolotl/core/trainers/mixins/optimizer.py +++ b/src/axolotl/core/trainers/mixins/optimizer.py @@ -104,7 +104,7 @@ class OptimizerMixin(Trainer): return optimizer_grouped_parameters - def create_optimizer(self): + def create_optimizer(self, model=None): if ( self.args.loraplus_lr_ratio is None and self.args.embedding_lr_scale is None @@ -112,9 +112,9 @@ class OptimizerMixin(Trainer): and self.args.lr_groups is None and self.optimizer_cls_and_kwargs is None ): - return super().create_optimizer() + return super().create_optimizer(model=model) - opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model + opt_model = self.model if model is None else model if ( not self.optimizer