chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,34 @@
+from torch.optim import Optimizer
+from transformers.trainer import Trainer as HfTrainer
+
+from .base import OptimizerCallback
+
+
+class LorapOptimizerCallback(OptimizerCallback):
+
+    def create_optimizer(self, model=None) -> Optimizer:
+        args = self.args
+        if model is None:
+            model = self.trainer.model
+        optimizer_grouped_parameters = None
+        if hasattr(model, 'create_optimizer_param_groups'):
+            # Lora+ parameter groups
+            optimizer_grouped_parameters = model.create_optimizer_param_groups(
+                lr=args.learning_rate, weight_decay=args.weight_decay)
+
+        if optimizer_grouped_parameters is None:
+            # Default parameter groups
+            decay_parameters = HfTrainer.get_decay_parameter_names(None, model)
+            optimizer_grouped_parameters = [
+                {
+                    'params': [p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad)],
+                    'weight_decay': args.weight_decay,
+                },
+                {
+                    'params':
+                    [p for n, p in model.named_parameters() if (n not in decay_parameters and p.requires_grad)],
+                    'weight_decay': 0.0,
+                },
+            ]
+        optimizer_cls, optimizer_kwargs = HfTrainer.get_optimizer_cls_and_kwargs(args)
+        return optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)