src/taoTrain/schedulers/linear_warmup.py

"""Linear warmup learning rate scheduler."""

import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR
from taoTrain.config import TrainingConfig
from .registry import register_scheduler


@register_scheduler("linearWarmup")
def create_linear_warmup(
    optimizer: optim.Optimizer,
    config: TrainingConfig,
    num_training_steps: int,
) -> LambdaLR:
    """
    Create a linear warmup scheduler.
    
    Linearly increases learning rate from 0 to peak over warmup steps,
    then keeps it constant.
    
    Args:
        optimizer: Optimizer instance
        config: TrainingConfig with scheduler configuration
        num_training_steps: Total number of training steps
    
    Returns:
        LambdaLR scheduler instance
    """
    scheduler_config = config.scheduler
    
    # Determine warmup steps
    if scheduler_config.warmup_steps > 0:
        warmup_steps = scheduler_config.warmup_steps
    else:
        warmup_steps = int(num_training_steps * scheduler_config.warmup_ratio)
    
    def lr_lambda(step):
        """Linear warmup learning rate schedule."""
        if step < warmup_steps:
            return float(step) / float(max(1, warmup_steps))
        return 1.0
    
    return LambdaLR(optimizer, lr_lambda, last_epoch=scheduler_config.last_epoch)