feat(core): add the PiecewiseLinear learning rate scheduler

@@ -531,11 +531,29 @@ Learning rate schedulers can be optionally configured under the ``optim.sched``     name: SquareRootAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.SquareRootAnnealing'>     name: PolynomialDecayAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.PolynomialDecayAnnealing'>     name: PolynomialHoldDecayAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.PolynomialHoldDecayAnnealing'>+    name: PiecewiseLinear, schedule: <class 'nemo.core.optim.lr_scheduler.PiecewiseLinear'>     name: StepLR, schedule: <class 'torch.optim.lr_scheduler.StepLR'>     name: ExponentialLR, schedule: <class 'torch.optim.lr_scheduler.ExponentialLR'>     name: ReduceLROnPlateau, schedule: <class 'torch.optim.lr_scheduler.ReduceLROnPlateau'>     name: CyclicLR, schedule: <class 'torch.optim.lr_scheduler.CyclicLR'> +A scheduler whose shape comes from the configuration, such as ``PiecewiseLinear``, takes its points under ``optim.sched`` as well.+Here a warmup over the first 45% of the training, a decay to 90% and a final decay, with relative boundaries that follow ``max_steps``:++.. code-block:: yaml++    optim:+        lr: 1.0+        sched:+            name: PiecewiseLinear+            relative: true+            boundaries: [0.0, 0.45, 0.9, 1.0]+            values: [1e-5, 4e-3, 1e-5, 1e-6]++``max_steps`` is the number of optimizer steps of the training, a positive integer:+``ModelPT`` passes it from the trainer, or ``sched.max_steps`` sets it explicitly.+Relative boundaries require it, step boundaries ignore it.+ Scheduler Params ~~~~~~~~~~~~~~~~ 

@@ -36,6 +36,7 @@     CosineAnnealingParams,     InverseSquareRootAnnealingParams,     NoamAnnealingParams,+    PiecewiseLinearParams,     PolynomialDecayAnnealingParams,     PolynomialHoldDecayAnnealingParams,     SchedulerParams,

@@ -15,7 +15,7 @@  from dataclasses import dataclass from functools import partial-from typing import Any, Dict, Optional+from typing import Any, Dict, List, Optional   @dataclass@@ -165,6 +165,23 @@ class PolynomialHoldDecayAnnealingParams(WarmupSchedulerParams):     cycle: bool = False  +@dataclass+class PiecewiseLinearParams(SchedulerParams):+    """+    Piecewise Linear parameter config+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).++    ``max_steps`` is the number of optimizer steps of the training, a positive integer;+    ``prepare_lr_scheduler`` fills it in from the trainer or from ``sched.max_steps``.+    Relative ``boundaries`` require it, step ``boundaries`` ignore it.+    """++    max_steps: int = 0+    boundaries: Optional[List[float]] = None+    values: Optional[List[float]] = None+    relative: bool = False++ """ Pytorch Optimizers """@@ -285,5 +302,6 @@ def get_scheduler_config(name: str, **kwargs: Optional[Dict[str, Any]]) -> Sched     'WarmupAnnealingParams': WarmupAnnealingParams,     'PolynomialDecayAnnealingParams': PolynomialDecayAnnealingParams,     'PolynomialHoldDecayAnnealingParams': PolynomialHoldDecayAnnealingParams,+    'PiecewiseLinearParams': PiecewiseLinearParams,     'ReduceLROnPlateauParams': ReduceLROnPlateauParams, }

@@ -22,6 +22,7 @@     CosineAnnealing,     InverseSquareRootAnnealing,     NoamAnnealing,+    PiecewiseLinear,     PolynomialDecayAnnealing,     PolynomialHoldDecayAnnealing,     SquareAnnealing,

@@ -17,13 +17,14 @@ # pylint: disable=C0301 # flake8: noqa +import bisect import copy import dataclasses import inspect import math import warnings from functools import partial-from typing import Any, Dict, Optional, Union+from typing import Any, Dict, List, Optional, Sequence, Union  import torch.optim as optim import torch.optim.lr_scheduler as pt_scheduler@@ -713,6 +714,97 @@ def _get_lr(self, step):         return new_lrs  +class PiecewiseLinear(_LRScheduler):+    """Learning rate interpolated linearly between given values at given points of the training.++    ``boundaries`` are increasing positions in the training+    and ``values`` the learning-rate multipliers at those positions;+    between two boundaries the multiplier is interpolated linearly,+    before the first and after the last boundary it stays at the nearest value.+    Positions are optimizer steps, or fractions of ``max_steps`` with ``relative=True``,+    so that one configuration describes the schedule for any training length.+    The multipliers scale the learning rate of each parameter group;+    with ``lr: 1.0`` they are the learning rates.++    Example, a warmup over the first 45% of the training and a decay to the end:++    .. code-block:: yaml++        sched:+          name: PiecewiseLinear+          relative: true+          boundaries: [0.0, 0.45, 0.9, 1.0]+          values: [1e-5, 4e-3, 1e-5, 1e-6]++    Args:+        optimizer: the wrapped optimizer+        boundaries: increasing positions, in optimizer steps or (``relative``) in fractions of ``max_steps``+        values: learning-rate multiplier at each boundary, one per boundary+        relative: interpret ``boundaries`` as fractions of ``max_steps``+        max_steps: the number of optimizer steps of the training, a positive integer;+            required with ``relative`` and ignored otherwise;+            :func:`prepare_lr_scheduler` fills it in from the trainer or from ``sched.max_steps``+        last_epoch: the index of the last optimizer step, -1 for a fresh schedule+    """++    def __init__(+        self,+        optimizer: optim.Optimizer,+        *,+        boundaries: Optional[Sequence[float]] = None,+        values: Optional[Sequence[float]] = None,+        relative: bool = False,+        max_steps: Optional[int] = None,+        last_epoch: int = -1,+    ) -> None:+        if boundaries is None or values is None:+            raise ValueError("PiecewiseLinear needs `boundaries` and `values`")+        boundaries = [float(boundary) for boundary in boundaries]+        values = [float(value) for value in values]+        if not boundaries or len(boundaries) != len(values):+            raise ValueError("`boundaries` and `values` must have the same length and at least one entry")+        if not all(math.isfinite(value) for value in boundaries + values):+            raise ValueError("`boundaries` and `values` must be finite")+        if any(later <= earlier for earlier, later in zip(boundaries, boundaries[1:])):+            raise ValueError("`boundaries` must be strictly increasing")+        if relative and not (isinstance(max_steps, int) and max_steps > 0):+            raise ValueError(f"Relative `boundaries` need a positive integer `max_steps`, got {max_steps!r}")+        self.boundaries = boundaries+        self.values = values+        self.relative = relative+        self.max_steps = max_steps+        super().__init__(optimizer, last_epoch)++    def multiplier(self, position: float) -> float:+        """The interpolated multiplier at ``position``, a step or (``relative``) a fraction of ``max_steps``.++        Args:+            position: the position in the training++        Returns:+            the multiplier of the base learning rate+        """+        index = bisect.bisect_right(self.boundaries, position)+        if index == 0:+            return self.values[0]+        if index == len(self.boundaries):+            return self.values[-1]+        x0, x1 = self.boundaries[index - 1], self.boundaries[index]+        y0, y1 = self.values[index - 1], self.values[index]+        return (y1 - y0) / (x1 - x0) * (position - x0) + y0++    def get_lr(self) -> List[float]:+        if not self._get_lr_called_within_step:+            warnings.warn(+                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning+            )++        step = self.last_epoch+        position = step / self.max_steps if self.relative else float(step)+        multiplier = self.multiplier(position)+        return [initial_lr * multiplier for initial_lr in self.base_lrs]++ def register_scheduler(name: str, scheduler: _LRScheduler, scheduler_params: SchedulerParams):     """     Checks if the scheduler name exists in the registry, and if it doesnt, adds it.@@ -1056,6 +1148,7 @@ def compute_max_steps(     'SquareRootAnnealing': SquareRootAnnealing,     'PolynomialDecayAnnealing': PolynomialDecayAnnealing,     'PolynomialHoldDecayAnnealing': PolynomialHoldDecayAnnealing,+    'PiecewiseLinear': PiecewiseLinear,     'StepLR': pt_scheduler.StepLR,     'ExponentialLR': pt_scheduler.ExponentialLR,     'ReduceLROnPlateau': pt_scheduler.ReduceLROnPlateau,

@@ -24,7 +24,7 @@ from lightning.pytorch.utilities import rank_zero_only  from nemo.core import config, optim-from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS+from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS, PiecewiseLinear from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS from nemo.utils import logging @@ -1180,3 +1180,145 @@ def train(             dataset_len=2000,             drop_last=True,         )+++def _sgd_with_group_lrs(*lrs: float) -> torch.optim.Optimizer:+    groups = [{'params': [torch.nn.Parameter(torch.zeros(2))], 'lr': lr} for lr in lrs]+    return torch.optim.SGD(groups, lr=lrs[0])+++def _learning_rates_per_step(scheduler, optimizer: torch.optim.Optimizer, steps: int) -> list:+    """The learning rate the optimizer uses at each of ``steps`` optimizer steps."""+    rates = []+    for _ in range(steps):+        rates.append(scheduler.get_last_lr()[0])+        optimizer.step()+        scheduler.step()+    return rates+++class TestPiecewiseLinear:+    @pytest.mark.unit+    def test_registry_and_config(self):+        assert AVAILABLE_SCHEDULERS['PiecewiseLinear'] is PiecewiseLinear+        sched_config = {+            'name': 'PiecewiseLinear',+            'max_steps': 20,+            'relative': True,+            'boundaries': [0.0, 0.5, 1.0],+            'values': [0.1, 1.0, 0.0],+        }+        for cfg in (sched_config, omegaconf.OmegaConf.create(sched_config)):+            setup = optim.lr_scheduler.prepare_lr_scheduler(_sgd_with_group_lrs(1.0), cfg)+            scheduler = setup['scheduler']+            assert isinstance(scheduler, PiecewiseLinear)+            assert setup['interval'] == 'step'+            assert scheduler.multiplier(0.25) == pytest.approx(0.55)++        by_class = {+            '_target_': 'nemo.core.config.PiecewiseLinearParams',+            'params': {'boundaries': [0, 10], 'values': [0.0, 1.0]},+            'max_steps': 20,+        }+        scheduler = optim.lr_scheduler.prepare_lr_scheduler(_sgd_with_group_lrs(1.0), by_class)['scheduler']+        assert isinstance(scheduler, PiecewiseLinear)+        assert not scheduler.relative+        assert scheduler.multiplier(5) == 0.5++        assert config.schedulers.PiecewiseLinearParams().boundaries is None++    @pytest.mark.unit+    def test_steps(self):+        optimizer = _sgd_with_group_lrs(1.0)+        scheduler = PiecewiseLinear(optimizer, boundaries=[0, 4, 8], values=[0.1, 1.0, 0.5])+        rates = _learning_rates_per_step(scheduler, optimizer, steps=11)+        assert rates == pytest.approx([0.1, 0.325, 0.55, 0.775, 1.0, 0.875, 0.75, 0.625, 0.5, 0.5, 0.5])++    @pytest.mark.unit+    def test_relative_to_max_steps(self):+        optimizer = _sgd_with_group_lrs(1.0)+        scheduler = PiecewiseLinear(+            optimizer, boundaries=[0.0, 0.5, 1.0], values=[0.0, 1.0, 0.0], relative=True, max_steps=20+        )+        rates = _learning_rates_per_step(scheduler, optimizer, steps=21)+        assert rates[0] == 0.0+        assert rates[5] == pytest.approx(0.5)+        assert rates[10] == pytest.approx(1.0)+        assert rates[15] == pytest.approx(0.5)+        assert rates[20] == pytest.approx(0.0)+        # the same schedule for a longer training+        longer = PiecewiseLinear(+            _sgd_with_group_lrs(1.0), boundaries=[0.0, 0.5, 1.0], values=[0.0, 1.0, 0.0], relative=True, max_steps=200+        )+        assert longer.multiplier(100 / 200) == pytest.approx(1.0)++    @pytest.mark.unit+    def test_values_before_and_after_the_boundaries(self):+        optimizer = _sgd_with_group_lrs(1.0)+        scheduler = PiecewiseLinear(optimizer, boundaries=[2, 4], values=[0.5, 1.0])+        rates = _learning_rates_per_step(scheduler, optimizer, steps=7)+        assert rates == pytest.approx([0.5, 0.5, 0.5, 0.75, 1.0, 1.0, 1.0])++    @pytest.mark.unit+    def test_multiplies_each_group_learning_rate(self):+        optimizer = _sgd_with_group_lrs(1.0, 0.1)+        scheduler = PiecewiseLinear(optimizer, boundaries=[0, 10], values=[0.0, 1.0])+        for _ in range(5):+            optimizer.step()+            scheduler.step()+        assert scheduler.get_last_lr() == pytest.approx([0.5, 0.05])+        assert [group['lr'] for group in optimizer.param_groups] == pytest.approx([0.5, 0.05])++    @pytest.mark.unit+    def test_state_dict_round_trip(self):+        optimizer = _sgd_with_group_lrs(1.0)+        scheduler = PiecewiseLinear(optimizer, boundaries=[0, 10], values=[0.0, 1.0])+        _learning_rates_per_step(scheduler, optimizer, steps=3)+        state = scheduler.state_dict()+        restored = PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=[0, 10], values=[0.0, 1.0])+        restored.load_state_dict(state)+        assert restored.last_epoch == 3+        assert restored.get_last_lr() == pytest.approx([0.3])++    @pytest.mark.unit+    def test_invalid_arguments(self):+        with pytest.raises(ValueError, match="needs `boundaries` and `values`"):+            PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=None, values=None)+        with pytest.raises(ValueError, match="same length"):+            PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=[0, 1], values=[1.0])+        with pytest.raises(ValueError, match="same length"):+            PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=[], values=[])+        with pytest.raises(ValueError, match="strictly increasing"):+            PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=[0, 1, 1], values=[1.0, 1.0, 1.0])+        for max_steps in (None, 0, -5, 10.5, "10"):+            with pytest.raises(ValueError, match="positive integer `max_steps`"):+                PiecewiseLinear(+                    _sgd_with_group_lrs(1.0), boundaries=[0, 1], values=[1.0, 1.0], relative=True, max_steps=max_steps+                )+        with pytest.raises(ValueError, match="finite"):+            PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=[0, float('nan')], values=[1.0, 1.0])+        with pytest.raises(ValueError, match="finite"):+            PiecewiseLinear(_sgd_with_group_lrs(1.0), boundaries=[0, 1], values=[1.0, float('inf')])++    @pytest.mark.unit+    def test_production_schedule(self):+        # a training of 88,004 steps with a warmup to 45%, a decay to 90% and a final decay,+        # as a trainer with `max_steps` and relative boundaries configures it+        max_steps = 88_004+        optimizer = _sgd_with_group_lrs(1.0)+        scheduler = PiecewiseLinear(+            optimizer,+            boundaries=[0.0, 0.45, 0.9, 1.0],+            values=[1e-5, 4e-3, 1e-5, 1e-6],+            relative=True,+            max_steps=max_steps,+        )+        rates = _learning_rates_per_step(scheduler, optimizer, steps=max_steps + 1)+        assert rates[0] == 1e-5+        assert rates[-1] == pytest.approx(1e-6)+        peak = max(range(len(rates)), key=rates.__getitem__)+        assert peak in (39601, 39602)  # 0.45 * 88004 = 39601.8+        assert rates[peak] == pytest.approx(4e-3, rel=1e-4)+        assert all(later >= earlier for earlier, later in zip(rates[:peak], rates[1 : peak + 1]))+        assert all(later <= earlier for earlier, later in zip(rates[peak:], rates[peak + 1 :]))+        assert rates[79204] == pytest.approx(1e-5, rel=1e-3)  # 0.9 * 88004 = 79203.6