Repository: huggingface/transformers (Python)
Task: Bug Fix (2024-05-30 03:10:04+00:00)

Issue:
get_wsd_schedule gets passed num_training_steps because not handled
getting:
```
TypeError: get_wsd_schedule() got an unexpected keyword argument 'num_training_steps'
```
because there's not a handling of ```WARMUP_STABLE_DECAY```, get_wsd_schedule gets passed default params.

https://github.com/huggingface/transformers/blob/a3c7b59e31710f76492146db1d87ba77e3ae265d/src/transformers/optimization.py#L547


Discussion:
cc @muellerzr 

Patch:
diff --git a/src/transformers/optimization.py b/src/transformers/optimization.py
--- a/src/transformers/optimization.py
+++ b/src/transformers/optimization.py
@@ -540,6 +540,9 @@ def scheduler_hook(param):
     if name == SchedulerType.INVERSE_SQRT:
         return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
 
+    if name == SchedulerType.WARMUP_STABLE_DECAY:
+        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, **scheduler_specific_kwargs)
+
     # 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.")


Test patch:
diff --git a/tests/optimization/test_optimization.py b/tests/optimization/test_optimization.py
--- a/tests/optimization/test_optimization.py
+++ b/tests/optimization/test_optimization.py
@@ -36,6 +36,7 @@
         get_inverse_sqrt_schedule,
         get_linear_schedule_with_warmup,
         get_polynomial_decay_schedule_with_warmup,
+        get_scheduler,
         get_wsd_schedule,
     )
 
@@ -176,6 +177,27 @@ def test_schedulers(self):
             lrs_2 = unwrap_and_save_reload_schedule(scheduler, self.num_steps)
             self.assertListEqual(lrs_1, lrs_2, msg=f"failed for {scheduler_func} in save and reload")
 
+    def test_get_scheduler(self):
+        test_params = [
+            {
+                "name": "warmup_stable_decay",
+                "optimizer": self.optimizer,
+                "num_warmup_steps": 2,
+                "scheduler_specific_kwargs": {"num_stable_steps": 1, "num_decay_steps": 3},
+            },
+            {
+                "name": "warmup_stable_decay",
+                "optimizer": self.optimizer,
+                "num_warmup_steps": 2,
+                "num_training_steps": 10,
+                "scheduler_specific_kwargs": {"num_stable_steps": 1, "num_decay_steps": 3},
+            },
+            {"name": "cosine", "optimizer": self.optimizer, "num_warmup_steps": 2, "num_training_steps": 10},
+        ]
+
+        for param in test_params:
+            self.assertTrue(get_scheduler(**param), msg=f"failed for {param['name']} in get_scheduler")
+
 
 class LambdaScheduleWrapper:
     """See https://github.com/huggingface/transformers/issues/21689"""