MultiStepLR (lr_scheduler/multistep.py)

Learning‑rate schedulers with multi‑step decay and optional warmup for PyTorch optimizers.

# Source: https://github.com/hpcaitech/ColossalAI | License: Apache-2.0 | Copyright 2021- HPC-AI Technology Inc. All rights reserved
# Provided via justclaw.it.com; keep this notice with the code.
from typing import List

from torch.optim.lr_scheduler import MultiStepLR as _MultiStepLR

from .delayed import WarmupScheduler


class MultiStepLR(_MultiStepLR):
    """Decays the learning rate of each parameter group by gamma once the
    number of epoch reaches one of the milestones. Notice that such decay can
    happen simultaneously with other changes to the learning rate from outside
    this scheduler. When last_epoch=-1, sets initial lr as lr.

    Args:
        optimizer (:class:`torch.optim.Optimizer`): Wrapped optimizer.
        total_steps (int): Number of total training steps.
     

... (truncated -- full source via MCP)