optimizer.py

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import torch
# ------------------------------------------------------------------------
# SlabAdamW: per-slab x per-subgroup param groups.
#
# Single torch.optim.AdamW, but with param_groups partitioned by
# (slab_id, subgroup_role). No extra memory vs a monolithic AdamW: the
# Adam state has the same total size, just split into logical groups.
# Scales automatically with num_slabs, no hardcoded "slab1/slab2" etc.
#
# Subgroup roles, per slab (populated from SlabParameterBank + submodules):
#   scan        : fused scan group q->(sa1, sa2, sa3, ig) (decay/injection gates)
#   readout     : readout_{w,b}, out_{w,b}              (TXSSM readout + slab outputs)
#   gating      : bp_{w,b}, wg_{w,b}                    (blueprint + write gate)
#   projection  : q/k/u/h/v/c weights+biases            (linear feature ops)
#   conv        : rmu.{local,cross,proj_x}, mix.{dw,pw} (all local convs)
#   vattn       : vattn.* except gq/gk                  (attention core)
#   vattn_gqk   : vattn.gq, vattn.gk                    (slow-warmup gain pair)
#
# Top-level, i.e. non-slab, parameters land in group "shared/misc":
# core_in_proj, core_out_proj, decoder, and, if enabled, container.
# ------------------------------------------------------------------------

def count_params(m: torch.nn.Module) -> int:
    return sum(p.numel() for p in m.parameters())


def _group_key(name: str):
    """Map a parameter name to (slab_or_'shared', subgroup_role) or
    ('shared', 'misc') for non-slab params. Returns None for anything
    that should be rejected (currently nothing)."""
    import re

    def _classify_shared(shared_name: str):
        # Keep shared params split by execution path so the logs stay readable.
        if shared_name == 'core.xmem_decay_logits':
            return ('shared', 'xmem_decay')
        if shared_name.startswith('core_in_proj.'):
            if '.ctx_dw.' in shared_name:
                return ('shared', 'input_ctx_dw')
            if '.kan.' in shared_name:
                return ('shared', 'input_kan')
            if '.ctx_pw.' in shared_name:
                return ('shared', 'input_ctx_pw')
            if '.content_proj.' in shared_name:
                return ('shared', 'input_content')
            # All three per-call-site LeaRNorm modules in CoreInputStem
            # land in the same group so the LR schedule treats them
            # uniformly while the log can still see their aggregate ||g||.
            if ('.norm_in.' in shared_name or '.norm_content.' in shared_name
                    or '.norm_context.' in shared_name):
                return ('shared', 'input_norm')
            return ('shared', 'input_other')
        if shared_name.startswith('core_out_proj.'):
            return ('shared', 'core_out_proj')
        if shared_name.startswith('decoder.'):
            tail = shared_name[len('decoder.'):]
            if tail.startswith('embedding.'):
                return ('shared', 'decoder_embedding')
            if tail.startswith('emb_proj.'):
                return ('shared', 'decoder_emb_proj')
            if tail.startswith('norm.'):
                return ('shared', 'decoder_norm')
            return ('shared', 'decoder_other')
        if shared_name.startswith('container.'):
            tail = shared_name[len('container.'):]
            m_layer = re.match(r'^route_layers\.(\d+)\.', tail)
            if m_layer:
                return ('shared', f'container_layer_{int(m_layer.group(1)) + 1}')
            m_skip = re.match(r'^route_skip_heads\.(\d+)\.', tail)
            if m_skip:
                return ('shared', f'container_layer_skip_{int(m_skip.group(1)) + 1}')
            if tail.startswith('alpha_head.'):
                return ('shared', 'container_alpha')
            if tail.startswith('slab_gate_head.'):
                return ('shared', 'container_slab_gate')
            if tail.startswith('neuron_gate_head.'):
                return ('shared', 'container_neuron_gate')
            if tail.startswith('tau_sa1_head.'):
                return ('shared', 'container_tau_sa1')
            if tail.startswith('tau_sa2_head.'):
                return ('shared', 'container_tau_sa2')
            if tail.startswith('tau_sa3_head.'):
                return ('shared', 'container_tau_sa3')
            if tail.startswith('state_proj.'):
                return ('shared', 'container_state_frozen')
            return ('shared', 'container_other')
        return ('shared', 'misc')

    m = re.match(r'^core\.slabs\.(slab\d+)\.(.*)$', name)
    if not m:
        return _classify_shared(name)
    slab = m.group(1)
    tail = m.group(2)
    # bank.params.{slab}_{role}_{w|b}
    m2 = re.match(r'^bank\.params\.' + re.escape(slab) + r'_([a-z0-9_]+)_([wb])$', tail)
    if m2:
        t = m2.group(1)
        # Fused groups: 'qku' (q,k,u); 'scan' (sa1,sa2,ig); 'bpwg' (bp,wg).
        if t == 'scan':
            return (slab, 'scan')
        if t == 'bpwg':
            return (slab, 'gating')
        if t in ('readout', 'readout_h', 'readout_c', 'col_read', 'out'):
            return (slab, 'readout')
        if t in ('qku', 'h', 'v', 'c'):
            return (slab, 'projection')
        return (slab, 'bank_other')
    if tail.startswith('vattn.'):
        leaf = tail[len('vattn.'):]
        if leaf in ('gq', 'gk'):
            return (slab, 'vattn_gqk')
        return (slab, 'vattn')
    # Per-slab learnable LeaRNorm (rmu.norm_proj.{gamma,beta}) gets its
    # own group, ahead of the generic 'rmu.' rule below, so the optimizer
    # and log can treat it separately from the conv stack.
    if tail.startswith('rmu.norm_proj.'):
        return (slab, 'norm')
    if tail.startswith('xmem.'):
        return (slab, 'xmem')
    if tail.startswith('dsl.'):
        return (slab, 'dsl')
    if tail.startswith('mix.') or tail.startswith('rmu.'):
        return (slab, 'conv')
    # PCSC predictor head + surprise-gain scalar live next to the slab's
    # bank but receive an independent LR (a small, fast-warming
    # predictor produces a more stable surprise signal; a large surprise
    # gain at init can drown the wg dynamics). Keep them in a single
    # group so the trainer can dampen both with one `predictor_lr_scale`.
    if tail.startswith('state_predictor.') or tail == 'surprise_gain':
        return (slab, 'predictor')
    return (slab, 'slab_other')


def build_groups(model, base_lr: float, weight_decay: float,
                            gqk_lr_scale: float,
                            extra_lr_scales: Optional[Dict[str, float]] = None,
                            ) -> List[Dict]:
    """Produce the optim param_groups list for SlabAdamW.

    Each group carries its own (lr, weight_decay, name, base_lr, is_gqk)
    so the LR schedule and per-group diagnostics can address them by
    name. Iteration order is deterministic, sorted by (slab, role),
    so checkpoint state matches across runs.

    `extra_lr_scales`: optional `{sub_name: multiplier}` table that
    multiplies `base_lr` for the named subgroup (e.g. `'decoder_emb_proj':
    0.3` to dampen the high-leverage factor→hidden lift, which sits on
    the forward path AND on the tied readout and so has ~3-5x the
    per-step loss-sensitivity of slab interior matrices). Applied AFTER
    `gqk_lr_scale`, so the two compose multiplicatively when both match.
    """
    extra_lr_scales = extra_lr_scales or {}
    buckets: Dict[Tuple[str, str], List[Tuple[str, torch.nn.Parameter]]] = {}
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        key = _group_key(name)
        buckets.setdefault(key, []).append((name, p))

    shared_order = {
        # Execution order in ViperLanguageModel.forward():
        # decoder.encode -> core_in_proj -> core(slabs/container)
        # -> core_out_proj -> decoder.norm/head.
        # IO-side shared modules are laid out close to that order.
        'input_norm': 9,        # LeaRNorm between input convs
        'input_kan': 10,
        'input_content': 10,
        'input_ctx_dw': 11,
        'input_ctx_pw': 12,
        'input_other': 13,
        'core_out_proj': 20,
        'decoder_embedding': 30,
        'decoder_emb_proj': 31,
        'decoder_norm': 32,
        'decoder_other': 33,
        'container_alpha': 40,
        'container_layer_in': 40,
        'container_slab_gate': 41,
        'container_neuron_gate': 42,
        'container_tau_sa1': 43,
        'container_tau_sa2': 44,
        'container_tau_sa3': 45,
        'container_state_frozen': 46,
        'container_other': 47,
        'xmem_decay': 15,
        'misc': 99,
    }
    slab_order = {
        'norm': 5,              # Per-slab LeaRNorm in RMU
        'scan': 10,
        'gating': 20,
        'projection': 30,
        'conv': 40,
        'vattn_gqk': 50,
        'vattn': 51,
        'readout': 60,
        'xmem': 55,
        'dsl': 56,
        'predictor': 65,
        'bank_other': 70,
        'slab_other': 80,
    }

    def _sort_key(kv):
        (slab, sub), _ = kv
        # Shared first, i.e. the IO path, then slabs in numeric order.
        if slab == 'shared':
            return (0, shared_order.get(sub, 999), sub)
        return (1, int(slab.replace('slab', '')), slab_order.get(sub, 999), sub)

    groups: List[Dict] = []
    for (slab, sub), named_params in sorted(buckets.items(), key=_sort_key):
        if not named_params:
            continue
        params = [p for _, p in named_params]
        param_names = [name for name, _ in named_params]
        is_gqk = (sub == 'vattn_gqk')
        wd = 0.0 if is_gqk else weight_decay
        scale = (gqk_lr_scale if is_gqk else 1.0) * float(extra_lr_scales.get(sub, 1.0))
        base = base_lr * scale
        groups.append({
            'params': params,
            'lr': base,
            'weight_decay': wd,
            'name': f'{slab}/{sub}',
            'base_lr': base,
            'is_gqk': is_gqk,
            'lr_scale': scale,
            'param_names': param_names,
        })
    return groups


@torch.no_grad()
def slab_stats(optimizer) -> Dict[str, Dict[str, float]]:
    """Per-slab aggregate of ||grad||2 and ||weight||2 across all its
    subgroups. One device->host sync per slab, NOT per param: scalar
    partial sums are reduced inside torch before `.item()`, so total
    overhead is O(num_slabs) syncs, not O(num_params). No flat-vector
    allocations.

    Returns {slab_name: {'g': float, 'w': float, 'n': int}}.
    """
    # Per (slab, device): list of scalar square-sums, still tensors on that device.
    partial: Dict[Tuple[str, torch.device], Dict[str, list]] = {}
    counts: Dict[str, int] = {}
    for g in optimizer.param_groups:
        name = g.get('name', '')
        slab = name.split('/', 1)[0] if '/' in name else name
        counts[slab] = counts.get(slab, 0)
        for p in g['params']:
            dev = p.device
            key = (slab, dev)
            ent = partial.setdefault(key, {'g': [], 'w': []})
            ent['w'].append(p.detach().float().pow(2).sum())
            g_sq = _grad_sq_sum(optimizer, p)
            if g_sq is not None:
                ent['g'].append(g_sq if g_sq.device == dev else g_sq.to(dev))
            counts[slab] += 1
    # Reduce the per-device partials, move the reduced scalar to CPU,
    # sum across devices, then call .item() once per slab.
    out: Dict[str, Dict[str, float]] = {}
    per_slab_dev: Dict[str, Dict[str, list]] = {}
    for (slab, _dev), ent in partial.items():
        d = per_slab_dev.setdefault(slab, {'g': [], 'w': []})
        if ent['w']:
            d['w'].append(torch.stack(ent['w']).sum().to('cpu', non_blocking=False))
        if ent['g']:
            d['g'].append(torch.stack(ent['g']).sum().to('cpu', non_blocking=False))
    for slab, d in per_slab_dev.items():
        w_sum = torch.stack(d['w']).sum() if d['w'] else torch.zeros(())
        g_sum = torch.stack(d['g']).sum() if d['g'] else torch.zeros(())
        out[slab] = {
            'w': float(w_sum.sqrt().item()),
            'g': float(g_sum.sqrt().item()),
            'n': counts.get(slab, 0),
        }
    return out


def _snapshot(optimizer):
    """Capture per-parameter pre-step copies for exact post-step delta.

    Also captures pre-step gradient square-sums so `slab_stats` /
    `group_stats` can be called AFTER `optimizer.step()` and still
    report meaningful `g` values. This matters for two configurations:

      1. Atom: ingests grads into `_ext_grads`, then `.step()` clears
         that buffer. After step, `p.grad` is None and `_ext_grads` is
         empty — naive `||p.grad||` reports 0.
      2. LowPrecGradAccumulator: post-accumulate hooks drain `p.grad`
         into a separate buffer during backward; `p.grad` is None
         throughout the rest of the step.

    The snapshot stashes a `_grad_snapshot` attribute on the optimizer
    that maps `id(p) -> g_sq_sum_tensor` (scalar tensor on the param's
    device, lifetime: until the next snapshot). `slab_stats` /
    `group_stats` consult it when present.
    """
    snap = []
    grad_snap: Dict[int, torch.Tensor] = {}
    for g in optimizer.param_groups:
        name = g.get('name', '')
        slab = name.split('/', 1)[0] if '/' in name else name
        for p in g['params']:
            snap.append((slab, name, p, p.detach().clone()))
            g_t = _live_grad(optimizer, p)
            if g_t is not None:
                grad_snap[id(p)] = g_t.detach().float().pow(2).sum()
    optimizer._grad_snapshot = grad_snap
    return snap


def _live_grad(optimizer, p: torch.nn.Parameter) -> Optional[torch.Tensor]:
    """Best-effort fetch of the current grad for `p` from any of the
    places the various code paths store it: the standard `.grad`
    attribute, Atom's `_ext_grads` ingestion buffer."""
    if p.grad is not None:
        return p.grad
    ext = getattr(optimizer, '_ext_grads', None)
    if ext is not None:
        g = ext.get(p)
        if g is not None:
            return g
    return None


def _grad_sq_sum(optimizer, p: torch.nn.Parameter) -> Optional[torch.Tensor]:
    """Pre-step grad ||·||² snapshot if available, else live grad."""
    snap = getattr(optimizer, '_grad_snapshot', None)
    if snap is not None:
        s = snap.get(id(p))
        if s is not None:
            return s
    g = _live_grad(optimizer, p)
    if g is None:
        return None
    return g.detach().float().pow(2).sum()


def slab_updates(snapshot) -> Dict[str, Dict[str, float]]:
    """Exact per-slab update ratio from a pre-step snapshot.

    The ratio is L1-over-L1:  sum(|Δw_i|) / sum(|w_i|).

    This measures the LITERAL fraction of aggregate weight mass that
    moved this step. It is linear in the number of parameters that
    actually moved (e.g. ATOM's `commit_frac`), so at `commit_frac=0.5`
    the displayed `upd/w` is exactly half of the per-committed
    per-parameter relative step — which is what "actual amount changed"
    means.

    (The legacy L2-over-L2 ratio we replaced, `||Δw||_2 / ||w||_2`,
    compresses that relationship to sqrt(commit_frac) and so overstated
    the effective move whenever the optimizer sparsely rejected blocks;
    it also mixed magnitudes of very differently-sized tensors less
    faithfully than L1 does.)

    Returns {slab: {'dw': sum(|Δw|), 'w': sum(|w_after|), 'upd_w': dw/w}}.
    """
    partial: Dict[Tuple[str, torch.device], Dict[str, list]] = {}
    for slab, _group_name, p, p_before in snapshot:
        dev = p.device
        key = (slab, dev)
        ent = partial.setdefault(key, {'dw': [], 'w': []})
        p_after = p.detach().float()
        p_prev = p_before.detach().float()
        ent['dw'].append((p_after - p_prev).abs().sum())
        ent['w'].append(p_after.abs().sum())

    per_slab_dev: Dict[str, Dict[str, list]] = {}
    for (slab, _dev), ent in partial.items():
        d = per_slab_dev.setdefault(slab, {'dw': [], 'w': []})
        if ent['dw']:
            d['dw'].append(torch.stack(ent['dw']).sum().to('cpu', non_blocking=False))
        if ent['w']:
            d['w'].append(torch.stack(ent['w']).sum().to('cpu', non_blocking=False))

    out: Dict[str, Dict[str, float]] = {}
    for slab, d in per_slab_dev.items():
        dw_sum = torch.stack(d['dw']).sum() if d['dw'] else torch.zeros(())
        w_sum = torch.stack(d['w']).sum() if d['w'] else torch.zeros(())
        w_val = float(w_sum.item())
        dw_val = float(dw_sum.item())
        out[slab] = {
            'dw': dw_val,
            'w': w_val,
            'upd_w': (dw_val / w_val) if w_val > 0 else 0.0,
        }
    return out


@torch.no_grad()
def group_stats(optimizer) -> Dict[str, Dict[str, float]]:
    """Per-group aggregate: ||grad||, ||weight||, param count, numel."""
    out: Dict[str, Dict[str, float]] = {}
    for g in optimizer.param_groups:
        name = g.get('name', 'unnamed')
        w_terms = []
        g_terms = []
        numel = 0
        nparams = 0
        for p in g['params']:
            nparams += 1
            numel += int(p.numel())
            w_terms.append(p.detach().float().pow(2).sum())
            g_sq = _grad_sq_sum(optimizer, p)
            if g_sq is not None:
                g_terms.append(g_sq)
        w_norm = float(torch.stack(w_terms).sum().sqrt().item()) if w_terms else 0.0
        g_norm = float(torch.stack(g_terms).sum().sqrt().item()) if g_terms else 0.0
        out[name] = {'w': w_norm, 'g': g_norm, 'params': nparams, 'numel': numel}
    return out


def group_updates(snapshot) -> Dict[str, Dict[str, float]]:
    """Exact per-group update ratio from a pre-step snapshot.

    Same L1/L1 semantics as `slab_updates` — see that docstring.
    """
    partial: Dict[Tuple[str, torch.device], Dict[str, list]] = {}
    for _slab, group_name, p, p_before in snapshot:
        dev = p.device
        key = (group_name, dev)
        ent = partial.setdefault(key, {'dw': [], 'w': []})
        p_after = p.detach().float()
        p_prev = p_before.detach().float()
        ent['dw'].append((p_after - p_prev).abs().sum())
        ent['w'].append(p_after.abs().sum())

    per_group_dev: Dict[str, Dict[str, list]] = {}
    for (group_name, _dev), ent in partial.items():
        d = per_group_dev.setdefault(group_name, {'dw': [], 'w': []})
        if ent['dw']:
            d['dw'].append(torch.stack(ent['dw']).sum().to('cpu', non_blocking=False))
        if ent['w']:
            d['w'].append(torch.stack(ent['w']).sum().to('cpu', non_blocking=False))

    out: Dict[str, Dict[str, float]] = {}
    for group_name, d in per_group_dev.items():
        dw_sum = torch.stack(d['dw']).sum() if d['dw'] else torch.zeros(())
        w_sum = torch.stack(d['w']).sum() if d['w'] else torch.zeros(())
        w_val = float(w_sum.item())
        dw_val = float(dw_sum.item())
        out[group_name] = {
            'dw': dw_val,
            'w': w_val,
            'upd_w': (dw_val / w_val) if w_val > 0 else 0.0,
        }
    return out