Repository: huggingface/transformers
Pull request: [`fix`] Pass adamw optimizer parameters to StableAdamW

# What does this PR do?

I was getting the following warning when training with `stable_adamw` and `decouple_lr=True`, even though the `TrainingArguments.weight_decay` was set to a different value.
```
optimi/optimizer.py:41: UserWarning: You are using weight_decay=0.01 which is potentially high for decouple_lr=True. Unlike decoupled weight decay, fully decoupled weight decay does not reduce weight decay by the learning rate.
```

I updated the `stable_adamw_kwargs` to include adamw-related TrainingArguments parameters.

https://optimi.benjaminwarner.dev/optimizers/stableadamw/
https://github.com/warner-benjamin/optimi/blob/main/optimi/stableadamw.py#L94

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@SunMarc @qgallouedec


Changed files:
filename: src/transformers/trainer.py
patch: @@ -1819,6 +1819,7 @@ def optimizer_hook(param):
             if kahan_sum is not None:
                 kahan_sum = bool(kahan_sum)
 
+            adam_kwargs["weight_decay"] = args.weight_decay
             stable_adamw_kwargs = {
                 "decouple_lr": bool(optim_args.pop("decouple_lr", False)),
                 "max_lr": max_lr,

filename: tests/trainer/test_trainer.py
patch: @@ -2569,6 +2569,38 @@ def test_stable_adamw_lr_display_without_scheduler(self):
         # reflects displayed lr in trainer
         self.assertEqual(trainer.get_learning_rates(), [learning_rate, learning_rate])
 
+    @require_torch_optimi
+    @require_torch_accelerator
+    def test_stable_adamw_trainer_adamw_args(self):
+        config = LlamaConfig(vocab_size=100, hidden_size=32, num_hidden_layers=3, num_attention_heads=4)
+        tiny_llama = LlamaForCausalLM(config)
+        x = torch.randint(0, 100, (128,))
+        train_dataset = RepeatDataset(x)
+
+        learning_rate = 1e-9
+        num_steps = 10
+
+        # Trainer without inf/nan filter
+        args = TrainingArguments(
+            self.get_auto_remove_tmp_dir(),
+            learning_rate=learning_rate,
+            logging_steps=5,
+            weight_decay=0.001,
+            adam_beta1=0.89,
+            adam_beta2=0.98,
+            adam_epsilon=1e-8,
+            optim="stable_adamw",
+            optim_target_modules=[r".*attn.*", r".*mlp.*"],
+        )
+        trainer = Trainer(tiny_llama, args, train_dataset=train_dataset)
+        trainer.create_optimizer_and_scheduler(num_training_steps=num_steps)
+
+        # check StableAdamW optimizer is created with the correct parameters
+        self.assertEqual(trainer.optimizer.defaults["beta1"], args.adam_beta1)
+        self.assertEqual(trainer.optimizer.defaults["beta2"], args.adam_beta2)
+        self.assertEqual(trainer.optimizer.defaults["eps"], args.adam_epsilon)
+        self.assertEqual(trainer.optimizer.defaults["weight_decay"], args.weight_decay)
+
     @require_torch_optimi
     @require_torch_accelerator
     def test_stable_adamw_lr_display_with_scheduler(self):