train_runtime.py

from __future__ import annotations

import os
from dataclasses import asdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import math

import torch
from torch.utils.checkpoint import checkpoint as torch_checkpoint

from inference import ChunkCarry, step


def ckpt_chunk(model, core_in, valid_mask, write_mask, carry, reset_mask=None):
    """Gradient-checkpointed chunked forward. Flattens ChunkCarry to a
    tuple of tensors so torch_checkpoint's tensor-only contract is
    satisfied, then reconstructs it. None-valued carry entries, i.e.
    the first chunk, are passed through as None and become zero-seed
    inside the per-slab scan primitives.
    """
    slab_names = tuple(model.core.slab_names)
    n = len(slab_names)

    # Layout of the flat_carry tuple (tensor-only contract for
    # torch.utils.checkpoint, which forbids non-tensor arguments):
    #   [0   .. n   )  -> rmu_m         per slab
    #   [n   .. 2n  )  -> vattn_m       per slab
    #   [2n  .. 3n  )  -> mix_hist      per slab
    #   [3n  .. 4n  )  -> predictor_pred per slab (PCSC)
    #   [4n  .. 5n  )  -> dsl_state per slab (Thought DSL)
    #   [5n]           -> xmem shared state (single tensor or None)
    #   [5n+1]         -> container (single tensor or None)
    carry_pred = carry.predictor_pred or {s: None for s in slab_names}
    carry_dsl = carry.dsl_state or {s: None for s in slab_names}

    def _run(core_in_, *flat_carry):
        rmu_m =    {s: flat_carry[i]           for i, s in enumerate(slab_names)}
        vattn_m =  {s: flat_carry[n + i]       for i, s in enumerate(slab_names)}
        mix_hist = {s: flat_carry[2 * n + i]   for i, s in enumerate(slab_names)}
        pred_in =  {s: flat_carry[3 * n + i]   for i, s in enumerate(slab_names)}
        dsl_in =   {s: flat_carry[4 * n + i]   for i, s in enumerate(slab_names)}
        xmem = flat_carry[5 * n]
        container_state = flat_carry[5 * n + 1]
        c = ChunkCarry(rmu_m=rmu_m, vattn_m=vattn_m,
                       mix_hist=mix_hist, xmem=xmem,
                       container=container_state,
                       predictor_pred=pred_in,
                       dsl_state=dsl_in)
        out_vol, new_c = step(
            model, core_in_, valid_mask, write_mask, c, reset_mask=reset_mask,
        )
        # PCSC/DSL aux losses are set as side effects on `model` by
        # step(). They live on the current chunk graph; surface them
        # through the checkpoint output so torch.checkpoint rewires them
        # correctly during backward recompute.
        pred_loss = model._predictor_step_loss
        dsl_exec_loss = model._dsl_exec_step_loss
        dsl_carry_loss = model._dsl_carry_step_loss
        dsl_src_loss = model._dsl_src_step_loss
        new_pred = new_c.predictor_pred or {s: None for s in slab_names}
        new_dsl = new_c.dsl_state or {s: None for s in slab_names}
        out_flat = tuple(new_c.rmu_m[s]    for s in slab_names) + \
                   tuple(new_c.vattn_m[s]  for s in slab_names) + \
                   tuple(new_c.mix_hist[s] for s in slab_names) + \
                   tuple(new_pred[s]       for s in slab_names) + \
                   tuple(new_dsl[s]        for s in slab_names) + \
                   (new_c.xmem,) + \
                   (new_c.container,) + \
                   (pred_loss,) + \
                   (dsl_exec_loss,) + \
                   (dsl_carry_loss,) + \
                   (dsl_src_loss,)
        return (out_vol,) + out_flat

    flat_in = (tuple(carry.rmu_m[s]    for s in slab_names) +
               tuple(carry.vattn_m[s]  for s in slab_names) +
               tuple(carry.mix_hist[s] for s in slab_names) +
               tuple(carry_pred[s]     for s in slab_names) +
               tuple(carry_dsl[s]      for s in slab_names) +
               (carry.xmem,) +
               (carry.container,))
    results = torch_checkpoint(_run, core_in, *flat_in, use_reentrant=False)
    out_vol = results[0]
    new_carry = ChunkCarry(
        rmu_m=   {s: results[1 + i]               for i, s in enumerate(slab_names)},
        vattn_m= {s: results[1 + n + i]           for i, s in enumerate(slab_names)},
        mix_hist={s: results[1 + 2 * n + i]       for i, s in enumerate(slab_names)},
        xmem=results[1 + 5 * n],
        container=results[1 + 5 * n + 1],
        predictor_pred={s: results[1 + 3 * n + i] for i, s in enumerate(slab_names)},
        dsl_state={s: results[1 + 4 * n + i] for i, s in enumerate(slab_names)},
    )
    predictor_loss = results[1 + 5 * n + 2]
    dsl_exec_loss = results[1 + 5 * n + 3]
    dsl_carry_loss = results[1 + 5 * n + 4]
    dsl_src_loss = results[1 + 5 * n + 5]
    return out_vol, new_carry, predictor_loss, dsl_exec_loss, dsl_carry_loss, dsl_src_loss


def _box_w() -> int:
    # Try to follow the current terminal width when available, but keep
    # deterministic safe bounds for redirected logs / no-tty runs.
    cols_s = os.environ.get('COLUMNS', '')
    try:
        cols = int(cols_s)
    except (TypeError, ValueError):
        cols = 120
    # Box uses "║ " + content + " ║", so the inner width should be a bit smaller.
    return max(78, min(140, cols - 2))


def _box_row(content: str) -> str:
    """Left-pads content inside a box row to exactly dynamic box width."""
    w = _box_w()
    if len(content) > w:
        content = content[:w - 1] + '…'
    return f"║ {content.ljust(w - 2)} ║"


def _box_top() -> str:
    w = _box_w()
    return '╔' + '═' * w + '╗'


def _box_div() -> str:
    w = _box_w()
    return '╠' + '═' * w + '╣'


def _box_end() -> str:
    w = _box_w()
    return '╚' + '═' * w + '╝'


def _fmt_upd(upd_w: float) -> str:
    """Human-readable relative update size per step."""
    ppm = upd_w * 1e6
    pct = upd_w * 100.0
    return f"upd/w={ppm:.2f} ppm ({pct:.5f}%/step)"


def _fmt_numel(n: int) -> str:
    """Compact numel formatter for log rows."""
    if n >= 1_000_000:
        return f'{n / 1_000_000:.1f}M'
    if n >= 1_000:
        return f'{n / 1_000:.1f}K'
    return str(n)


# Shared-layer rows shown in the per-step log box. Each tuple is
# (category, sub, group_name), where group_name is the optimizer
# param-group key produced by `optimizer._group_key` / `build_groups`.
# The order follows forward execution: input -> core_out -> decoder
# head -> container heads. Rows whose group is absent, e.g. fp32-only
# setups that disable a subgroup, are silently skipped.
SHARED_LAYER_ROWS: List[Tuple[str, str, str]] = [
    ('INPUT',     'norm',           'shared/input_norm'),       # LeaRNormND between convs
    ('INPUT',     'kan',            'shared/input_kan'),
    ('INPUT',     'content',        'shared/input_content'),
    ('INPUT',     'depth',          'shared/input_ctx_dw'),     # depthwise conv
    ('INPUT',     'point',          'shared/input_ctx_pw'),     # pointwise conv
    ('INPUT',     'other',          'shared/input_other'),
    ('CORE_OUT',  'proj',           'shared/core_out_proj'),
    ('DECODER',   'embed',          'shared/decoder_embedding'),
    ('DECODER',   'proj',           'shared/decoder_emb_proj'), # emb projection
    ('DECODER',   'norm',           'shared/decoder_norm'),     # layer norm
    ('DECODER',   'other',          'shared/decoder_other'),
]


def collect_shared_layer_stats(
    grp_stats: Dict[str, Dict[str, float]],
    grp_upd_stats: Dict[str, Dict[str, float]],
) -> List[Tuple[str, str, float, float, float, int]]:
    """Project the optimizer's per-group stats onto the named shared
    layers shown in the log box. Order follows `SHARED_LAYER_ROWS`.

    Returns a list of `(category, sub, g, w, upd_w, numel)`. Empty rows
    (missing group, or numel=0) are dropped so the box stays compact.
    """
    rows: List[Tuple[str, str, float, float, float, int]] = []
    for cat, sub, key in SHARED_LAYER_ROWS:
        ent = grp_stats.get(key)
        if not ent or int(ent.get('numel', 0)) == 0:
            continue
        upd = grp_upd_stats.get(key, {}).get('upd_w', 0.0)
        rows.append((cat, sub, float(ent['g']), float(ent['w']),
                     float(upd), int(ent['numel'])))
    return rows


def render_log(
    step: int,
    loss: float,
    lr: float,
    N: int,
    row_len: int,
    tok_per_sec: float,
    grad_tokens_total: int,
    seq_len: int,
    grid: Tuple[int, int, int],
    stats: Dict[str, Dict[str, float]],
    grp_stats: Dict[str, Dict[str, float]],
    actual_update_stats: Optional[Dict[str, Dict[str, float]]],
    learnorm_stats: Optional[Dict[str, float]],
    per_task_losses: Dict[str, float],
    container_k: int,
    container_runtime_stats: Optional[Dict[str, Dict[str, float]]],
    n_tasks: int,
    ckpt_saved: bool,
    skipped: bool,
    val_loss: Optional[float] = None,
    xmem_retention: Optional[float] = None,
) -> str:
    """Multi-line box report of one training step."""
    W = 85  # Fixed width for all lines
    row_text_w = W - 2

    def _top() -> str:
        return '╔' + '═' * W + '╗'

    def _end() -> str:
        return '╚' + '═' * W + '╝'

    def _div() -> str:
        return '╠' + '═' * W + '╣'

    def _div_label(label: str) -> str:
        # Center the label with brackets inside the divider, total width
        # W+2 including the borders.
        inner = '[' + label + ']'  # label with brackets
        total_fill = W - len(inner)  # space to fill with ═
        left = total_fill // 2
        right = total_fill - left
        return '╠' + '═' * left + inner + '═' * right + '╣'

    def _short_div() -> str:
        return '╠═' + ' ' * (W - 1) + '║'

    def _row(text: str) -> str:
        # Pad text to fit within W, leaving space for ║ on both sides
        # plus one extra space for alignment.
        content = text[:W-3]
        return '║ ' + content + ' ' * (W - 2 - len(content)) + ' ║'

    def _wrap_pipe_line(text: str) -> List[str]:
        """Wrap a `a | b | c` style line to row width without truncation."""
        if len(text) <= row_text_w:
            return [text]
        parts = [p.strip() for p in text.split('|')]
        out: List[str] = []
        cur = ""
        for p in parts:
            item = p if not cur else f"{cur} | {p}"
            if len(item) <= row_text_w:
                cur = item
            else:
                if cur:
                    out.append(cur)
                    cur = p
                else:
                    out.append(p[:row_text_w])
                    cur = ""
        if cur:
            out.append(cur)
        return out if out else [""]

    def _pack_metric_lines(items: List[str], extra_tail: str, max_lines: Optional[int] = None) -> List[str]:
        """Greedy pack of `a | b | c` items. If max_lines is None, do not cap rows."""
        if max_lines is not None and max_lines < 1:
            return [extra_tail] if extra_tail else []
        lines: List[str] = []
        cur = ""
        for item in items:
            cand = item if not cur else f"{cur} | {item}"
            if len(cand) <= row_text_w:
                cur = cand
                continue
            if cur:
                lines.append(cur)
                cur = item
            else:
                lines.append(item[:row_text_w])
                cur = ""
        if cur:
            lines.append(cur)
        if not lines:
            lines = [""]

        if max_lines is not None and len(lines) > max_lines:
            head = lines[:max_lines - 1]
            tail_items = lines[max_lines - 1:]
            tail = " | ".join(x for x in tail_items if x)
            if len(tail) > row_text_w:
                tail = tail[:row_text_w - 1] + '…'
            lines = head + [tail]

        if extra_tail:
            last = lines[-1]
            cand = extra_tail if not last else f"{last} | {extra_tail}"
            if len(cand) <= row_text_w:
                lines[-1] = cand
            elif max_lines is None or len(lines) < max_lines:
                lines.append(extra_tail)
            else:
                reserve = len(extra_tail) + 3
                keep = max(0, row_text_w - reserve - 1)
                prefix = last[:keep].rstrip()
                if prefix:
                    lines[-1] = f"{prefix}… | {extra_tail}"
                else:
                    lines[-1] = extra_tail[:row_text_w]
        if max_lines is not None:
            while len(lines) < max_lines:
                lines.append("")
            return lines[:max_lines]
        return lines

    def _stat(name: str, g: float, w: float, upd_w: float) -> str:
        return f"{name:<11s} g={g:.1e} | w={w:.1e} | upd={upd_w*1e6:.1f}ppm"

    def _slab(name: str, g: float, w: float, upd_w: float, suffix: str = "") -> str:
        base = f"{name:<11s} g={g:.1e} | w={w:.1e} | upd={upd_w*1e6:.1f}ppm"
        if suffix:
            base += f" | {suffix}"
        return base

    lines = [_top()]

    # STEP header + losses
    tasks_items = [f"{k}={v:.3f}" for k, v in per_task_losses.items()]
    if val_loss is not None:
        tasks_items.append(f"val={val_loss:.3f}")
    tasks_str = ' | '.join(tasks_items)
    finite = 'yes' if not skipped else 'NO'
    ckpt_str = 'yes' if ckpt_saved else 'no'
    header = (f"STEP {step} | lss={loss:.4f} | lr={lr:.2e} | "
              f"N={N} | row={row_len} | Tokens={grad_tokens_total} | {tok_per_sec:.0f} tok/s")
    for hline in _wrap_pipe_line(header):
        lines.append(_row(hline))
    metric_lines = _pack_metric_lines(
        tasks_items,
        extra_tail=f"finite={finite} | ckpt={ckpt_str}",
        max_lines=None,
    )
    for line in metric_lines:
        lines.append(_row(line))
    lines.append(_div())

    packed_str = 'yes' if n_tasks > 0 else 'no'
    g0, g1, g2 = grid

    # --- INPUT block ---
    lines.append(_row(f"INPUT | seq={seq_len} | grid=({g0},{g1},{g2}) | packed={packed_str} | reset=yes"))
    lines.append(_short_div())

    input_names = {
        'shared/input_norm': 'NORM',
        'shared/input_kan': 'KAN.Proj',
        'shared/input_content': 'Proj',
        'shared/input_ctx_dw': 'Ctx.Depth',
        'shared/input_ctx_pw': 'Ctx.Point',
    }
    for key, name in input_names.items():
        if key not in grp_stats:
            continue
        ent = grp_stats[key]
        upd = actual_update_stats.get(key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
        lines.append(_row(_stat(name, ent['g'], ent['w'], upd)))

    # --- CONTAINER block ---
    if container_k > 0:
        lines.append(_div_label('CONTAINER'))
        cont_names = {
            'shared/container_alpha': 'R. Head',
            'shared/container_slab_gate': 'Sl. Gate',
            'shared/container_neuron_gate': 'N. Gate',
            'shared/container_tau_sa1': 'TXSSM.fst',
            'shared/container_tau_sa2': 'TXSSM.lng',
            'shared/container_tau_sa3': 'TXSSM.col',
        }
        for key, name in cont_names.items():
            if key not in grp_stats:
                continue
            ent = grp_stats[key]
            upd = actual_update_stats.get(key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            lines.append(_row(_stat(name, ent['g'], ent['w'], upd)))
        # Optional self-skip container trunk layers (early-exit style).
        layer_keys = sorted(
            [k for k in grp_stats if k.startswith('shared/container_layer_') and 'container_layer_skip_' not in k],
            key=lambda x: int(x.split('container_layer_')[1]),
        )
        layer_on_fracs = []
        if container_runtime_stats and '__container__' in container_runtime_stats:
            layer_on_fracs = container_runtime_stats['__container__'].get('layer_on_fracs', [])
        for i, key in enumerate(layer_keys, start=1):
            ent = grp_stats[key]
            upd = actual_update_stats.get(key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            on = 'on'
            if i - 1 < len(layer_on_fracs):
                v = layer_on_fracs[i - 1]
                on = 'on' if (float(v.item()) if torch.is_tensor(v) else float(v)) > 0.0 else 'off'
            lines.append(_row(_slab(f"C. Layer {i}", ent['g'], ent['w'], upd, on)))

    # Helper to aggregate slab updates from all subgroups.
    def _slab_upd(slab_name: str) -> float:
        if not actual_update_stats:
            return 0.0
        # Sum updates from all groups belonging to this slab:
        # slab_name/other, slab_name/norm, etc.
        total_dw = 0.0
        total_w = 0.0
        prefix = slab_name + '/'
        for key, vals in actual_update_stats.items():
            if key.startswith(prefix):
                total_dw += vals.get('dw', 0.0)
                total_w += vals.get('w', 0.0)
        return (total_dw / total_w) if total_w > 0 else 0.0

    # --- SLABS block ---
    lines.append(_div_label('SLABS'))
    slab_keys = sorted([s for s in stats if s.startswith('slab')],
                       key=lambda s: int(s.replace('slab', '')))
    for i, s in enumerate(slab_keys):
        ent = stats[s]
        idx = s.replace('slab', '')
        upd = _slab_upd(s) if actual_update_stats else \
              (lr * ent['g'] / ent['w']) if ent['w'] > 0 else 0.0
        suffix = ""
        if container_runtime_stats and s in container_runtime_stats:
            c = container_runtime_stats[s]
            def _f(v) -> float:
                return float(v.item()) if torch.is_tensor(v) else float(v)
            active_frac = max(0.0, min(1.0, _f(c.get('active_frac', 0.0))))
            if active_frac <= 0.0:
                active = 'off'
            elif active_frac >= 1.0:
                active = 'on'
            else:
                active = f'{100.0 * active_frac:.1f}%'
            off_pct = 100.0 * _f(c.get('neuron_off_frac', 0.0))
            suffix = f"act={active} off={off_pct:.1f}%"
        elif container_k > 0:
            suffix = "off=0%"
        lines.append(_row(_slab(f"SLAB{idx}", ent['g'], ent['w'], upd, suffix)))
        # NORM row
        norm_key = f'{s}/norm'
        if norm_key in grp_stats:
            norm_ent = grp_stats[norm_key]
            norm_upd = actual_update_stats.get(norm_key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            lines.append(_row(_slab(f"NORM{idx}", norm_ent['g'], norm_ent['w'], norm_upd)))
        # PCSC predictor row
        pred_key = f'{s}/predictor'
        if pred_key in grp_stats:
            pred_ent = grp_stats[pred_key]
            pred_upd = actual_update_stats.get(pred_key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            lines.append(_row(_slab(f"PRED{idx}", pred_ent['g'], pred_ent['w'], pred_upd)))
        # VAttn row
        vattn_key = f'{s}/vattn'
        if vattn_key in grp_stats:
            vattn_ent = grp_stats[vattn_key]
            vattn_upd = actual_update_stats.get(vattn_key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            lines.append(_row(_slab(f"VAttn{idx}", vattn_ent['g'], vattn_ent['w'], vattn_upd)))
        lines.append(_short_div())

    # --- XMEM block (only when xmem groups exist) ---
    xmem_decay_key = 'shared/xmem_decay'
    has_xmem = xmem_decay_key in grp_stats or any(
        k.endswith('/xmem') for k in grp_stats
    )
    if has_xmem:
        ret_str = f'retain={xmem_retention:.3f}' if xmem_retention is not None else ''
        lines.append(_div_label('XMEM'))
        if xmem_decay_key in grp_stats:
            ent = grp_stats[xmem_decay_key]
            upd = actual_update_stats.get(xmem_decay_key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            lines.append(_row(_slab('Decay', ent['g'], ent['w'], upd, ret_str)))
        for s in slab_keys:
            xk = f'{s}/xmem'
            if xk not in grp_stats:
                continue
            idx = s.replace('slab', '')
            ent = grp_stats[xk]
            upd = actual_update_stats.get(xk, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
            lines.append(_row(_stat(f'XMem{idx}', ent['g'], ent['w'], upd)))

    # --- OUT block (core_out) ---
    lines.append(_div_label('OUT'))
    if 'shared/core_out_proj' in grp_stats:
        ent = grp_stats['shared/core_out_proj']
        upd = actual_update_stats.get('shared/core_out_proj', {}).get('upd_w', 0.0) if actual_update_stats else 0.0
        lines.append(_row(_stat('projection', ent['g'], ent['w'], upd)))

    # --- DECODER block ---
    lines.append(_div_label('DECODER'))
    decoder_names = [
        ('shared/decoder_norm', 'NORM'),
        ('shared/decoder_embedding', 'embedding'),
        ('shared/decoder_emb_proj', 'projection'),
    ]
    for key, name in decoder_names:
        if key not in grp_stats:
            continue
        ent = grp_stats[key]
        upd = actual_update_stats.get(key, {}).get('upd_w', 0.0) if actual_update_stats else 0.0
        lines.append(_row(_stat(name, ent['g'], ent['w'], upd)))

    lines.append(_end())
    return '\n'.join(lines)


def render_dsl_log(
    dsl_stats: Optional[Dict[str, Dict[str, object]]],
) -> str:
    """Render a compact Thought DSL program box.

    The raw program tensors are [steps, B, lanes], so this view shows:
      * trace receiver/source vocabulary for each slab,
      * answer/edit/state RMS values,
      * the exact chunk-level tensor program words from batch row 0.
    """
    if not dsl_stats:
        return ''
    W = 85

    def _top() -> str:
        return '╔' + '═' * W + '╗'

    def _end() -> str:
        return '╚' + '═' * W + '╝'

    def _div_label(label: str) -> str:
        inner = '[' + label + ']'
        total_fill = W - len(inner)
        left = total_fill // 2
        right = total_fill - left
        return '╠' + '═' * left + inner + '═' * right + '╣'

    def _row(text: str) -> str:
        content = text[:W - 3]
        return '║ ' + content + ' ' * (W - 2 - len(content)) + ' ║'

    def _f(v) -> float:
        if torch.is_tensor(v):
            return float(v.detach().float().mean().cpu())
        return float(v)

    def _op_name(i: int) -> str:
        return ('scan', 'mean', 'diff', 'read')[max(0, min(3, int(i)))]

    def _op_ch(i: int) -> str:
        return ('s', 'm', 'd', 'r')[max(0, min(3, int(i)))]

    def _alias(i: int) -> str:
        alphabet = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
        if i < len(alphabet):
            return alphabet[i]
        return f'@{i}'

    def _short_name(name: str) -> str:
        return (name
                .replace('slab', 'S')
                .replace('cml_', 'c')
                .replace('bank_', 'b')
                .replace('vattn_', 'va_')
                .replace('_proj', 'p'))

    def _wrap_rows(prefix: str, text: str, width: int = 80) -> List[str]:
        rows: List[str] = []
        cur = prefix
        for part in text.split(' '):
            if not part:
                continue
            cand = f'{prefix}{part}' if cur == prefix else f'{cur} {part}'
            if len(cand) <= width:
                cur = cand
                continue
            rows.append(cur)
            cur = f'{prefix}{part}'
        rows.append(cur)
        return rows

    lines = [_top(), _row('THOUGHT DSL | internal tensor-program log')]
    for slab in sorted(dsl_stats.keys(), key=lambda s: int(s.replace('slab', '')) if s.startswith('slab') else 999):
        st = dsl_stats[slab]
        prog = st.get('program')
        names = list(st.get('trace_names', []))
        if prog is None:
            continue
        src = prog['source_id'].detach().cpu()      # [S,B,N]
        op = prog['op_id'].detach().cpu()
        src_p = prog['source_prob'].detach().float().cpu()
        op_p = prog['op_prob'].detach().float().cpu()
        ret = prog['retention'].detach().float().cpu()
        gate = prog['lane_gate'].detach().float().cpu()
        gate_p = prog.get('lane_prob', gate).detach().float().cpu()
        carry = prog.get('carry_gate', torch.zeros_like(gate)).detach().float().cpu()
        carry_p = prog.get('carry_prob', carry).detach().float().cpu()
        expose = prog.get('expose_gate', torch.zeros_like(gate)).detach().float().cpu()
        expose_p = prog.get('expose_prob', expose).detach().float().cpu()
        ans = prog.get('answer_rms', None)
        ans_mean = float(ans.detach().float().mean().cpu()) if torch.is_tensor(ans) else 0.0
        S, B, N = src.shape
        lines.append(_div_label(slab.upper()))
        lines.append(_row(
            f"sources={int(_f(st.get('sources', 0)))} lanes={N} steps={S} "
            f"edit={_f(st.get('edit_rms', 0)):.1e} ctx={_f(st.get('context_rms', 0)):.1e} "
            f"bp={_f(st.get('bp_rms', 0)):.1e} wg={_f(st.get('wg_rms', 0)):.1e}"
        ))
        lines.append(_row(
            f"state={_f(st.get('state_rms', 0)):.1e} ans={ans_mean:.1e}"
            f" use={_f(st.get('exec_loss', 0)):.2f}/{_f(st.get('carry_loss', 0)):.2f}"
            f"/{_f(st.get('src_loss', 0)):.2f}"
        ))
        if names:
            vocab = ' '.join(f'{_alias(i)}={_short_name(n)}' for i, n in enumerate(names[:32]))
            for row in _wrap_rows('recv ', vocab):
                lines.append(_row(row))
        active_sources = prog.get('active_sources', None)
        if torch.is_tensor(active_sources):
            active_sources = active_sources.detach().float().cpu()
        for si in range(S):
            active = gate[si] > 0.5
            active_count = int(active.sum().item())
            carry_count = int(((carry[si] > 0.5) & active).sum().item())
            expose_count = int(((expose[si] > 0.5) & active).sum().item())
            b0_active = int(active[0].sum().item()) if B > 0 else 0
            if torch.is_tensor(active_sources):
                src_n_avg = float(active_sources[si].mean().item())
                src_n_str = f' src_n={src_n_avg:.2f}'
            else:
                src_n_str = ''
            if active_count > 0:
                src_flat = src[si][active].reshape(-1)
                op_flat = op[si][active].reshape(-1)
                src_counts = torch.bincount(src_flat, minlength=max(1, len(names))).float()
                op_counts = torch.bincount(op_flat, minlength=4).float()
                top_src = int(src_counts.argmax().item())
                top_op = int(op_counts.argmax().item())
                top_src_name = names[top_src] if 0 <= top_src < len(names) else f'#{top_src}'
                lines.append(_row(
                    f"step{si}: act={active_count}/{B * N} b0={b0_active}/{N} car={carry_count} exp={expose_count} "
                    f"top={_alias(top_src)}:{_short_name(top_src_name)} op={_op_name(top_op)} "
                    f"p={float(gate_p[si].mean()):.2f}/{float(carry_p[si].mean()):.2f}/{float(expose_p[si].mean()):.2f}"
                    f"{src_n_str}"
                ))
            else:
                lines.append(_row(
                    f"step{si}: noop act=0/{B * N} b0=0/{N} p={float(gate_p[si].mean()):.2f}/{float(carry_p[si].mean()):.2f}/{float(expose_p[si].mean()):.2f}"
                    f"{src_n_str}"
                ))
            lane_words = [
                (
                    f"{_alias(int(src[si, 0, lane]))}{_op_ch(int(op[si, 0, lane]))}"
                    f"{'>' if float(carry[si, 0, lane]) > 0.5 else ''}"
                    f"{'!' if float(expose[si, 0, lane]) > 0.5 else ''}"
                )
                if float(gate[si, 0, lane]) > 0.5 else '--'
                for lane in range(N)
            ]
            prog_line = ','.join(lane_words)
            for row in _wrap_rows(f'p{si}b0 ', prog_line):
                lines.append(_row(row))
            var = prog.get('variable_rms', None)
            exp = prog.get('expose_rms', None)
            pred = prog.get('predicate_mean', None)
            if torch.is_tensor(var) and torch.is_tensor(exp) and torch.is_tensor(pred):
                var_v = var[si, 0].detach().float().cpu()
                exp_v = exp[si, 0].detach().float().cpu()
                pred_v = pred[si, 0].detach().float().cpu()
                detail = ' '.join(
                    f'{lane}:var={float(var_v[lane]):.1e}/exp={float(exp_v[lane]):.1e}/pred={float(pred_v[lane]):.2f}'
                    for lane in range(min(N, 8))
                )
                for row in _wrap_rows('q ', detail):
                    lines.append(_row(row))
    lines.append(_end())
    return '\n'.join(lines)


def mem_mb(
    batch_size: int,
    seq_len: int,
    dec_hidden: int,
    core_hidden: int,
    params_total: int,
    train: bool = True,
) -> Dict[str, float]:
    # Rough estimate: params + grads + Adam states + a few activation buffers.
    bytes_per = 4.0
    params_mb = params_total * bytes_per / (1024.0 * 1024.0)
    grads_mb = params_mb if train else 0.0
    adam_mb = params_mb * 2.0 if train else 0.0

    # Activations, very rough: keep several core-volume / ray tensors across 3 slabs.
    elems_dec = batch_size * seq_len * dec_hidden
    elems_core = batch_size * seq_len * core_hidden
    act_elems = (4 * elems_dec) + (24 * elems_core)
    acts_mb = act_elems * bytes_per / (1024.0 * 1024.0)
    total_mb = params_mb + grads_mb + adam_mb + acts_mb
    return {
        "params_mb": params_mb,
        "grads_mb": grads_mb,
        "adam_mb": adam_mb,
        "acts_mb_est": acts_mb,
        "total_mb_est": total_mb,
    }


def save_ckpt(path: Path, model, optimizer: torch.optim.Optimizer, step: int, epoch: int, loss: float,
              data_pos: Optional[Dict] = None):
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        'step': step,
        'epoch': epoch,
        'loss': loss,
        'model_state': model.state_dict(),
        'optimizer_state': optimizer.state_dict(),
        'rmu_cfg': asdict(model.rmu_cfg),
        'dec_cfg': {
            **asdict(model.dec_cfg),
            'dtype': str(model.dec_cfg.dtype),
        },
        'container_cfg': (
            {
                'allow_slab_skip': bool(getattr(model.container, 'allow_slab_skip', False)),
                'max_unit_skip_frac': float(getattr(model.container, 'max_unit_skip_frac', 0.2)),
                'cont_layers': int(getattr(model.container, 'cont_layers', 1)),
            }
            if hasattr(model, 'container') else None
        ),
        'data_pos': data_pos,
    }
    torch.save(payload, path)


def msg_stream(dl):
    """Cross-iteration generator that repeatedly restarts the DataLoader.

    Retained for backward-compat. New code should use TrackedDataStream
    to enable resume-from-position.
    """
    while True:
        for mb in dl:
            yield from mb


class TrackedDataStream:
    """Iterator over a DataLoader that tracks consumed chars/tokens.

    Supports:
      * `skip_chars`: skip N chars/tokens at start (for resume). Partial
        skip is supported for str items and List[int] token sequences.
        Messages (List[dict]) are skipped whole-item only.
      * `start_epoch`: initial value for DataLoader-pass counter; used to
        seed DistributedSampler.set_epoch so resume uses the same shuffle
        as the saved run.
      * Sampler.set_epoch is called automatically on every DataLoader
        restart (one full pass through the dataset). Do NOT call
        sampler.set_epoch externally when using this stream.

    Counter `chars_in_epoch` is reset on each DataLoader restart and
    represents chars yielded BEFORE the most recent pull (counted in the
    resume point so skip(chars_in_epoch) lands exactly on the next item).
    """

    def __init__(self, dl, sampler=None, start_epoch: int = 0,
                 skip_chars: int = 0):
        self.dl = dl
        self.sampler = sampler
        self.dl_epoch = int(start_epoch)
        self.chars_in_epoch = 0
        self._skip = int(skip_chars)
        self._iter = None

    @staticmethod
    def _size(item) -> int:
        if isinstance(item, str):
            return len(item)
        if isinstance(item, (list, tuple)):
            if not item:
                return 0
            first = item[0]
            if isinstance(first, int):
                return len(item)
            if isinstance(first, dict):
                return sum(
                    len(m.get('content', '')) for m in item
                    if isinstance(m, dict) and isinstance(m.get('content'), str)
                )
        return 1

    def _gen(self):
        while True:
            if self.sampler is not None:
                self.sampler.set_epoch(self.dl_epoch)
            self.chars_in_epoch = 0
            skip = self._skip
            self._skip = 0  # only the first pass skips
            for mb in self.dl:
                for item in mb:
                    sz = self._size(item)
                    if skip >= sz and skip > 0:
                        skip -= sz
                        self.chars_in_epoch += sz
                        continue
                    if skip > 0:
                        if isinstance(item, str) or (
                            isinstance(item, (list, tuple)) and item
                            and isinstance(item[0], int)
                        ):
                            item = item[skip:]
                            self.chars_in_epoch += skip
                            skip = 0
                            sz = self._size(item)
                        else:
                            # messages: no partial slice; drop whole item
                            self.chars_in_epoch += sz
                            skip = max(0, skip - sz)
                            continue
                    self.chars_in_epoch += sz
                    yield item
            self.dl_epoch += 1

    def __iter__(self):
        if self._iter is None:
            self._iter = self._gen()
        return self

    def __next__(self):
        if self._iter is None:
            self._iter = self._gen()
        return next(self._iter)


def lr_at(step_idx: int, lr: float, lr_min: float, warmup_steps: int, max_steps: int) -> float:
    """Linear warmup -> long hold -> smooth late decay.

    Shape:
      1) warmup: [0, lr] over warmup_steps
      2) hold:   keep lr until ~5/7 of total steps
      3) decay:  polynomial (power=1.5) from lr to lr_min until max_steps
    """
    if step_idx <= warmup_steps and warmup_steps > 0:
        return lr * step_idx / max(1, warmup_steps)
    hold_end = max(warmup_steps + 1, int((5.0 / 7.0) * max_steps))
    if step_idx <= hold_end:
        return lr
    decay_total = max(1, max_steps - hold_end)
    t = min(1.0, (step_idx - hold_end) / decay_total)
    # Slower than linear at the beginning of decay, then catches up late.
    decay = (1.0 - t) ** 1.5
    return lr_min + (lr - lr_min) * decay