dtypes.py

"""Compute-dtype configuration for the volumetric RMU stack.

One place to declare which tensors run in which precision so train.py,
inference.py and chat.py stay in lock-step. It is safe to lower precision
inside the core slabs because every non-linearity in activations.py is
a bounded saturating map (siloid/bound/scan_affine on [-1, 1]),
so over/underflow cannot cascade the way it does with unbounded
ReLU/GELU stacks. Outside the core we keep fp32 for numerically
sensitive ops: vocab log-sum-exp in loss_chunked, optimizer states,
normalization statistics, and the scan's analytic backward.

CLI convention:
    --dtype {fp32, fp16, bf16, fp8_e4m3, fp8_e5m2}

Internally we expose two things:
  * ``parse_compute_dtype(name)`` -> torch.dtype or fp8 marker tuple.
  * ``autocast_ctx(device, dtype)`` -> context manager that wraps the
    forward call. For fp32 it is a no-op. For bf16/fp16 it is
    ``torch.autocast`` with the matching dtype. For fp8 it raises a
    clear NotImplementedError pointing at the required follow-up
    (per-tensor absmax pack/unpack around slab boundaries).

Design note on fp8: PyTorch exposes ``torch.float8_e4m3fn`` and
``torch.float8_e5m2`` as storage dtypes but, as of this codebase's
torch version) no production-grade autocast driver — matmul still
requires an upstream library (TransformerEngine, or a manual
quant-compute-dequant dance with per-tensor scales). Rather than ship a
silent bf16 fallback that pretends to be fp8, we reserve the name and
raise with exact instructions, and keep the CLI option so training
resumes won't reject the flag once fp8 lands.
"""

from contextlib import contextmanager, nullcontext
from typing import Optional, Union

import torch


# ---------- public names -------------------------------------------------

DTYPE_NAMES = ('fp32', 'fp16', 'bf16', 'fp8_e4m3', 'fp8_e5m2')

_TORCH_DTYPES = {
    'fp32': torch.float32,
    'fp16': torch.float16,
    'bf16': torch.bfloat16,
}

# fp8 marker: the raw torch dtype exists as storage, but compute
# requires additional plumbing, see the module docstring.
# parse_compute_dtype returns it as a sentinel string rather than a
# torch.dtype so call sites can route to the pack/unpack path without
# touching bf16 logic.
_FP8_DTYPES = {
    'fp8_e4m3': getattr(torch, 'float8_e4m3fn', None),
    'fp8_e5m2': getattr(torch, 'float8_e5m2', None),
}


ComputeDtype = Union[torch.dtype, str]  # str only for fp8_* markers


def parse_compute_dtype(name: str) -> ComputeDtype:
    """Resolve a --dtype string to a torch.dtype or an fp8 marker."""
    key = name.lower().strip()
    if key in _TORCH_DTYPES:
        return _TORCH_DTYPES[key]
    if key in _FP8_DTYPES:
        if _FP8_DTYPES[key] is None:
            raise RuntimeError(
                f'{key} requested but this torch build does not expose '
                f'the storage dtype (needs >=2.1 with fp8 types).'
            )
        return key  # return marker string; callers route via is_fp8().
    raise ValueError(f'unknown dtype {name!r}; expected one of {DTYPE_NAMES}')


def is_fp8(dt: ComputeDtype) -> bool:
    return isinstance(dt, str) and dt.startswith('fp8_')


def is_low_precision(dt: ComputeDtype) -> bool:
    """True for any non-fp32 compute dtype."""
    if is_fp8(dt):
        return True
    return dt in (torch.float16, torch.bfloat16)


# ---------- autocast -----------------------------------------------------

def autocast_ctx(device: Union[str, torch.device], dtype: ComputeDtype):
    """Context manager wrapping the safe low-precision compute region.

    fp32 => nullcontext, zero overhead.
    bf16 / fp16 => torch.autocast, parameters stay fp32 while autocast
        casts activations inside matmul / conv / linear on entry.
    fp8 => NotImplementedError with a pointer to the required follow-up.
        The flag still parses so run_config roundtrips; swap to a real
        pack/unpack path once slab-boundary scales are in place.
    """
    if dtype is torch.float32:
        return nullcontext()
    if dtype in (torch.float16, torch.bfloat16):
        device_type = torch.device(device).type if not isinstance(device, str) \
            else torch.device(device).type
        return torch.autocast(device_type=device_type, dtype=dtype)
    if is_fp8(dtype):
        raise NotImplementedError(
            f'Compute dtype {dtype} requires per-tensor absmax pack/unpack '
            'around every slab boundary (RMU in-proj / out-proj / vAttn qk). '
            'Expose that path via fp8_quantize() + fp8_dequantize() and '
            'remove this guard to enable. Until then pass --dtype bf16 '
            '(recommended) for memory savings at stable loss.'
        )
    raise TypeError(f'unsupported compute dtype {dtype!r}')


# ---------- fp8 pack / unpack (stubs, wired when fp8 kernel lands) -------

def fp8_quantize(x: torch.Tensor, fmt: str = 'fp8_e4m3'):
    """Per-tensor absmax quantize fp32/bf16 -> (fp8_tensor, scale fp32).

    Stub for future use: the scale is the absmax / max-representable
    magnitude of the fp8 format, and dequantize is a simple cast * scale.
    Intentionally raises until fp8 path is activated so no caller mixes
    dtypes silently.
    """
    raise NotImplementedError('fp8_quantize: enable with fp8 compute path')


def fp8_dequantize(x_fp8: torch.Tensor, scale: torch.Tensor, out_dtype: torch.dtype):
    raise NotImplementedError('fp8_dequantize: enable with fp8 compute path')


class LowPrecGradAccumulator:
    """Per-parameter low-precision grad offloader + accumulator.

    Uses ``Tensor.register_post_accumulate_grad_hook`` (PyTorch >= 2.1)
    to drain each parameter's ``grad`` the instant it is finalized by
    autograd, i.e. *during* backward, into a low-precision buffer on a
    (possibly different) device, and immediately sets ``param.grad``
    back to ``None``. This way fp32 grad tensors never co-reside on the
    compute GPU: each one exists only for the brief window between
    "autograd finalised it" and "hook copied and nulled it".

    Across the multi-task mixer's 3 backwards the hook also accumulates
    (sum) into the same low-prec buffer on the offload device, so one
    per-param buffer survives the whole step. Before ``optimizer.step``
    you call ``flush_to_fp32()`` which materialises the accumulated
    grads back into ``param.grad`` as fp32 on the param's device.

    Blocking copies are used on purpose: an earlier
    non_blocking=True version leaked cached blocks on the source GPU
    because the async copy kept the source tensor alive in the
    allocator's bookkeeping.
    """

    def __init__(self, params, accum_dtype: torch.dtype,
                 accum_device: Optional[Union[str, torch.device]] = None):
        assert accum_dtype in (torch.float16, torch.bfloat16), \
            f'accum_dtype must be fp16 or bf16, got {accum_dtype}'
        self.accum_dtype = accum_dtype
        self.accum_device = torch.device(accum_device) if accum_device is not None else None
        self.params = [p for p in params if p.requires_grad]
        self._buf: dict = {id(p): None for p in self.params}
        self._hooks = []
        if not hasattr(torch.Tensor, 'register_post_accumulate_grad_hook'):
            raise RuntimeError(
                'LowPrecGradAccumulator needs PyTorch >= 2.1 for '
                'register_post_accumulate_grad_hook.'
            )
        for p in self.params:
            self._hooks.append(p.register_post_accumulate_grad_hook(self._make_hook()))

    def _target_device(self, p: torch.Tensor) -> torch.device:
        return self.accum_device if self.accum_device is not None else p.device

    def _make_hook(self):
        # Closure capturing self. Fires once per param per backward once
        # autograd has finalized param.grad. We drain synchronously so
        # the compute GPU reclaims the fp32 grad block before the next
        # autograd node allocates its own grad tensor.
        def hook(param: torch.Tensor):
            g = param.grad
            if g is None:
                return
            dev = self._target_device(param)
            # Step 1: cross-device blocking copy, fp32 -> fp32 on target.
            if dev != g.device:
                g_on_dev = g.detach().to(dev)  # blocking, sync copy
            else:
                g_on_dev = g.detach()
            # Step 2: cast to low precision on the target device.
            g_lp = g_on_dev.to(self.accum_dtype)
            del g_on_dev  # free the intermediate fp32 on target device
            # Step 3: accumulate into persistent buffer.
            key = id(param)
            if self._buf[key] is None:
                self._buf[key] = g_lp
            else:
                self._buf[key].add_(g_lp)
                del g_lp
            # Step 4: release the fp32 grad on the compute device NOW.
            param.grad = None
        return hook

    def flush_to_param_grads(self, out_dtype: Optional[torch.dtype] = torch.float32):
        """Move buffered low-prec grads back into ``param.grad`` on the
        parameter device, then reset the internal buffer.

        If ``out_dtype`` is None, keeps the accumulator dtype.
        """
        for p in self.params:
            b = self._buf[id(p)]
            if b is None:
                continue
            tgt_dtype = self.accum_dtype if out_dtype is None else out_dtype
            p.grad = b.to(device=p.device, dtype=tgt_dtype)  # blocking
            self._buf[id(p)] = None

    def flush_to_fp32(self):
        """Backward-compatible alias for legacy callers."""
        self.flush_to_param_grads(out_dtype=torch.float32)

    def reset(self):
        for k in self._buf:
            self._buf[k] = None

    def close(self):
        for h in self._hooks:
            h.remove()
        self._hooks.clear()


class OffloadedOptim:
    """Generic offloaded optimizer whose master fp32 weights and
    optimizer states live on a different device from the compute params.

    Supports any torch.optim.Optimizer subclass (AdamW, Atom, etc.).

    Memory accounting for a P-param model:
      * Compute GPU: params fp32 (4P) stays for forward/backward.
      * Offload device: master fp32 (4P) + optimizer states (~8-12P depending on optimizer).
      * Offload device also holds the grad buffer (bf16, 2P) if paired
        with LowPrecGradAccumulator(accum_device == optim_device).

    Net saving on the compute GPU vs a vanilla optimizer:
      * optimizer states freed: ~8-12P
      * fp32 grad window freed: 4P, because ingest pulls grads to the
        master side and nulls compute .grad immediately.

    Total: ~12-16P bytes removed from the compute GPU, at the cost of
    cross-device copies of (a) grads on ingest, (b) updated params on
    sync. Both are O(P) per optimizer step, trivially less than a full
    forward/backward.

    Usage:
        opt = OffloadedOptim(optim_groups, optim_device='cuda:1',
                             optim_class=torch.optim.AdamW,
                             lr=..., weight_decay=..., betas=..., eps=...)
        ...
        # after backward(s) and optional LowPrecGradAccumulator usage:
        opt.ingest_grads(grad_accum=grad_accum)   # or grad_accum=None
        grad_norm = opt.clip_grad_norm(args.grad_clip)
        opt.step()             # runs optimizer on optim_device + syncs params back
        opt.zero_grad()

    The object exposes ``.param_groups`` / ``.state_dict`` /
    ``.load_state_dict`` that forward to the internal master optimizer,
    so existing LR schedulers and checkpoint IO keep working unchanged.
    """

    def __init__(self, compute_param_groups, optim_device: Union[str, torch.device],
                 optim_class=torch.optim.AdamW,
                 master_dtype: torch.dtype = torch.float32,
                 **optim_kwargs):
        self.optim_device = torch.device(optim_device)
        self._optim_class = optim_class
        self.master_dtype = master_dtype

        self._compute_params: list = []
        self._master_params: list = []
        master_groups = []
        for g in compute_param_groups:
            ms = []
            for p in g['params']:
                self._compute_params.append(p)
                mp = torch.nn.Parameter(
                    p.detach().to(self.optim_device, dtype=self.master_dtype).clone()
                )
                mp.requires_grad_(True)
                self._master_params.append(mp)
                ms.append(mp)
            new_g = {k: v for k, v in g.items() if k != 'params'}
            new_g['params'] = ms
            master_groups.append(new_g)

        self._opt = optim_class(master_groups, **optim_kwargs)

    # --- torch.optim.Optimizer-ish forwarded interface -----------------

    @property
    def param_groups(self):
        # Master groups. External code (LR scheduler) writes here; keeps
        # working unchanged because master is what actually steps.
        return self._opt.param_groups

    def state_dict(self):
        return self._opt.state_dict()

    def load_state_dict(self, sd):
        self._opt.load_state_dict(sd)
        # Re-sync compute params to master so forward sees loaded weights.
        self._sync_master_to_compute()

    # --- custom grad ingest / clip / step ------------------------------

    def ingest_grads(self, grad_accum: Optional['LowPrecGradAccumulator'] = None) -> None:
        """Pull grads onto the master params on optim_device.

        If ``grad_accum`` is provided, its bf16/fp16 buffer is drained
        into ``master.grad`` as fp32 (buffer is cleared in place).
        Otherwise we read ``compute_param.grad`` fp32 from the compute
        device, copy to optim_device, and null out compute .grad.
        """
        dev = self.optim_device
        if grad_accum is not None:
            for cp, mp in zip(self._compute_params, self._master_params):
                buf = grad_accum._buf.get(id(cp))
                if buf is None:
                    mp.grad = None
                else:
                    mp.grad = buf.to(device=dev, dtype=mp.dtype)
                    grad_accum._buf[id(cp)] = None
        else:
            for cp, mp in zip(self._compute_params, self._master_params):
                g = cp.grad
                if g is None:
                    mp.grad = None
                else:
                    mp.grad = g.detach().to(device=dev, dtype=mp.dtype)
                    cp.grad = None

    def clip_grad_norm(self, max_norm: float):
        if max_norm is None or max_norm <= 0:
            return None
        return torch.nn.utils.clip_grad_norm_(self._master_params, max_norm)

    def _sync_master_to_compute(self) -> None:
        for cp, mp in zip(self._compute_params, self._master_params):
            cp.data.copy_(mp.data.to(cp.device))

    def step(self) -> None:
        self._opt.step()
        self._sync_master_to_compute()

    def zero_grad(self, set_to_none: bool = True) -> None:
        self._opt.zero_grad(set_to_none=set_to_none)
        # Also null compute .grad (safety: hooks from LowPrecGradAccumulator
        # already do this, but external code may have set it).
        for cp in self._compute_params:
            if set_to_none:
                cp.grad = None
            elif cp.grad is not None:
                cp.grad.zero_()

    # Expose last_stats for Atom compatibility
    @property
    def last_stats(self):
        """Access optimizer-specific stats (e.g., Atom's trust_mean, commit_frac)."""
        return getattr(self._opt, 'last_stats', {})


# Backward compatibility alias
class OffloadedAdamW(OffloadedOptim):
    """Backward compatibility alias for OffloadedOptim with AdamW."""

    def __init__(self, compute_param_groups, optim_device: Union[str, torch.device],
                 **adamw_kwargs):
        super().__init__(compute_param_groups, optim_device,
                         optim_class=torch.optim.AdamW, **adamw_kwargs)


__all__ = [
    'DTYPE_NAMES',
    'ComputeDtype',
    'parse_compute_dtype',
    'is_fp8',
    'is_low_precision',
    'autocast_ctx',
    'fp8_quantize',
    'fp8_dequantize',
    'LowPrecGradAccumulator',
    'OffloadedOptim',
    'OffloadedAdamW',
]