src/utils/training_utils.py

import os
import wandb
import torch
import traceback
from collections import OrderedDict
from rich import print
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import (
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_linear_schedule_with_warmup,
)


# ---------------------------------------------------------------------------
# General-purpose utilities
# ---------------------------------------------------------------------------

def print_rank0(*args, **kwargs) -> None:
    """Print only from rank-0 in a distributed setting, or unconditionally otherwise."""
    if dist.is_initialized():
        if dist.get_rank() == 0:
            print(*args, **kwargs)
    else:
        print(*args, **kwargs)


def format_number(num: int | float) -> str:
    """Format a large number as a human-readable string with B/M/K suffixes.

    Args:
        num (int | float): The number to format.

    Returns:
        str: Human-readable representation (e.g. '1.50B', '3.20M', '512.00K').
    """
    if num >= 1_000_000_000:
        return f"{num / 1_000_000_000:.2f}B"
    elif num >= 1_000_000:
        return f"{num / 1_000_000:.2f}M"
    elif num >= 1_000:
        return f"{num / 1_000:.2f}K"
    return str(num)


# ---------------------------------------------------------------------------
# Optimizer and learning rate scheduler creation
# ---------------------------------------------------------------------------

def create_optimizer(
    model: torch.nn.Module,
    weight_decay: float,
    learning_rate: float,
    betas: tuple[float, float],
) -> tuple[torch.optim.AdamW, dict[str, torch.nn.Parameter], dict[str, torch.nn.Parameter]]:
    """Build an AdamW optimizer with separate weight-decay groups for the model.

    1D parameters (biases, norms) and any parameter flagged with
    ``_no_weight_decay`` are placed in a zero-decay group; all others receive
    the specified weight decay.

    Args:
        model (torch.nn.Module): The model whose parameters will be optimized.
        weight_decay (float): Weight decay applied to multi-dimensional parameters.
        learning_rate (float): Base learning rate.
        betas (tuple[float, float]): AdamW beta coefficients.

    Returns:
        tuple: (optimizer, optimized_param_dict, all_param_dict) where
            optimized_param_dict contains only trainable parameters and
            all_param_dict contains every named parameter.
    """
    all_param_dict = {name: param for name, param in model.named_parameters()}
    optimized_param_dict = {name: param for name, param in all_param_dict.items() if param.requires_grad}

    decay_params, nodecay_params = [], []
    for name, param in optimized_param_dict.items():
        if param.dim() == 1 or getattr(param, '_no_weight_decay', False):
            nodecay_params.append(param)
        else:
            decay_params.append(param)
    optim_groups = [
        {'params': decay_params, 'weight_decay': weight_decay},
        {'params': nodecay_params, 'weight_decay': 0.0}
    ]
    optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)

    if dist.is_initialized():
        if dist.get_rank() == 0:
            def get_module_name(name):
                """Returns a two-level module prefix from a dotted parameter name."""
                parts = name.split('.')
                if len(parts) > 2 and parts[0] == 'module':
                    return parts[1] + '.' + parts[2]
                return parts[0]  # Fallback to first part if no 'module.' prefix
            print(f'Optimizer: AdamW, learning rate: {learning_rate}, weight decay: {weight_decay}, betas: {betas}')
            total_params = sum(p.numel() for p in model.parameters())
            trainable_params = sum(p.numel() for p in optimized_param_dict.values())
            optim_module_names = sorted(set(get_module_name(name) for name in optimized_param_dict.keys()))
            frozen_module_names = sorted(set(get_module_name(name) for name in set(all_param_dict.keys()) - set(optimized_param_dict.keys())))

            print(f'Total parameters: {format_number(total_params)}, Trainable parameters: {format_number(trainable_params)}')
            print(f'Optimized parameters: {optim_module_names}')
            print(f'Frozen parameters: {frozen_module_names}')

    return optimizer, optimized_param_dict, all_param_dict


def create_lr_scheduler(
    optimizer: torch.optim.Optimizer,
    param_update_steps: int,
    warm_up_steps: int,
    scheduler_type: str = 'cosine',
) -> torch.optim.lr_scheduler.LRScheduler:
    """Create a learning rate scheduler with linear warmup.

    Args:
        optimizer (torch.optim.Optimizer): The optimizer to schedule.
        param_update_steps (int): Total number of parameter update steps.
        warm_up_steps (int): Number of warmup steps at the start of training.
        scheduler_type (str): One of 'linear', 'cosine', or 'constant'.

    Returns:
        torch.optim.lr_scheduler.LRScheduler: The configured scheduler.

    Raises:
        ValueError: If scheduler_type is not one of the supported values.
    """
    if scheduler_type == 'linear':
        scheduler = get_linear_schedule_with_warmup(optimizer, warm_up_steps, param_update_steps)
    elif scheduler_type == 'cosine':
        scheduler = get_cosine_schedule_with_warmup(optimizer, warm_up_steps, param_update_steps)
    elif scheduler_type == 'constant':
        scheduler = get_constant_schedule_with_warmup(optimizer, warm_up_steps)
    else:
        raise ValueError(f'Invalid scheduler type: {scheduler_type}')
    return scheduler


# ---------------------------------------------------------------------------
# Checkpoint utilities
# ---------------------------------------------------------------------------

def find_checkpoints(load_path: str) -> list[str]:
    """Return sorted checkpoint paths found at load_path.

    Args:
        load_path (str): Either a directory containing .pt files, or a direct
            path to a single .pt file.

    Returns:
        list[str]: Sorted list of absolute checkpoint file paths.
    """
    if os.path.isdir(load_path):
        ckpt_names = [file_name for file_name in os.listdir(load_path) if file_name.endswith(".pt")]
        ckpt_names = sorted(ckpt_names, key=lambda x: x)
        ckpt_paths = [os.path.join(load_path, ckpt_name) for ckpt_name in ckpt_names]
    else:
        if load_path.endswith(".pt"):
            ckpt_paths = [load_path]
        else:
            ckpt_paths = []
    return ckpt_paths


def auto_resume_job(
    load_path: str,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
    reset_training_state: bool,
) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler, int, int]:
    """Resume training from the latest checkpoint in the specified directory.

    Args:
        load_path (str): If a directory, loads the last checkpoint in it;
            otherwise treats the path as a direct checkpoint file.
        model (torch.nn.Module): Model whose weights will be loaded.
        optimizer (torch.optim.Optimizer): Optimizer to restore.
        lr_scheduler (torch.optim.lr_scheduler.LRScheduler): Scheduler to restore.
        reset_training_state (bool): If True, only model weights are restored and
            optimizer/scheduler state is left at initialization.

    Returns:
        tuple: (optimizer, lr_scheduler, forward_pass_step, param_update_step).
    """
    forward_pass_step = 0
    param_update_step = 0
    all_ckpt_paths = find_checkpoints(load_path)
    if len(all_ckpt_paths) == 0:
        print_rank0(f"No checkpoint found in {load_path}, we will start from scratch")
        return optimizer, lr_scheduler, forward_pass_step, param_update_step
    try:
        ckpt_path = all_ckpt_paths[-1]
        checkpoint = torch.load(ckpt_path, map_location="cpu")
    except:
        traceback.print_exc()
        print_rank0(f"Failed to load {ckpt_path}, we will start from scratch")
        return optimizer, lr_scheduler, forward_pass_step, param_update_step

    if isinstance(model, DDP):
        status = model.module.load_state_dict(checkpoint['model'], strict=False)
    else:
        status = model.load_state_dict(checkpoint['model'], strict=False)
    print_rank0(f"Loaded model from {os.path.abspath(ckpt_path)}, the status is {status}")

    if not reset_training_state:
        try:
            optimizer.load_state_dict(checkpoint["optimizer"])
            lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
            forward_pass_step = checkpoint["fwdbwd_pass_step"]
            param_update_step = checkpoint["param_update_step"]
            print_rank0(f"Resumed optimizer and lr_scheduler from {ckpt_path}")
        except:
            traceback.print_exc()
            print_rank0(f"Failed to load optimizer and lr_scheduler from {ckpt_path}")

    return optimizer, lr_scheduler, forward_pass_step, param_update_step


def save_checkpoint(
    model: torch.nn.Module,
    ema: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
    cur_step: int,
    param_step: int,
    checkpoint_dir: str,
) -> None:
    """Serialize model/EMA/optimizer/scheduler to a timestamped checkpoint file."""
    checkpoint = {
        "model": strip_module_prefix(model.state_dict()),
        "ema": ema.state_dict(),
        "optimizer": optimizer.state_dict(),
        "lr_scheduler": lr_scheduler.state_dict(),
        "fwdbwd_pass_step": cur_step,
        "param_update_step": param_step,
    }
    os.makedirs(checkpoint_dir, exist_ok=True)
    ckpt_path = os.path.join(checkpoint_dir, f"ckpt_{cur_step:016}.pt")
    torch.save(checkpoint, ckpt_path)
    print(f"Saved checkpoint at step {cur_step} to {os.path.abspath(ckpt_path)}")


# ---------------------------------------------------------------------------
# Model parameter utilities
# ---------------------------------------------------------------------------

@torch.no_grad()
def update_ema(ema_model: torch.nn.Module, model: torch.nn.Module, decay: float = 0.999) -> None:
    """Step the EMA model towards the current model."""
    ema_params = OrderedDict(ema_model.named_parameters())
    model_params = OrderedDict(model.named_parameters())
    for name, param in model_params.items():
        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)


def requires_grad(model: torch.nn.Module, flag: bool = True) -> None:
    """Set requires_grad on all parameters of a model.

    Args:
        model (torch.nn.Module): Model to modify.
        flag (bool): Value to assign to requires_grad on every parameter.
    """
    for p in model.parameters():
        p.requires_grad = flag


def strip_module_prefix(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Remove DDP/FSDP/torch.compile key prefixes from a state dict.

    Strips '_checkpoint_wrapped_module.', '_orig_mod.', and any number of
    leading 'module.' prefixes so that weights can be loaded into an unwrapped
    model.

    Args:
        state_dict (dict[str, torch.Tensor]): The raw state dict to clean.

    Returns:
        dict[str, torch.Tensor]: State dict with prefixes removed.
    """
    new_state_dict = {}
    for key, value in state_dict.items():
        key = key.replace("_checkpoint_wrapped_module.", "")
        key = key.replace("_orig_mod.", "")
        while key.startswith("module."):
            key = key[len("module."):]
        new_state_dict[key] = value
    return new_state_dict


# ---------------------------------------------------------------------------
# Logging utilities
# ---------------------------------------------------------------------------

def log_to_console(
    epoch: int,
    cur_step: int,
    param_step: int,
    iter_time: float,
    lr: float,
    loss_dict: dict[str, float],
    print_every: int,
    start_step: int,
) -> None:
    """Print training progress when the logging criteria are met."""
    if cur_step % print_every != 0 and cur_step >= start_step + 100:
        return
    loss_str = " | ".join(f"{k}: {v:.6f}" for k, v in loss_dict.items())
    print(
        f"[Epoch {epoch:>3d}] | "
        f"Forward step: {cur_step:>6d} (Param update step: {param_step:>6d}) | "
        f"Iter time: {iter_time:.2f}s | LR: {lr:.6f}\n"
        + loss_str
    )


def log_to_wandb(
    cur_step: int,
    param_step: int,
    iter_time: float,
    lr: float,
    grad_norm: float,
    loss_dict: dict[str, float],
    wandb_log_every: int,
    start_step: int,
) -> None:
    """Log metrics to W&B when the logging criteria are met."""
    if cur_step % wandb_log_every != 0 and cur_step >= start_step + 200:
        return
    log_dict = {
        "iter": cur_step,
        "forward_pass_step": cur_step,
        "param_update_step": param_step,
        "lr": lr,
        "iter_time": iter_time,
        "grad_norm": grad_norm,
    }
    log_dict.update({"train/" + k: v for k, v in loss_dict.items()})
    wandb.log(log_dict, step=cur_step)