mini_v41/utils.py

"""Seeding, device/dtype selection, LR schedules, parameter groups, logging."""

from __future__ import annotations

import json
import math
import os
import random
import time
from pathlib import Path

import numpy as np
import torch

__all__ = [
    "set_seed",
    "resolve_device",
    "torch_dtype",
    "build_param_groups",
    "build_scheduler",
    "grad_norm",
    "group_grad_norms",
    "MetricLogger",
    "estimate_mfu",
    "peak_flops",
    "human",
]


def set_seed(seed: int, deterministic: bool = False) -> None:
    random.seed(seed)
    np.random.seed(seed % (2**32))
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if deterministic:
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
        torch.use_deterministic_algorithms(True, warn_only=True)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def resolve_device(spec: str = "auto") -> torch.device:
    if spec == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(spec)


def torch_dtype(name: str) -> torch.dtype:
    return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[name]


def human(n: float) -> str:
    for unit, scale in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("K", 1e3)):
        if abs(n) >= scale:
            return f"{n / scale:.2f}{unit}"
    return f"{n:.0f}"


# --------------------------------------------------------------------------- #
# Optimizer
# --------------------------------------------------------------------------- #


def build_param_groups(model: torch.nn.Module, optim_cfg) -> list[dict]:
    """Three groups: decayed matrices, undecayed vectors, Engram tables.

    Engram tables are separated because they are sparsely touched -- a row is
    updated only on the steps where some n-gram hashed to it. Weight decay would
    shrink every untouched row on every step, which is exactly wrong for a
    memory. They also get their own LR multiplier.
    """
    decay, no_decay, engram = [], [], []
    seen: set[int] = set()
    for name, param in model.named_parameters():
        if not param.requires_grad or id(param) in seen:
            continue
        seen.add(id(param))
        if name.startswith("engrams.") and ".table." in name:
            engram.append(param)
        elif param.ndim >= 2:
            decay.append(param)
        else:  # norms, biases, gates, sinks
            no_decay.append(param)

    groups = [
        {"params": decay, "weight_decay": optim_cfg.weight_decay, "name": "decay"},
        {"params": no_decay, "weight_decay": 0.0, "name": "no_decay"},
    ]
    if engram:
        groups.append(
            {
                "params": engram,
                "weight_decay": 0.0,
                "lr": optim_cfg.lr * optim_cfg.engram_lr_mult,
                "name": "engram",
            }
        )
    return [g for g in groups if g["params"]]


def build_scheduler(optimizer, optim_cfg, total_steps: int):
    """Warmup + cosine/linear decay to `min_lr_ratio * lr`."""
    warmup = max(0, optim_cfg.warmup_steps)
    floor = optim_cfg.min_lr_ratio
    kind = optim_cfg.lr_schedule

    def lr_lambda(step: int) -> float:
        if warmup and step < warmup:
            return (step + 1) / warmup
        if kind == "constant":
            return 1.0
        progress = (step - warmup) / max(1, total_steps - warmup)
        progress = min(1.0, max(0.0, progress))
        if kind == "linear":
            scale = 1.0 - progress
        else:  # cosine
            scale = 0.5 * (1.0 + math.cos(math.pi * progress))
        return floor + (1.0 - floor) * scale

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)


def grad_norm(params) -> float:
    total = 0.0
    for p in params:
        if p.grad is not None:
            total += p.grad.detach().float().norm(2).item() ** 2
    return total**0.5


def group_grad_norms(model: torch.nn.Module) -> dict[str, float]:
    """Gradient norms per component, so a dead Engram or a dead router is
    visible in the logs rather than inferred from the loss curve."""
    groups: dict[str, float] = {}
    for name, p in model.named_parameters():
        if p.grad is None:
            continue
        if ".ffn.router." in name:
            key = "router"
        elif ".ffn.experts." in name:
            key = "routed_experts"
        elif ".ffn.shared_expert." in name:
            key = "shared_expert"
        elif name.startswith("engrams.") and ".table." in name:
            key = "engram_table"
        elif name.startswith("engrams."):
            key = "engram_proj"
        elif name.startswith("global_kv"):
            key = "global_kv"
        elif ".attn." in name:
            key = "attention"
        elif name.startswith("embed") or name.startswith("lm_head"):
            key = "embedding"
        else:
            key = "other"
        groups[key] = groups.get(key, 0.0) + p.grad.detach().float().norm(2).item() ** 2
    return {k: v**0.5 for k, v in groups.items()}


# --------------------------------------------------------------------------- #
# MFU
# --------------------------------------------------------------------------- #

# Dense BF16 peak, TFLOP/s. Used only for a rough MFU estimate; unknown GPUs
# report MFU as None rather than a made-up number. Dense, not the marketing
# "with sparsity" figure -- 989 for H100/H200 is the 2:4-sparsity number, and
# mixing conventions would make H100 MFU read ~2x lower than A100 MFU.
_PEAK_BF16_TFLOPS = {
    "H200": 494.7, "H100": 494.7, "A100": 312.0,
    "RTX 5090": 210.0, "RTX 5080": 170.0, "RTX 5060 Ti": 90.0,
    "RTX 4090": 165.2, "L40S": 181.0,
}


def peak_flops(device: torch.device) -> float | None:
    if device.type != "cuda":
        return None
    name = torch.cuda.get_device_name(device)
    for key, tflops in _PEAK_BF16_TFLOPS.items():
        if key in name:
            return tflops * 1e12
    return None


def estimate_mfu(
    active_params: int, tokens_per_second: float, device: torch.device
) -> float | None:
    """Rough MFU from the 6*N*tokens rule (fwd+bwd), attention FLOPs excluded.

    For an MoE model N is the *active* parameter count, which is why
    parameter_report() computes it. Attention's quadratic term is omitted; at
    L=2048 with these hidden sizes it is a few percent of the total, so treat
    this as a lower bound.
    """
    peak = peak_flops(device)
    if not peak or tokens_per_second <= 0:
        return None
    return (6.0 * active_params * tokens_per_second) / peak


# --------------------------------------------------------------------------- #
# Logging
# --------------------------------------------------------------------------- #


class MetricLogger:
    """TensorBoard + JSONL + console, all optional, same call site."""

    def __init__(self, out_dir: str | Path, use_tensorboard: bool = True,
                 use_wandb: bool = False, wandb_project: str = "mini-v41", config: dict | None = None):
        self.out_dir = Path(out_dir)
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.jsonl = (self.out_dir / "metrics.jsonl").open("a", encoding="utf-8")
        self.writer = None
        self.wandb = None
        self.t0 = time.time()

        if use_tensorboard:
            try:
                from torch.utils.tensorboard import SummaryWriter

                self.writer = SummaryWriter(str(self.out_dir / "tb"))
            except Exception as exc:  # pragma: no cover
                print(f"[log] tensorboard unavailable: {exc}")
        if use_wandb:
            try:
                import wandb

                wandb.init(project=wandb_project, dir=str(self.out_dir), config=config)
                self.wandb = wandb
            except Exception as exc:  # pragma: no cover
                print(f"[log] wandb unavailable: {exc}")

    def log(self, step: int, metrics: dict, prefix: str = "") -> None:
        flat = {}
        for key, value in metrics.items():
            if value is None:
                continue
            name = f"{prefix}{key}"
            if isinstance(value, (int, float)) and math.isfinite(float(value)):
                flat[name] = float(value)
            elif isinstance(value, (int, float)):
                flat[name] = float(value)  # keep inf/nan in jsonl; skip for TB
        record = {"step": step, "wall": time.time() - self.t0, **flat}
        self.jsonl.write(json.dumps(record) + "\n")
        self.jsonl.flush()
        if self.writer is not None:
            for name, value in flat.items():
                if math.isfinite(value):
                    self.writer.add_scalar(name, value, step)
        if self.wandb is not None:
            self.wandb.log(record, step=step)

    def close(self) -> None:
        self.jsonl.close()
        if self.writer is not None:
            self.writer.close()