src/optim.py

from __future__ import annotations

import torch

from configs import cfg

def fused_adamw_preflight(logger) -> bool:
    if not torch.cuda.is_available():
        logger.info("  Optimizer:  fused AdamW requested but CUDA is unavailable; using standard AdamW")
        return False

    try:
        probe = torch.nn.Parameter(torch.ones(8, device="cuda", dtype=torch.bfloat16))
        probe_optim = torch.optim.AdamW([probe], lr=1.0e-4, fused=True)
        loss = probe.float().square().sum()
        loss.backward()
        probe_optim.step()
        probe_optim.zero_grad(set_to_none=True)
        del loss, probe_optim, probe
        return True
    except Exception as exc:
        logger.warning(f"  Optimizer:  fused AdamW preflight failed ({exc}); using standard AdamW")
        return False

def build_adamw_optimizer(params: list[torch.nn.Parameter], logger, allow_fused: bool) -> torch.optim.Optimizer:
    kwargs = {
        "lr": cfg.training.learning_rate,
        "weight_decay": cfg.training.weight_decay,
        "betas": (0.9, 0.999),
    }
    fused_requested = bool(getattr(cfg.training, "fused_adamw", False)) and allow_fused
    if fused_requested and fused_adamw_preflight(logger):
        try:
            optimizer = torch.optim.AdamW(params, **kwargs, fused=True)
            logger.info("  Optimizer:  AdamW fused=True")
            return optimizer
        except Exception as exc:
            logger.warning(f"  Optimizer:  fused AdamW construction failed ({exc}); using standard AdamW")
    elif bool(getattr(cfg.training, "fused_adamw", False)) and not allow_fused:
        logger.info("  Optimizer:  fused AdamW disabled for DeepSpeed")

    optimizer = torch.optim.AdamW(params, **kwargs)
    logger.info("  Optimizer:  AdamW standard")
    return optimizer