Repository: mosaicml/composer

Issue:
Warmup schedulers crash if `t_warmup` == `max_duration`
## To reproduce

Steps to reproduce the behavior:

1. Use a scheduler like `cosine_decay_with_warmup` and set `t_warmup` == `max_duration` == `100ba`. 
2. Launch a training run, it will run until the last step and then crash with a `ZeroDivisionError`:

```
Traceback (most recent call last):
  File "/root/composer/examples/run_composer_trainer.py", line 67, in <module>
    main()
  File "/root/composer/examples/run_composer_trainer.py", line 63, in main
    trainer.fit()
  File "/root/composer/composer/trainer/trainer.py", line 1289, in fit
    self._train_loop()
  File "/root/composer/composer/trainer/trainer.py", line 1497, in _train_loop
    scheduler.step()
  File "/usr/local/lib/python3.9/dist-packages/torch/optim/lr_scheduler.py", line 154, in step
    values = self.get_lr()
  File "/usr/local/lib/python3.9/dist-packages/torch/optim/lr_scheduler.py", line 252, in get_lr
    return [base_lr * lmbda(self.last_epoch)
  File "/usr/local/lib/python3.9/dist-packages/torch/optim/lr_scheduler.py", line 252, in <listcomp>
    return [base_lr * lmbda(self.last_epoch)
  File "/root/composer/composer/optim/scheduler.py", line 190, in scheduler_fn
    return scheduler(state, ssr)
  File "/root/composer/composer/optim/scheduler.py", line 697, in __call__
    frac_of_total = ((current_time - t_warmup) / (t_max - t_warmup)).value
  File "/root/composer/composer/core/time.py", line 308, in __truediv__
    return Time(self.value / other.value, TimeUnit.DURATION)
ZeroDivisionError: division by zero
```

## Expected behavior

We should either allow this setting, and just not attempt to increment the schedule past the warmup... or we should catch this edge case and raise a `ValueError` out on the scheduler's `__init__`.




Discussion:

Resolving patch:
diff --git a/composer/optim/scheduler.py b/composer/optim/scheduler.py
--- a/composer/optim/scheduler.py
+++ b/composer/optim/scheduler.py
@@ -132,13 +132,12 @@ def _convert_time(time: Union[str, Time[int], Time[float]], state: State, ssr: f
     assert state.max_duration is not None, 'max_duration should be set whenever schedulers are invoked'
 
     if time.unit == TimeUnit.DURATION:
-        if state.dataloader_len is None:
-            raise RuntimeError('Cannot convert time, as state.dataloader_len is None.')
         if state.max_duration.unit == TimeUnit.EPOCH:
+            if state.dataloader_len is None:
+                raise RuntimeError('Cannot convert time, as state.dataloader_len is None.')
             return Time(int(time.value * int(state.dataloader_len) * state.max_duration.value), TimeUnit.BATCH)
         return Time(int(time.value * state.max_duration.value), state.max_duration.unit)
-
-    if time.unit == TimeUnit.EPOCH:
+    elif time.unit == TimeUnit.EPOCH:
         # Epochs do not provide sufficient granularity for SSR scaling
         # e.g. if max_duration = 1ep, then any SSR would result in a new duration of 0.
         # so, convert the time into batches
@@ -693,7 +692,8 @@ def __call__(self, state: State, ssr: float = 1.0):
 
         t_max = _convert_time(self.t_max, state, ssr=ssr)
         current_time = state.timestamp.get(t_warmup.unit)
-        frac_of_total = min(1.0, ((current_time - t_warmup) / (t_max - t_warmup)).value)
+        frac_of_total = ((current_time - t_warmup) / (t_max - t_warmup)).value if (t_max > t_warmup) else 0.0
+        frac_of_total = min(1.0, frac_of_total)
 
         current_factor = self.alpha_i + frac_of_total * (self.alpha_f - self.alpha_i)
 
@@ -760,7 +760,8 @@ def __call__(self, state: State, ssr: float = 1.0):
 
         t_max = _convert_time(self.t_max, state, ssr=ssr)
         current_time = state.timestamp.get(t_warmup.unit)
-        frac_of_total = ((current_time - t_warmup) / (t_max - t_warmup)).value
+        frac_of_total = ((current_time - t_warmup) / (t_max - t_warmup)).value if (t_max > t_warmup) else 0.0
+        frac_of_total = min(1.0, frac_of_total)
 
         return _cosine_anneal(x=frac_of_total, min_y=self.alpha_f)
 
@@ -829,7 +830,8 @@ def __call__(self, state: State, ssr: float = 1.0):
 
         t_max = _convert_time(self.t_max, state, ssr=ssr)
         current_time = state.timestamp.get(t_warmup.unit)
-        frac_of_total = ((current_time - t_warmup) / (t_max - t_warmup)).value
+        frac_of_total = ((current_time - t_warmup) / (t_max - t_warmup)).value if (t_max > t_warmup) else 0.0
+        frac_of_total = min(1.0, frac_of_total)
 
         coeff = (1 - frac_of_total)**self.power
         current_factor = self.alpha_f + coeff * (1.0 - self.alpha_f)


Test patch:
diff --git a/tests/optim/test_scheduler.py b/tests/optim/test_scheduler.py
--- a/tests/optim/test_scheduler.py
+++ b/tests/optim/test_scheduler.py
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import contextlib
-from typing import Iterable, List, Optional, Type, cast
+from typing import Iterable, List, Optional, Type
 
 import pytest
 import torch
@@ -40,11 +40,11 @@ def dummy_schedulers_state(dummy_model: torch.nn.Module, rank_zero_seed: int):
                  [1.0, 0.8, 0.512]),
     pytest.param(StepScheduler(step_size='1ep', gamma=0.5), 1.0, ['500ba', '1500ba', '3500ba'], [1.0, 0.5, 0.125]),
     pytest.param(StepScheduler(step_size='10ba', gamma=0.5), 0.5, ['3ba', '8ba', '18ba'], [1.0, 0.5, 0.125]),
-    pytest.param(MultiStepScheduler(milestones=cast(List, ['10ba', '30ba', '70ba'])), 1.0,
-                 ['5ba', '20ba', '50ba', '100ba'], [1.0, 0.1, 0.01, 0.001]),
-    pytest.param(MultiStepScheduler(milestones=cast(List, ['100ba', '1ep', '0.01dur']), gamma=0.5), 1.0,
+    pytest.param(MultiStepScheduler(milestones=['10ba', '30ba', '70ba']), 1.0, ['5ba', '20ba', '50ba', '100ba'],
+                 [1.0, 0.1, 0.01, 0.001]),
+    pytest.param(MultiStepScheduler(milestones=['100ba', '1ep', '0.01dur'], gamma=0.5), 1.0,
                  ['50ba', '500ba', '5000ba', '50000ba'], [1.0, 0.5, 0.25, 0.125]),
-    pytest.param(MultiStepScheduler(milestones=cast(List, ['100ba', '1ep', '0.01dur']), gamma=0.5), 4.0,
+    pytest.param(MultiStepScheduler(milestones=['100ba', '1ep', '0.01dur'], gamma=0.5), 4.0,
                  ['200ba', '2000ba', '20000ba', '200000ba'], [1.0, 0.5, 0.25, 0.125]),
     pytest.param(LinearScheduler(), 1.0, ['100000ba', '200000ba', '400000ba'], [0.9, 0.8, 0.6]),
     pytest.param(LinearScheduler(alpha_i=0.0, alpha_f=2.0), 1.0, ['100000ba', '200000ba', '400000ba'], [0.2, 0.4, 0.8]),
@@ -72,14 +72,14 @@ def dummy_schedulers_state(dummy_model: torch.nn.Module, rank_zero_seed: int):
                  [1.0, 0.905, 0.82, 0.625]),
     pytest.param(PolynomialScheduler(power=2.0, t_max='100ba', alpha_f=0.5), 0.5, ['0ba', '10ba', '20ba', '50ba'],
                  [1.0, 0.82, 0.68, 0.5]),
-    pytest.param(MultiStepWithWarmupScheduler(t_warmup='10ba', milestones=cast(List, ['20ba', '40ba'])), 1.0,
+    pytest.param(MultiStepWithWarmupScheduler(t_warmup='10ba', milestones=['20ba', '40ba']), 1.0,
                  ['0ba', '5ba', '15ba', '25ba', '45ba'], [0.0, 0.5, 1.0, 0.1, 0.01]),
-    pytest.param(MultiStepWithWarmupScheduler(t_warmup='10ba', milestones=cast(List, ['2ep', '4ep']), gamma=0.5), 0.5,
+    pytest.param(MultiStepWithWarmupScheduler(t_warmup='10ba', milestones=['2ep', '4ep'], gamma=0.5), 0.5,
                  ['0ba', '5ba', '15ba', '1500ba', '2500ba'], [0.0, 0.5, 1.0, 0.5, 0.25]),
-    pytest.param(
-        MultiStepWithWarmupScheduler(
-            t_warmup='10ba', milestones=cast(List, ['2ep', '4ep']), gamma=0.5, scale_warmup=True), 0.5,
-        ['0ba', '5ba', '15ba', '1500ba', '2500ba'], [0.0, 1.0, 1.0, 0.5, 0.25]),
+    pytest.param(MultiStepWithWarmupScheduler(t_warmup='10ba', milestones=['2ep', '4ep'], gamma=0.5, scale_warmup=True),
+                 0.5, ['0ba', '5ba', '15ba', '1500ba', '2500ba'], [0.0, 1.0, 1.0, 0.5, 0.25]),
+    pytest.param(MultiStepWithWarmupScheduler(t_warmup='1000ep', milestones=[]), 1.0, ['0ep', '100ep', '1000ep'],
+                 [0.0, 0.1, 1.0]),
     pytest.param(ConstantWithWarmupScheduler(t_warmup='500ep'), 1.0,
                  ['0ba', '250000ba', '500000ba', '750000ba', '1000000ba'], [0.0, 0.5, 1.0, 1.0, 1.0]),
     pytest.param(ConstantWithWarmupScheduler(t_warmup='500ep', alpha=3.0), 1.0,
@@ -90,6 +90,7 @@ def dummy_schedulers_state(dummy_model: torch.nn.Module, rank_zero_seed: int):
                  [0.0, 0.5, 1.0, 1.0]),
     pytest.param(ConstantWithWarmupScheduler(t_warmup='500ep', alpha=3.0, t_max='501000ep'), 0.5,
                  ['0ba', '250000ba', '500000ba', '500500ba', '501000ba', '502000ba'], [0.0, 1.5, 3.0, 3.0, 3.0, 3.0]),
+    pytest.param(ConstantWithWarmupScheduler(t_warmup='1000ep'), 1.0, ['0ep', '100ep', '1000ep'], [0.0, 0.1, 1.0]),
     pytest.param(LinearWithWarmupScheduler(t_warmup='500ep'), 1.0,
                  ['0ba', '250000ba', '500000ba', '750000ba', '1000000ba'], [0.0, 0.5, 1.0, 0.5, 0.0]),
     pytest.param(LinearWithWarmupScheduler(t_warmup='500ep', alpha_i=3.0, alpha_f=2.0, t_max='1002ep'), 0.5,
@@ -100,18 +101,22 @@ def dummy_schedulers_state(dummy_model: torch.nn.Module, rank_zero_seed: int):
                  ['0ba', '250ba', '500ba', '249875ba'], [0.0, 0.5, 1.0, 0.5]),
     pytest.param(LinearWithWarmupScheduler(t_warmup='500ba', scale_warmup=True), 0.5,
                  ['0ba', '125ba', '250ba', '249875ba'], [0.0, 0.5, 1.0, 0.5]),
+    pytest.param(LinearWithWarmupScheduler(t_warmup='1000ep'), 1.0, ['0ep', '100ep', '1000ep'], [0.0, 0.1, 1.0]),
     pytest.param(CosineAnnealingWithWarmupScheduler(t_warmup='0.9dur'), 1.0,
                  ['0ba', '450000ba', '900000ba', '933333ba', '950000ba', '1000000ba'], [0.0, 0.5, 1.0, 0.75, 0.5, 0.0]),
     pytest.param(CosineAnnealingWithWarmupScheduler(t_warmup='0.9dur', alpha_f=0.5), 0.01,
                  ['0ba', '4500ba', '9000ba', '9333ba', '9500ba', '10000ba'], [0.0, 0.5, 1.0, 0.875, 0.75, 0.5]),
     pytest.param(CosineAnnealingWithWarmupScheduler(t_warmup='0.9dur', alpha_f=0.5, scale_warmup=True), 0.01,
                  ['0ba', '4500ba', '9000ba', '9333ba', '9500ba', '10000ba'], [0.0, 0.5, 1.0, 0.875, 0.75, 0.5]),
+    pytest.param(CosineAnnealingWithWarmupScheduler(t_warmup='1000ep'), 1.0, ['0ep', '100ep', '1000ep'],
+                 [0.0, 0.1, 1.0]),
     pytest.param(PolynomialWithWarmupScheduler(t_warmup='0.9dur'), 1.0,
                  ['0ba', '450000ba', '900000ba', '913397ba', '929289ba', '1000000ba'], [0.0, 0.5, 1.0, 0.75, 0.5, 0.0]),
     pytest.param(PolynomialWithWarmupScheduler(t_warmup='0.9dur', alpha_f=0.5), 0.01,
                  ['0ba', '4500ba', '9000ba', '9134ba', '9293ba', '10000ba'], [0.0, 0.5, 1.0, 0.875, 0.75, 0.5]),
     pytest.param(PolynomialWithWarmupScheduler(t_warmup='0.9dur', alpha_f=0.5, scale_warmup=True), 0.01,
                  ['0ba', '4500ba', '9000ba', '9134ba', '9293ba', '10000ba'], [0.0, 0.5, 1.0, 0.875, 0.75, 0.5]),
+    pytest.param(PolynomialWithWarmupScheduler(t_warmup='1000ep'), 1.0, ['0ep', '100ep', '1000ep'], [0.0, 0.1, 1.0]),
 ])
 def test_scheduler_init(scheduler: ComposerScheduler, ssr: float, test_times: List[str], expected_lrs: List[float],
                         dummy_schedulers_state: State):