Repository: microsoft/DeepSpeed
File: tests/unit/runtime/test_ds_initialize.py (Python)

Code before:
        def _my_lambda(epoch):
            return epoch // 10

        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
            return LambdaLR(optimizer, _my_lambda)

        config_dict = {'train_batch_size': 1}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)

        client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

        if scheduler_type is None:
            config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
        elif scheduler_type == _LRScheduler:
            client_scheduler = LambdaLR(client_optimizer, _my_lambda)
        else:
            client_scheduler = _lr_scheduler_callable

        _, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
                                                        model=model,
                                                        model_parameters=list(model.parameters()),
                                                        optimizer=client_optimizer,
                                                        lr_scheduler=client_scheduler)
        if client_scheduler is None:
            assert not isinstance(ds_lr_scheduler, LambdaLR)
            assert isinstance(ds_lr_scheduler, WarmupLR)
        else:
            assert isinstance(ds_lr_scheduler, LambdaLR)
            assert not isinstance(ds_lr_scheduler, WarmupLR)


Review comment (suggestion):
I think a more rigorous test would be to check that each of the three conditions of `scheduler_type` in line 337-342 is correctly handled by the `deepspeed.initialize()` call

```suggestion
        if scheduler_type is None:
```

Code after:
        def _my_lambda(epoch):
            return epoch // 10

        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
            return LambdaLR(optimizer, _my_lambda)

        config_dict = {'train_batch_size': 1}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)

        client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

        if scheduler_type is None:
            config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
            client_scheduler = None
        elif scheduler_type == _LRScheduler:
            client_scheduler = LambdaLR(client_optimizer, _my_lambda)
        else:
            client_scheduler = _lr_scheduler_callable

        _, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
                                                        model=model,
                                                        model_parameters=list(model.parameters()),
                                                        optimizer=client_optimizer,
                                                        lr_scheduler=client_scheduler)
        if scheduler_type is None:
            # in this case, we initialize from config
            assert not isinstance(ds_lr_scheduler, LambdaLR)
            assert isinstance(ds_lr_scheduler, WarmupLR)
        else:
            # in this case, we initialize from passed-in scheduler
            assert isinstance(ds_lr_scheduler, LambdaLR)
            assert not isinstance(ds_lr_scheduler, WarmupLR)

    def test_diff_lrscheler_and_callable(self, scheduler_type):
        """
        In this test,
        the LambdaLR will be used for lrscheduler type
        and the StepLR will be used for callable type
        """

        from torch.optim.lr_scheduler import StepLR

        def _my_lambda(epoch):
            return epoch // 10

        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
            return StepLR(optimizer, step_size=30)

        config_dict = {'train_batch_size': 1}