transformers/tests/trainer/test_trainer_optimizers.py

# Copyright 2018 the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Trainer optimizer and LR scheduler tests: custom optimizers, LR scheduler kwargs, cosine-with-min-lr,
reduce-on-plateau, Adafactor, bitsandbytes (RMSProp, AdEMAMix), LOMO, GrokAdamW, schedule-free,
GaLore, Apollo, Stable AdamW, Liger kernel, optimizer choice resolution, factory pattern detection,
and model parameter inspection.
"""

import tempfile

import numpy as np
from parameterized import parameterized

from transformers import (
    GPT2Config,
    GPT2LMHeadModel,
    LlamaConfig,
    LlamaForCausalLM,
    Trainer,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import (
    TestCasePlus,
    require_apollo_torch,
    require_bitsandbytes,
    require_galore_torch,
    require_grokadamw,
    require_lomo,
    require_schedulefree,
    require_torch,
    require_torch_accelerator,
    require_torch_optimi,
)
from transformers.trainer_utils import check_target_module_exists

from .trainer_test_utils import (
    BasicTextGenerationModel,
    RegressionDataset,
    RegressionModel,
    RepeatDataset,
    TorchTracemalloc,
    TrainerIntegrationCommon,
    TstLayer,
    bytes2megabytes,
    get_regression_trainer,
)


if is_torch_available():
    import torch
    from torch import nn

_ATTN_MLP_TARGET_MODULES = [r".*attn.*", r".*mlp.*"]


@require_torch
class TrainerOptimizerIntegrationTest(TestCasePlus, TrainerIntegrationCommon):
    def setUp(self):
        super().setUp()
        args = TrainingArguments("..")
        self.n_epochs = args.num_train_epochs
        self.batch_size = args.train_batch_size

    # ---------------------------------------------------------------------------
    # Helpers
    # ---------------------------------------------------------------------------

    def _get_llama_and_dataset(self):
        config = LlamaConfig(vocab_size=100, hidden_size=32, num_hidden_layers=3, num_attention_heads=4)
        model = LlamaForCausalLM(config)
        train_dataset = RepeatDataset(torch.randint(0, 100, (128,)))
        return model, train_dataset

    def _get_gpt2_and_dataset(self):
        config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4)
        model = GPT2LMHeadModel(config)
        train_dataset = RepeatDataset(torch.randint(0, 100, (128,)))
        return model, train_dataset

    def _train_with_llama(self, optim, optim_target_modules=None, **extra_kwargs):
        """Smoke-test: tiny Llama + RepeatDataset with the given optimizer."""
        tiny_llama, train_dataset = self._get_llama_and_dataset()
        kwargs = {"learning_rate": 1e-9, "logging_steps": 5, "optim": optim}
        if optim_target_modules is not None:
            kwargs["optim_target_modules"] = optim_target_modules
        kwargs.update(extra_kwargs)
        args = TrainingArguments(self.get_auto_remove_tmp_dir(), **kwargs)
        trainer = Trainer(tiny_llama, args, train_dataset=train_dataset)
        trainer.train()
        return trainer

    def _check_lr_display_without_scheduler(self, optim, optim_target_modules):
        """Verify that LR is correctly reported without an LR scheduler."""
        tiny_llama, train_dataset = self._get_llama_and_dataset()
        learning_rate = 1e-9
        args = TrainingArguments(
            self.get_auto_remove_tmp_dir(),
            learning_rate=learning_rate,
            logging_steps=5,
            optim=optim,
            optim_target_modules=optim_target_modules,
        )
        trainer = Trainer(tiny_llama, args, train_dataset=train_dataset)
        trainer.create_optimizer_and_scheduler(num_training_steps=10)
        self.assertEqual(trainer.get_learning_rates(), [learning_rate, learning_rate])

    def _check_lr_display_with_scheduler(self, optim, optim_target_modules, num_train_epochs=2):
        """Verify warmup + cosine LR schedule: increases then decreases."""
        tiny_llama, train_dataset = self._get_llama_and_dataset()
        learning_rate = 2e-4
        num_warmup_steps = 5
        args = TrainingArguments(
            self.get_auto_remove_tmp_dir(),
            num_train_epochs=num_train_epochs,
            learning_rate=learning_rate,
            warmup_steps=num_warmup_steps,
            lr_scheduler_type="cosine",
            logging_steps=1,
            optim=optim,
            optim_target_modules=optim_target_modules,
        )
        trainer = Trainer(tiny_llama, args, train_dataset=train_dataset)
        trainer.train()
        logs = trainer.state.log_history[1:-1]

        self.assertTrue(logs[num_warmup_steps - 1]["learning_rate"] == learning_rate)
        self.assertTrue(np.allclose(logs[-1]["learning_rate"], 0, atol=5e-6))

        increasing_lrs = [
            logs[i]["learning_rate"] < logs[i + 1]["learning_rate"]
            for i in range(len(logs))
            if i < num_warmup_steps - 1
        ]
        decreasing_lrs = [
            logs[i]["learning_rate"] > logs[i + 1]["learning_rate"]
            for i in range(len(logs) - 1)
            if i >= num_warmup_steps - 1
        ]

        self.assertTrue(all(increasing_lrs))
        self.assertTrue(all(decreasing_lrs))
        self.assertTrue(len(decreasing_lrs) > len(increasing_lrs))

    # ---------------------------------------------------------------------------
    # adafactor optmizer test
    # ---------------------------------------------------------------------------

    def test_adafactor_lr_none(self):
        # test the special case where lr=None, since Trainer can't not have lr_scheduler

        from transformers.optimization import Adafactor, AdafactorSchedule

        train_dataset = RegressionDataset()
        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(tmp_dir)
            model = RegressionModel()
            optimizer = Adafactor(
                model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None
            )
            lr_scheduler = AdafactorSchedule(optimizer)
            trainer = Trainer(model, args, train_dataset=train_dataset, optimizers=(optimizer, lr_scheduler))
            trainer.train()

            # Train a default model to compare against
            default_trainer = get_regression_trainer(learning_rate=0.1, output_dir=tmp_dir)
            default_trainer.train()

            self.assertFalse(torch.allclose(trainer.model.a, default_trainer.model.a))
            self.assertFalse(torch.allclose(trainer.model.b, default_trainer.model.b))
            self.assertGreater(trainer.optimizer.state_dict()["param_groups"][0]["lr"], 0)

    # ---------------------------------------------------------------------------
    # BNB optimizer tests
    # ---------------------------------------------------------------------------

    @parameterized.expand(["rmsprop_bnb", "ademamix", "ademamix_8bit", "rmsprop_bnb_8bit", "rmsprop_bnb_32bit"])
    @require_bitsandbytes
    def test_bnb_optim(self, optim):
        tiny_gpt2, train_dataset = self._get_gpt2_and_dataset()
        args = TrainingArguments(
            self.get_auto_remove_tmp_dir(),
            learning_rate=1e-9,
            logging_steps=5,
            logging_nan_inf_filter=False,
            optim=optim,
        )
        Trainer(tiny_gpt2, args, train_dataset=train_dataset).train()

    @require_bitsandbytes
    def test_bnb_8bit_optimizer_skip_embedding(self):
        model = BasicTextGenerationModel(8, 4)
        with tempfile.TemporaryDirectory() as tmp_dir:
            for name_optim in ["rmsprop_bnb_8bit", "adamw_8bit"]:
                args = TrainingArguments(
                    output_dir=tmp_dir,
                    optim=name_optim,
                )
                trainer = Trainer(model=model, args=args)
                optimizer = trainer.create_optimizer()
                modules = optimizer.mng.module_weight_config_triple
                self.assertNotEqual(len(modules), 0)
                module, name, config = modules[0]
                self.assertIsInstance(module, torch.nn.Embedding)
                self.assertEqual(name, "weight")
                self.assertDictEqual(config, {"optim_bits": 32})

    # ---------------------------------------------------------------------------
    # LOMO tests
    # ---------------------------------------------------------------------------

    @require_lomo
    @require_torch_accelerator
    def test_lomo(self):
        tiny_llama, train_dataset = self._get_llama_and_dataset()
        previous_params = {n: p.clone() for n, p in tiny_llama.named_parameters()}

        args = TrainingArguments(
            self.get_auto_remove_tmp_dir(), learning_rate=1e-2, logging_steps=5, optim="lomo", max_steps=20
        )
        Trainer(tiny_llama, args, train_dataset=train_dataset).train()

        for name, param in tiny_llama.named_parameters():
            self.assertFalse(torch.allclose(param, previous_params[name].to(param.device), rtol=1e-12, atol=1e-12))

    @require_lomo
    @require_torch_accelerator
    def test_adalomo(self):
        self._train_with_llama("adalomo")

    # ---------------------------------------------------------------------------
    # GrokAdamW test
    # ---------------------------------------------------------------------------

    @require_grokadamw
    @require_torch_accelerator
    def test_grokadamw(self):
        self._train_with_llama("grokadamw", learning_rate=2e-5, max_steps=20)

    # ---------------------------------------------------------------------------
    # Schedule-free tests
    # ---------------------------------------------------------------------------

    @parameterized.expand([("schedule_free_adamw",), ("schedule_free_radam",)])
    @require_schedulefree
    @require_torch_accelerator
    def test_schedulefree(self, optim):
        self._train_with_llama(optim, lr_scheduler_type="constant")

    # ---------------------------------------------------------------------------
    # GaLore tests
    # ---------------------------------------------------------------------------

    def test_galore_matched_modules(self):
        regex_patterns = [r".*.attn.*", r".*.mlp.*"]

        module_names = [
            "model.transformer.h.0.ln_1",
            "model.transformer.h.0.attn.q_proj",
            "model.lm_head",
            "model.transformer.h.0.mlp.up_proj",
        ]
        expected_values = [False, True, False, True]

        for expected_value, module_name in zip(expected_values, module_names):
            is_module_matched, is_regex = check_target_module_exists(regex_patterns, module_name, return_is_regex=True)
            self.assertTrue(is_module_matched == expected_value)
            if is_module_matched:
                self.assertTrue(is_regex)

        exact_patterns = ["q_proj", "up_proj"]

        module_names = [
            "model.transformer.h.0.ln_1",
            "model.transformer.h.0.attn.q_proj",
            "model.lm_head",
            "model.transformer.h.0.mlp.up_proj",
        ]
        expected_values = [False, True, False, True]

        for expected_value, module_name in zip(expected_values, module_names):
            is_module_matched, is_regex = check_target_module_exists(exact_patterns, module_name, return_is_regex=True)
            self.assertTrue(is_module_matched == expected_value)
            if is_module_matched:
                self.assertFalse(is_regex)

        simple_regex = r".*.attn.*"

        module_names = [
            "model.transformer.h.0.ln_1",
            "model.transformer.h.0.attn.q_proj",
            "model.lm_head",
            "model.transformer.h.0.mlp.up_proj",
        ]
        expected_values = [False, True, False, False]

        for expected_value, module_name in zip(expected_values, module_names):
            is_module_matched, is_regex = check_target_module_exists(simple_regex, module_name, return_is_regex=True)
            self.assertTrue(is_module_matched == expected_value)
            if is_module_matched:
                self.assertTrue(is_regex)

        simple_regex = "model.transformer.h.0.attn.q_proj"

        module_names = [
            "model.transformer.h.0.ln_1",
            "model.transformer.h.0.attn.q_proj",
            "model.lm_head",
            "model.transformer.h.0.mlp.up_proj",
        ]
        expected_values = [False, True, False, False]

        for expected_value, module_name in zip(expected_values, module_names):
            is_module_matched, is_regex = check_target_module_exists(simple_regex, module_name, return_is_regex=True)
            self.assertTrue(is_module_matched == expected_value)
            if is_module_matched:
                self.assertFalse(is_regex)

        target_modules = ["attn", "mlp"]

        module_names = [
            "model.transformer.h.0.ln_1",
            "model.transformer.h.0.attn.q_proj",
            "model.lm_head",
            "model.transformer.h.0.mlp.up_proj",
        ]
        expected_values = [False, True, False, True]

        for expected_value, module_name in zip(expected_values, module_names):
            is_module_matched, is_regex = check_target_module_exists(target_modules, module_name, return_is_regex=True)
            self.assertTrue(is_module_matched == expected_value)
            if is_module_matched:
                self.assertFalse(is_regex)

    @parameterized.expand([("galore_adamw",), ("galore_adamw_layerwise",), ("galore_adamw_8bit",)])
    @require_galore_torch
    @require_torch_accelerator
    def test_galore(self, optim):
        self._train_with_llama(optim, optim_target_modules=_ATTN_MLP_TARGET_MODULES)

    @require_galore_torch
    @require_torch_accelerator
    def test_galore_extra_args(self):
        self._train_with_llama(
            "galore_adamw",
            optim_target_modules=_ATTN_MLP_TARGET_MODULES,
            optim_args="rank=64, update_proj_gap=100, scale=0.10",
        )

    @require_galore_torch
    @require_torch_accelerator
    def test_galore_layerwise_with_scheduler(self):
        self._train_with_llama(
            "galore_adamw_layerwise",
            optim_target_modules=_ATTN_MLP_TARGET_MODULES,
            lr_scheduler_type="cosine",
        )

    @parameterized.expand(
        [
            (_ATTN_MLP_TARGET_MODULES,),
            (["q_proj", "k_proj", "v_proj"],),
            ("all-linear",),
        ]
    )
    @require_galore_torch
    @require_torch_accelerator
    def test_galore_adafactor(self, optim_target_modules):
        upper_bound_pm = 700
        lower_bound_pm = 650
        tiny_llama, train_dataset = self._get_llama_and_dataset()

        with tempfile.TemporaryDirectory() as tmpdir, TorchTracemalloc() as tracemalloc:
            args = TrainingArguments(
                tmpdir,
                learning_rate=1e-9,
                logging_steps=5,
                optim="galore_adafactor",
                optim_target_modules=optim_target_modules,
            )
            Trainer(tiny_llama, args, train_dataset=train_dataset).train()

        galore_peak_memory = tracemalloc.peaked + bytes2megabytes(tracemalloc.begin)
        self.assertTrue(galore_peak_memory < upper_bound_pm)
        self.assertTrue(lower_bound_pm < galore_peak_memory)

    @require_galore_torch
    @require_torch_accelerator
    def test_galore_lr_display_without_scheduler(self):
        self._check_lr_display_without_scheduler("galore_adamw", _ATTN_MLP_TARGET_MODULES)

    @require_galore_torch
    @require_torch_accelerator
    def test_galore_lr_display_with_scheduler(self):
        self._check_lr_display_with_scheduler("galore_adamw", _ATTN_MLP_TARGET_MODULES)

    # ---------------------------------------------------------------------------
    # Apollo tests
    # ---------------------------------------------------------------------------

    @parameterized.expand([("apollo_adamw",), ("apollo_adamw_layerwise",)])
    @require_apollo_torch
    @require_torch_accelerator
    def test_apollo(self, optim):
        self._train_with_llama(optim, optim_target_modules=_ATTN_MLP_TARGET_MODULES)

    @require_apollo_torch
    @require_torch_accelerator
    def test_apollo_extra_args(self):
        self._train_with_llama(
            "apollo_adamw",
            optim_target_modules=_ATTN_MLP_TARGET_MODULES,
            optim_args="proj=random,scale_type=tensor,rank=1,update_proj_gap=100,scale=128.0",
        )

    @require_apollo_torch
    @require_torch_accelerator
    def test_apollo_layerwise_with_scheduler(self):
        self._train_with_llama(
            "apollo_adamw_layerwise",
            optim_target_modules=_ATTN_MLP_TARGET_MODULES,
            lr_scheduler_type="cosine",
        )

    @require_apollo_torch
    @require_torch_accelerator
    def test_apollo_lr_display_without_scheduler(self):
        self._check_lr_display_without_scheduler("apollo_adamw", _ATTN_MLP_TARGET_MODULES)

    @require_apollo_torch
    @require_torch_accelerator
    def test_apollo_lr_display_with_scheduler(self):
        self._check_lr_display_with_scheduler("apollo_adamw", _ATTN_MLP_TARGET_MODULES, num_train_epochs=10)

    # ---------------------------------------------------------------------------
    # Stable AdamW tests
    # ---------------------------------------------------------------------------

    @require_torch_optimi
    @require_torch_accelerator
    def test_stable_adamw(self):
        self._train_with_llama("stable_adamw", optim_target_modules=_ATTN_MLP_TARGET_MODULES)

    @require_torch_optimi
    @require_torch_accelerator
    def test_stable_adamw_extra_args(self):
        self._train_with_llama(
            "stable_adamw",
            optim_target_modules=_ATTN_MLP_TARGET_MODULES,
            optim_args="decouple_lr=True,max_lr=1e-3,kahan_sum=True",
        )

    @require_torch_optimi
    @require_torch_accelerator
    def test_stable_adamw_trainer_adamw_args(self):
        tiny_llama, train_dataset = self._get_llama_and_dataset()
        args = TrainingArguments(
            self.get_auto_remove_tmp_dir(),
            learning_rate=1e-9,
            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=_ATTN_MLP_TARGET_MODULES,
        )
        trainer = Trainer(tiny_llama, args, train_dataset=train_dataset)
        trainer.create_optimizer_and_scheduler(num_training_steps=10)

        # 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_without_scheduler(self):
        self._check_lr_display_without_scheduler("stable_adamw", _ATTN_MLP_TARGET_MODULES)

    @require_torch_optimi
    @require_torch_accelerator
    def test_stable_adamw_lr_display_with_scheduler(self):
        self._check_lr_display_with_scheduler("stable_adamw", _ATTN_MLP_TARGET_MODULES, num_train_epochs=10)

    # ---------------------------------------------------------------------------
    # Misc optimizer tests
    # ---------------------------------------------------------------------------

    def test_optimizer_factory_pattern(self):
        """Test that is_optimizer_factory correctly identifies factory classes vs optimizer classes."""
        from transformers.trainer_optimizer import is_optimizer_factory

        # Create a mock optimizer class
        class MockComplexOptimizer(torch.optim.Optimizer):
            def __init__(self, params, lr=1e-3):
                defaults = {"lr": lr}
                super().__init__(params, defaults)

            def step(self, closure=None):
                pass

        # Create a factory class (simulates Muon/Dion pattern)
        class MockOptimizerFactory:
            def __call__(self, opt_model, **optimizer_kwargs):
                all_params = list(opt_model.parameters())
                return MockComplexOptimizer(all_params, **optimizer_kwargs)

        # Verify is_optimizer_factory correctly identifies factories vs optimizer classes
        self.assertFalse(is_optimizer_factory(MockComplexOptimizer))  # Optimizer class should return False
        self.assertTrue(is_optimizer_factory(MockOptimizerFactory))  # Factory class should return True

    # ---------------------------------------------------------------------------
    # Optimizer group and learning rate inspection tests
    # ---------------------------------------------------------------------------

    def test_get_optimizer_group(self):
        model = nn.Sequential(nn.Linear(128, 64))
        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = Trainer(model=model, args=TrainingArguments(output_dir=tmp_dir))
            # ValueError is raised if optimizer is None
            with self.assertRaises(ValueError):
                trainer.get_optimizer_group()
            trainer.create_optimizer()
            # Get groups
            num_groups = len(trainer.get_optimizer_group())
            self.assertEqual(num_groups, 2)
            # Get group of parameter
            param = next(model.parameters())
            group = trainer.get_optimizer_group(param)
            self.assertIn(param, group["params"])


# ---------------------------------------------------------------------------
# Custom optimizer and LR scheduler tests
# ---------------------------------------------------------------------------


class TrainerOptimizerTest(TestCasePlus):
    def test_get_optimizer_group(self):
        model = nn.Sequential(nn.Linear(128, 64))
        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = Trainer(model=model, args=TrainingArguments(output_dir=tmp_dir))
            # ValueError is raised if optimizer is None
            with self.assertRaises(ValueError):
                trainer.get_optimizer_group()
            trainer.create_optimizer()
            # Get groups
            num_groups = len(trainer.get_optimizer_group())
            self.assertEqual(num_groups, 2)
            # Get group of parameter
            param = next(model.parameters())
            group = trainer.get_optimizer_group(param)
            self.assertIn(param, group["params"])

    def test_optimizer_factory_pattern(self):
        """Test that is_optimizer_factory correctly identifies factory classes vs optimizer classes."""
        from transformers.trainer_optimizer import is_optimizer_factory

        # Create a mock optimizer class
        class MockComplexOptimizer(torch.optim.Optimizer):
            def __init__(self, params, lr=1e-3):
                defaults = {"lr": lr}
                super().__init__(params, defaults)

            def step(self, closure=None):
                pass

        # Create a factory class (simulates Muon/Dion pattern)
        class MockOptimizerFactory:
            def __call__(self, opt_model, **optimizer_kwargs):
                all_params = list(opt_model.parameters())
                return MockComplexOptimizer(all_params, **optimizer_kwargs)

        # Verify is_optimizer_factory correctly identifies factories vs optimizer classes
        self.assertFalse(is_optimizer_factory(MockComplexOptimizer))  # Optimizer class should return False
        self.assertTrue(is_optimizer_factory(MockOptimizerFactory))  # Factory class should return True

    def test_custom_optimizer(self):
        train_dataset = RegressionDataset()
        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(tmp_dir)
            model = RegressionModel()
            optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
            lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: 1.0)
            trainer = Trainer(model, args, train_dataset=train_dataset, optimizers=(optimizer, lr_scheduler))
            trainer.train()

            # Train a default model to compare against
            default_trainer = get_regression_trainer(learning_rate=0.1, output_dir=tmp_dir)
            default_trainer.train()

            self.assertFalse(torch.allclose(trainer.model.a, default_trainer.model.a))
            self.assertFalse(torch.allclose(trainer.model.b, default_trainer.model.b))
            self.assertEqual(trainer.optimizer.state_dict()["param_groups"][0]["lr"], 1.0)

    # ---------------------------------------------------------------------------
    # Weight decay parameter groups
    # ---------------------------------------------------------------------------

    def test_no_wd_param_group(self):
        model = nn.Sequential(TstLayer(128), nn.ModuleList([TstLayer(128), TstLayer(128)]))
        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = Trainer(model=model, args=TrainingArguments(output_dir=tmp_dir))
            trainer.create_optimizer_and_scheduler(10)
            wd_names = ['0.linear1.weight', '0.linear2.weight', '1.0.linear1.weight', '1.0.linear2.weight', '1.1.linear1.weight', '1.1.linear2.weight']  # fmt: skip
            wd_params = [p for n, p in model.named_parameters() if n in wd_names]
            no_wd_params = [p for n, p in model.named_parameters() if n not in wd_names]
            self.assertListEqual(trainer.optimizer.param_groups[0]["params"], wd_params)
            self.assertListEqual(trainer.optimizer.param_groups[1]["params"], no_wd_params)


@require_torch
class TrainerLRTest(TestCasePlus):
    def test_get_learning_rates(self):
        model = nn.Sequential(nn.Linear(128, 64))
        with tempfile.TemporaryDirectory() as tmp_dir:
            trainer = Trainer(model=model, args=TrainingArguments(output_dir=tmp_dir))
            with self.assertRaises(ValueError):
                trainer.get_learning_rates()
            trainer.create_optimizer()
            self.assertEqual(trainer.get_learning_rates(), [5e-05, 5e-05])

    def test_lr_scheduler_kwargs(self):
        from transformers import get_polynomial_decay_schedule_with_warmup

        # test scheduler kwargs passed via TrainingArguments
        train_dataset = RegressionDataset()
        model = RegressionModel()
        num_steps, num_warmup_steps = 10, 2
        extra_kwargs = {"power": 5.0, "lr_end": 1e-5}  # Non-default arguments
        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(
                tmp_dir,
                lr_scheduler_type="polynomial",
                lr_scheduler_kwargs=extra_kwargs,
                learning_rate=0.2,
                warmup_steps=num_warmup_steps,
            )
            trainer = Trainer(model, args, train_dataset=train_dataset)
            trainer.create_optimizer_and_scheduler(num_training_steps=num_steps)

            # Checking that the scheduler was created
            self.assertIsNotNone(trainer.lr_scheduler)

            # Checking that the correct args were passed
            sched1 = trainer.lr_scheduler
            sched2 = get_polynomial_decay_schedule_with_warmup(
                trainer.optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_steps, **extra_kwargs
            )
            self.assertEqual(sched1.lr_lambdas[0].args, sched2.lr_lambdas[0].args)
            self.assertEqual(sched1.lr_lambdas[0].keywords, sched2.lr_lambdas[0].keywords)

    def test_cosine_with_min_lr_scheduler(self):
        train_dataset = RegressionDataset()
        model = RegressionModel()
        num_steps, num_warmup_steps = 10, 2
        extra_kwargs = {"min_lr": 1e-5}  # Non-default arguments
        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(
                tmp_dir,
                lr_scheduler_type="cosine_with_min_lr",
                lr_scheduler_kwargs=extra_kwargs,
                learning_rate=0.2,
                warmup_steps=num_warmup_steps,
            )
            trainer = Trainer(model, args, train_dataset=train_dataset)
            trainer.create_optimizer_and_scheduler(num_training_steps=num_steps)

            # Checking that the scheduler was created
            self.assertIsNotNone(trainer.lr_scheduler)

            # Check the last learning rate
            for _ in range(num_steps):
                trainer.lr_scheduler.step()
            self.assertEqual(trainer.lr_scheduler.get_last_lr()[0], 1e-5)

    def test_cosine_with_min_lr_schedule_with_warmup_lr_rate(self):
        train_dataset = RegressionDataset()
        model = RegressionModel()
        num_steps, num_warmup_steps = 10, 2
        extra_kwargs = {"min_lr": 1e-5}  # Non-default arguments
        args = TrainingArguments(
            "./regression",
            lr_scheduler_type="cosine_warmup_with_min_lr",
            lr_scheduler_kwargs=extra_kwargs,
            learning_rate=0.2,
            warmup_steps=num_warmup_steps,
        )
        trainer = Trainer(model, args, train_dataset=train_dataset)
        trainer.create_optimizer_and_scheduler(num_training_steps=num_steps)

        # Checking that the scheduler was created
        self.assertIsNotNone(trainer.lr_scheduler)

        # Check the last learning rate
        step_lrs = []
        for _ in range(num_steps):
            step_lrs.append(trainer.optimizer.param_groups[0]["lr"])
            trainer.lr_scheduler.step()
        self.assertEqual(step_lrs[0], 0.1)
        self.assertEqual(step_lrs[1], 0.2)
        self.assertEqual(step_lrs[-1], 1e-05)

    def test_reduce_lr_on_plateau_args(self):
        # test passed arguments for a custom ReduceLROnPlateau scheduler
        train_dataset = RegressionDataset(length=64)
        eval_dataset = RegressionDataset(length=64)
        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(
                tmp_dir,
                eval_strategy="epoch",
                metric_for_best_model="eval_loss",
            )
            model = RegressionModel()
            optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
            lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, cooldown=2)
            trainer = Trainer(
                model,
                args,
                train_dataset=train_dataset,
                eval_dataset=eval_dataset,
                optimizers=(optimizer, lr_scheduler),
            )
            trainer.train()

            self.assertIsInstance(trainer.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau)
            self.assertEqual(trainer.lr_scheduler.factor, 0.2)
            self.assertEqual(trainer.lr_scheduler.patience, 5)
            self.assertEqual(trainer.lr_scheduler.cooldown, 2)

    def test_reduce_lr_on_plateau(self):
        # test the ReduceLROnPlateau scheduler

        class TrainerWithLRLogs(Trainer):
            def log(self, logs):
                # the LR is computed after metrics and does not exist for the first epoch
                if hasattr(self.lr_scheduler, "_last_lr"):
                    logs["learning_rate"] = self.lr_scheduler._last_lr[0]
                super().log(logs)

        train_dataset = RegressionDataset(length=64)
        eval_dataset = RegressionDataset(length=64)

        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(
                tmp_dir,
                lr_scheduler_type="reduce_lr_on_plateau",
                eval_strategy="epoch",
                metric_for_best_model="eval_loss",
                num_train_epochs=10,
                learning_rate=0.2,
            )
            model = RegressionModel()
            trainer = TrainerWithLRLogs(model, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
            trainer.train()

            self.assertIsInstance(trainer.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau)
            patience = trainer.lr_scheduler.patience

            logs = trainer.state.log_history[1:]
            best_loss = logs[0]["eval_loss"]
            bad_epochs = 0
            for i, log in enumerate(logs[:-1]):  # Compare learning rate to next epoch's
                loss = log["eval_loss"]
                just_decreased = False
                if loss > best_loss:
                    bad_epochs += 1
                    if bad_epochs > patience:
                        self.assertLess(logs[i + 1]["learning_rate"], log["learning_rate"])
                        just_decreased = True
                        bad_epochs = 0
                else:
                    best_loss = loss
                    bad_epochs = 0
                if not just_decreased:
                    self.assertEqual(logs[i + 1]["learning_rate"], log["learning_rate"])

    def test_greedy_lr_args(self):
        # test passed arguments for a custom GreedyLR scheduler
        from transformers.optimization import GreedyLR

        train_dataset = RegressionDataset(length=64)
        eval_dataset = RegressionDataset(length=64)
        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(
                tmp_dir,
                eval_strategy="epoch",
                metric_for_best_model="eval_loss",
            )
            model = RegressionModel()
            optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
            lr_scheduler = GreedyLR(optimizer, factor=0.8, patience=5, cooldown=2)
            trainer = Trainer(
                model,
                args,
                train_dataset=train_dataset,
                eval_dataset=eval_dataset,
                optimizers=(optimizer, lr_scheduler),
            )
            trainer.train()

            self.assertIsInstance(trainer.lr_scheduler, GreedyLR)
            self.assertEqual(trainer.lr_scheduler.factor, 0.8)
            self.assertEqual(trainer.lr_scheduler.patience, 5)
            self.assertEqual(trainer.lr_scheduler.cooldown, 2)

    def test_greedy_lr(self):
        # test the GreedyLR scheduler
        from transformers.optimization import GreedyLR

        class TrainerWithLRLogs(Trainer):
            def log(self, logs):
                if hasattr(self.lr_scheduler, "_last_lr"):
                    logs["learning_rate"] = self.lr_scheduler._last_lr[0]
                super().log(logs)

        train_dataset = RegressionDataset(length=64)
        eval_dataset = RegressionDataset(length=64)

        with tempfile.TemporaryDirectory() as tmp_dir:
            args = TrainingArguments(
                tmp_dir,
                lr_scheduler_type="greedy",
                lr_scheduler_kwargs={"patience": 1, "factor": 0.5},
                eval_strategy="epoch",
                metric_for_best_model="eval_loss",
                num_train_epochs=10,
                learning_rate=0.2,
            )
            model = RegressionModel()
            trainer = TrainerWithLRLogs(model, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
            trainer.train()

            self.assertIsInstance(trainer.lr_scheduler, GreedyLR)
            # Verify LR was adjusted at least once during training
            logs = trainer.state.log_history[1:]
            lr_values = [log["learning_rate"] for log in logs if "learning_rate" in log]
            self.assertTrue(len(set(lr_values)) > 1, "GreedyLR should have adjusted the LR at least once")