train.py

# (c) 2026 Ivan K (srose69, SimpleRose). Licensed under AGPL-3.0-only.
# Viper-LLM: Volumetric Language Model with Triangle Cross-Scan State Modelling.
# See TRADEMARKS.md for project naming and origin policy.

import argparse
import contextlib
import json
import os
import subprocess
import sys
import time
from dataclasses import asdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np

# expandable_segments=True lets the CUDA caching allocator grow a single
# segment instead of carving fixed blocks, which eliminates the
# fragmentation-induced OOMs that hit on 8-12 GB GPUs around step 2-3
# (free memory exists but no contiguous block of the requested size).
# Set BEFORE torch is imported so the allocator picks it up.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.autograd.graph import save_on_cpu
from datasets import Dataset, load_dataset, load_from_disk
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from transformers import AutoTokenizer

from decoder import DecoderConfig, ViperTiedDecoder
from cml import CoreConfig, Pipeline
from input_proj import InputProj
from inference import step, init_carry
from container import Container
from packing import (
    factorize_near_cube, triangle_sample_n, pack_messages_to_batch,
    apply_fim_to_packed, apply_nat_to_packed, PackedBatch,
)
from dtypes import (
    DTYPE_NAMES, parse_compute_dtype, autocast_ctx, is_low_precision,
    LowPrecGradAccumulator, OffloadedOptim, OffloadedAdamW,
)
from atom.atom import ATOM
from runtime_data import (
    FIM_WARMUP_FRAC,
    FIM_MIN_SPAN_FRAC,
    FIM_MAX_SPAN_FRAC,
)
from converter.main import (
    _get_hf_id,
    _to_messages,
    _to_str_ids,
    _load_text,
    _ensure_messages,
)
from optimizer import (
    count_params,
    build_groups,
    slab_stats,
    _snapshot,
    slab_updates,
    group_stats,
    group_updates,
)
from train_runtime import (
    ckpt_chunk, render_log, mem_mb, save_ckpt, lr_at, msg_stream,
    TrackedDataStream, render_dsl_log,
)


def _load_validation_dataset(ds_path: Path, dataset_type: str, text_column: str = 'auto'):
    """Load validation dataset from files/dirs with 'validation' in name.
    
    Searches for files/directories containing 'validation' (case-insensitive).
    Returns None if no validation files found.
    """
    if not ds_path.is_dir():
        return None
    
    # Find files/dirs with 'validation' in name
    val_items = []
    for item in ds_path.iterdir():
        if 'validation' in item.name.lower():
            val_items.append(item)
    
    if not val_items:
        return None
    
    # Load based on file type
    parquet_files = [f for f in val_items if f.is_file() and f.suffix.lower() == '.parquet']
    if parquet_files:
        return load_dataset('parquet', data_files=[str(f) for f in sorted(parquet_files)], split='train')
    
    # Check for arrow dataset directory
    for item in val_items:
        if item.is_dir():
            try:
                return load_from_disk(str(item))
            except Exception:
                pass
    
    # Try text files
    txt_files = [f for f in val_items if f.is_file() and f.suffix.lower() == '.txt']
    if txt_files:
        text = txt_files[0].read_text(encoding='utf-8')
        return Dataset.from_dict({'text': [text]})
    
    return None


def _validation_indices_path(ds_path: Path) -> Path:
    return ds_path / 'validation_indices.pt'


def _validation_auto_parquet_path(ds_path: Path) -> Path:
    return ds_path / 'validation_auto.parquet'


def _load_validation_indices(ds_path: Path) -> List[int]:
    idx_path = _validation_indices_path(ds_path)
    if not idx_path.is_file():
        return []
    payload = torch.load(idx_path, map_location='cpu', weights_only=False)
    indices = payload.get('indices', [])
    return [int(i) for i in indices]


def _exclude_validation_rows(ds, ds_path: Path):
    held_out = _load_validation_indices(ds_path)
    if not held_out:
        return ds
    n = len(ds)
    drop = set(i for i in held_out if 0 <= i < n)
    if not drop:
        return ds
    keep = [i for i in range(n) if i not in drop]
    return ds.select(keep)


def _maybe_create_validation_split(
    ds,
    ds_path: Path,
    val_fraction: float,
    val_seed: int,
    is_main: bool,
):
    """Auto-create held-out validation parquet plus index sidecar.

    Returns `(train_ds, val_ds, created_now)`.
    """
    if not ds_path.is_dir():
        return ds, None, False
    if _load_validation_dataset(ds_path, dataset_type='messages') is not None:
        return _exclude_validation_rows(ds, ds_path), _load_validation_dataset(ds_path, dataset_type='messages'), False
    if not (0.0 < float(val_fraction) < 1.0):
        return ds, None, False

    n = len(ds)
    if n < 2:
        return ds, None, False

    n_val = max(1, int(round(n * float(val_fraction))))
    n_val = min(n_val, n - 1)
    gen = torch.Generator()
    gen.manual_seed(int(val_seed))
    perm = torch.randperm(n, generator=gen).tolist()
    val_idx = sorted(perm[:n_val])
    train_idx = sorted(perm[n_val:])
    val_ds = ds.select(val_idx)
    train_ds = ds.select(train_idx)
    val_path = _validation_auto_parquet_path(ds_path)
    idx_path = _validation_indices_path(ds_path)
    val_ds.to_parquet(str(val_path))
    torch.save({'indices': val_idx, 'seed': int(val_seed), 'fraction': float(val_fraction)}, idx_path)
    if is_main:
        print(f'[data] validation auto-created file={val_path.name} size={len(val_ds)} '
              f'fraction={float(val_fraction):.6f}')
    return train_ds, val_ds, True


def _dist_barrier(local_rank: int) -> None:
    if not dist.is_available() or not dist.is_initialized():
        return
    backend = dist.get_backend()
    if backend == 'nccl':
        dist.barrier(device_ids=[int(local_rank)])
        return
    dist.barrier()


def _activation_offload_ctx(mem_device: str):
    """Autograd saved-tensor offload for activations.

    Only CPU is supported here intentionally: PyTorch provides a stable,
    correct path via save_on_cpu(). Cross-GPU saved-tensor offload would need
    custom hooks and tends to be much more fragile/perf-sensitive.
    """
    dev = (mem_device or '').strip().lower()
    if dev == 'cpu':
        return save_on_cpu(pin_memory=True)
    return contextlib.nullcontext()


def _activation_offload_enabled(args) -> bool:
    # save_on_cpu() and checkpoint(reentrant/non-reentrant) can disagree on
    # saved tensor metadata in this graph; keep the safe path by disabling
    # activation offload when core checkpointing is active.
    return bool((args.optim_mem_device or '').strip().lower() == 'cpu' and not bool(args.checkpoint_core))


def _auto_omp_threads(nproc: int) -> int:
    total = os.cpu_count() or 1
    if nproc <= 0:
        return 1
    # Keep one thread per rank at minimum, but avoid oversubscribing CPU
    # badly on multi-GPU launches. Cap to a modest value because packing
    # and tokenization already use Python/DataLoader workers separately.
    return max(1, min(8, total // nproc))


class ViperLanguageModel(torch.nn.Module):
    def __init__(
        self,
        rmu_cfg: CoreConfig,
        dec_cfg: DecoderConfig,
        loss_chunk_size: int,
        checkpoint_core: bool = True,
        num_slabs: int = 3,
        io_device: str = 'cuda:0',
        slab_devices: Tuple[str, ...] = None,
        container_k: int = 0,
        container_layers: int = 1,
        container_skip_slabs: bool = False,
        container_skip_units: float = 0.2,
    ):
        super().__init__()
        if (rmu_cfg.grid_g, rmu_cfg.grid_h, rmu_cfg.grid_w) != (dec_cfg.grid_g, dec_cfg.grid_h, dec_cfg.grid_w):
            raise ValueError("grid mismatch between RMU and decoder")
        if num_slabs < 1:
            raise ValueError(f"num_slabs must be >= 1, got {num_slabs}")
        if slab_devices is not None and len(slab_devices) != num_slabs:
            raise ValueError(
                f"slab_devices length {len(slab_devices)} != num_slabs {num_slabs}"
            )

        self.rmu_cfg = rmu_cfg
        self.dec_cfg = dec_cfg
        self.loss_chunk_size = loss_chunk_size
        self.checkpoint_core = checkpoint_core
        self.num_slabs = num_slabs
        self.io_device = torch.device(io_device)

        slab_names = tuple(f"slab{i + 1}" for i in range(num_slabs))
        self.core_in_proj = InputProj(
            dec_cfg.hidden_dim,
            rmu_cfg.channels,
            grid=(rmu_cfg.grid_g, rmu_cfg.grid_h, rmu_cfg.grid_w),
        )
        self.core = Pipeline(
            rmu_cfg, slab_names=slab_names, slab_devices=slab_devices,
        )
        self.core_out_proj = torch.nn.Conv3d(rmu_cfg.channels, dec_cfg.hidden_dim, kernel_size=1, bias=True)
        # Sub-Xavier fan-in init for the volume->dec_hidden projection that
        # feeds decoder._norm/LM head. Pytorch's Kaiming-uniform default is
        # ~sqrt(2)x larger than fan-in 1.0, which (post-LeaRNorm in decode)
        # over-shoots the unit-variance logit contract at init. gain=0.51
        # matches the rest of the model.
        _coreout_w = self.core_out_proj.weight
        _coreout_fanin, _ = torch.nn.init._calculate_fan_in_and_fan_out(_coreout_w)
        try:
            _coreout_gain = float(os.getenv("VIPER_CORE_OUT_GAIN", "0.9"))
        except Exception:
            _coreout_gain = 0.51
        _coreout_std = _coreout_gain / max(1.0, float(_coreout_fanin)) ** 0.5
        torch.nn.init.trunc_normal_(_coreout_w, mean=0.0, std=_coreout_std,
                                    a=-2.0 * _coreout_std, b=2.0 * _coreout_std)
        torch.nn.init.zeros_(self.core_out_proj.bias)
        self.decoder = ViperTiedDecoder(dec_cfg)
        # Container: optional per-chunk internal router. Disabled when
        # container_k==0 (absent attribute; `getattr(model, 'container',
        # None)` in _chunk_forward_with_carry returns None -> identity
        # path, bit-equivalent to pre-Container behavior).
        if container_k > 0:
            self.container = Container(
                slab_names=slab_names, ray_dim=rmu_cfg.ray_dim, K=container_k,
                cont_layers=container_layers,
                allow_slab_skip=container_skip_slabs,
                max_unit_skip_frac=container_skip_units,
            )

        # See `_enforce_conv3d_channels_last_3d` for why. Apply on init so
        # the layout is correct before the first forward; `_apply` below
        # re-applies after every subsequent `.to()`/`.cuda()` move.
        self._enforce_conv3d_channels_last_3d()

    def _enforce_conv3d_channels_last_3d(self) -> None:
        """Force all nn.Conv3d weight tensors into channels_last_3d strides.

        cuDNN on SM>=80 emits `∂L/∂W` for 3-D convolutions in channels_last_3d
        layout: strides `(C_in*D*H*W, 1, C_in*H*W, C_in*W, C_in)`. PyTorch's
        DDP reducer builds its bucket view from `param.stride()` snapshotted
        at wrap time; if the weight was stored as default-contiguous
        `(C_in*D*H*W, D*H*W, H*W, W, 1)` the grad stride disagrees and DDP
        logs:
            "Grad strides do not match bucket view strides. [...] may
             impair performance."
        plus performs an extra device-side copy per bucket per reduction.
        Aligning the PARAMETER storage with the grad layout both removes
        the warning and avoids that copy. Forward also goes through
        cuDNN's channels-last kernels which is the preferred fast path
        on Ampere+ for 1×1×1 point-wise convs that dominate the core.

        Subtlety addressed here:
          * `.to(memory_format=torch.channels_last_3d)` is a no-op for
            weights with a size-1 non-batch dim (depthwise kernels have
            `C_in_per_group=1`): channels_last_3d and contiguous are
            metadata-indistinguishable in that case. We explicitly build
            an `empty_strided` target and copy the data, so `stride()`
            returns exactly `(C*D*H*W, 1, C*H*W, C*W, C)`.
          * `.cuda()` / `.to(device=...)` DROPS this layout on singleton-
            dim tensors because the default `_apply` path clones with
            default memory format. This method is invoked after every
            `self._apply(...)` (see override below) so the layout holds
            across moves/dtype-casts.
        """
        # Enumerate every module with a 5-D `weight` tensor — this covers
        # `nn.Conv3d` AND the project's `MaskedCausalConv3d` (which is a
        # plain `nn.Module` that calls `F.conv3d` on its own weight, so
        # isinstance(nn.Conv3d) does NOT match it). Both hit the same
        # cuDNN backward path and emit channels_last_3d grad strides.
        for _m in self.modules():
            _w = getattr(_m, 'weight', None)
            if not isinstance(_w, torch.Tensor) or _w.dim() != 5:
                continue
            _wd = _w.data.contiguous()
            _, _C, _D, _H, _W = _wd.shape
            # Depthwise / grouped kernels (C_in == 1): autograd emits grad
            # in plain contiguous layout, so keep the param contiguous.
            if _C == 1:
                if _w.data.stride() != _wd.stride():
                    _m.weight.data = _wd
                continue
            _cl3d_strides = (_C * _D * _H * _W, 1, _C * _H * _W, _C * _W, _C)
            if _wd.stride() == _cl3d_strides:
                continue
            _new = torch.empty_strided(
                _wd.shape, _cl3d_strides, dtype=_wd.dtype, device=_wd.device,
            )
            _new.copy_(_wd)
            _m.weight.data = _new

    def _apply(self, fn, *args, **kwargs):
        # nn.Module._apply underlies .to(), .cuda(), .float(), etc.
        # Re-enforce the layout every time params are re-materialized so
        # `param.stride()` keeps matching the grad stride cuDNN produces.
        out = super()._apply(fn, *args, **kwargs)
        self._enforce_conv3d_channels_last_3d()
        return out

    @staticmethod
    def _make_write_mask(valid_mask: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
        # Write at every valid position, including BOS (pos 0) and EOS (last
        # valid token). Cross-sample isolation in single-sample legacy mode is a
        # non-issue; scan reset is handled by reset_mask, not write exclusion.
        return valid_mask.clone()

    @property
    def L_native(self) -> int:
        return self.rmu_cfg.grid_g * self.rmu_cfg.grid_h * self.rmu_cfg.grid_w

    def _forward_chunks(self, input_ids: torch.Tensor,
                         valid_mask: torch.Tensor,
                         write_mask: torch.Tensor,
                         reset_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Unified chunked training forward (handles N>=1 chunks).

        Row length L must be a positive multiple of L_native. For each
        chunk we project to core space, run `step`
        (which threads RMU/vAttn/mix carries, runs Container routing + update
        if attached, and honors `reset_mask` to zero scan decay at sample
        boundaries), and decode to seq. Seqs are concatenated along L.

        The per-chunk carry is a Python dataclass of tensors; gradients
        flow through the carry chain across chunks, giving full BPTT over
        the row. `checkpoint_core=1` wraps each chunk's core pass in
        `torch.utils.checkpoint` to trade forward-activation VRAM for
        backward-time recomputation of that chunk.
        """
        B, L = input_ids.shape
        Ln = self.L_native
        if L % Ln != 0 or L < Ln:
            raise ValueError(f'seq_len {L} must be a positive multiple of L_native={Ln}')
        N = L // Ln

        carry = init_carry(self)
        rays: List[torch.Tensor] = []
        # PCSC: accumulate per-chunk predictor MSE across the row.
        # Threaded as an explicit tensor through the checkpoint boundary
        # (see `ckpt_chunk`) so the backward recompute rewires it; the
        # non-checkpoint path reads it from `self._predictor_step_loss`.
        pred_loss_accum = torch.zeros((), device=input_ids.device, dtype=torch.float32)
        dsl_exec_loss_accum = torch.zeros((), device=input_ids.device, dtype=torch.float32)
        dsl_carry_loss_accum = torch.zeros((), device=input_ids.device, dtype=torch.float32)
        dsl_src_loss_accum = torch.zeros((), device=input_ids.device, dtype=torch.float32)
        pred_loss_chunks = 0
        dsl_loss_chunks = 0
        for ci in range(N):
            s, e = ci * Ln, (ci + 1) * Ln
            toks = input_ids[:, s:e]
            vm = valid_mask[:, s:e]
            wm = write_mask[:, s:e]
            rm = reset_mask[:, s:e] if reset_mask is not None else None
            # Single sync per chunk: if this chunk has no sample boundary,
            # null out reset_mask so all reset-aware paths (RayMix.step,
            # RMU.prepare, VAttn._build_gates) take their fast no-reset
            # branch without any further per-slab .item() calls. In
            # packed batches, only chunks that actually contain a sample
            # boundary need the slower reset-aware path.
            if rm is not None and not bool(rm.any()):
                rm = None
            volume, _ = self.decoder.encode(toks, valid_mask=vm)
            core_in = self.core_in_proj(volume)
            if self.training and self.checkpoint_core:
                (core_out_vol, carry, chunk_pred_loss,
                 chunk_dsl_exec_loss, chunk_dsl_carry_loss,
                 chunk_dsl_src_loss) = ckpt_chunk(
                    self, core_in, vm, wm, carry, reset_mask=rm,
                )
            else:
                core_out_vol, carry = step(
                    self, core_in, vm, wm, carry, reset_mask=rm,
                )
                chunk_pred_loss = self._predictor_step_loss
                chunk_dsl_exec_loss = self._dsl_exec_step_loss
                chunk_dsl_carry_loss = self._dsl_carry_step_loss
                chunk_dsl_src_loss = self._dsl_src_step_loss
            pred_loss_accum = pred_loss_accum + chunk_pred_loss.float()
            dsl_exec_loss_accum = dsl_exec_loss_accum + chunk_dsl_exec_loss.float()
            dsl_carry_loss_accum = dsl_carry_loss_accum + chunk_dsl_carry_loss.float()
            dsl_src_loss_accum = dsl_src_loss_accum + chunk_dsl_src_loss.float()
            pred_loss_chunks += 1
            dsl_loss_chunks += 1
            # Avoid forcing fp32 activations during normal mixed-precision
            # training. Only enforce dtype/device match when autocast is off
            # (e.g. direct probe callers), where Conv3d requires exact match.
            _core_out_w = self.core_out_proj.weight
            if core_out_vol.device != _core_out_w.device:
                core_out_vol = core_out_vol.to(device=_core_out_w.device)
            if (not torch.is_autocast_enabled()) and (core_out_vol.dtype != _core_out_w.dtype):
                core_out_vol = core_out_vol.to(dtype=_core_out_w.dtype)
            core_out = self.core_out_proj(core_out_vol)
            rays.append(self.decoder.vol_to_seq(core_out))
        # Mean over chunks for the row; the trainer multiplies by
        # `predictor_lambda` and adds to the total loss in `forward`.
        denom = float(max(1, pred_loss_chunks))
        self._predictor_row_loss = pred_loss_accum / denom
        dsl_denom = float(max(1, dsl_loss_chunks))
        self._dsl_exec_row_loss = dsl_exec_loss_accum / dsl_denom
        self._dsl_carry_row_loss = dsl_carry_loss_accum / dsl_denom
        self._dsl_src_row_loss = dsl_src_loss_accum / dsl_denom
        return torch.cat(rays, dim=1)

    def forward(self, input_ids: torch.Tensor,
                target_ids: Optional[torch.Tensor] = None,
                valid_mask: Optional[torch.Tensor] = None,
                write_mask: Optional[torch.Tensor] = None,
                reset_mask: Optional[torch.Tensor] = None,
                learnorm_lambda: float = 0.0,
                predictor_lambda: float = 0.0,
                vector_error_lambda: float = 0.0,
                dsl_exec_lambda: float = 0.0,
                dsl_carry_lambda: float = 0.0,
                dsl_source_lambda: float = 0.0,
                return_loss_parts: bool = False):
        """Training forward.

        Two modes:
          (a) Packed mode [preferred for throughput]: caller provides
              `valid_mask`, `write_mask`, `reset_mask`, `target_ids`
              pre-computed by `packing.pack_messages_to_batch`. All four
              tensors describe the row's multi-sample layout explicitly.
              Boundary labels in `target_ids` are already pad-id so
              `loss_chunked`'s `ignore_index=pad_id` silences cross-sample
              next-token predictions without any extra work here.

          (b) Legacy mode [backward compat]: if valid_mask / write_mask
              are omitted, they are derived from `lengths_from_tokens` on
              target_ids, producing the single-sample behavior used
              before packing. `reset_mask` stays None, so all scan paths
              behave bit-identically to the pre-packing training setup.
        """
        if target_ids is None:
            target_ids = input_ids
        B, L = input_ids.shape

        if valid_mask is None:
            lengths = self.decoder.lengths_from_tokens(target_ids)
            idx = torch.arange(L, device=input_ids.device).unsqueeze(0)
            valid_mask = idx < lengths.unsqueeze(1)
            if write_mask is None:
                write_mask = self._make_write_mask(valid_mask, lengths)
        elif write_mask is None:
            # valid provided but no write: derive a safe default from
            # valid_mask minus position 0 and the last valid position.
            lengths = valid_mask.sum(dim=1).long()
            write_mask = self._make_write_mask(valid_mask, lengths)

        seq = self._forward_chunks(input_ids, valid_mask, write_mask, reset_mask=reset_mask)
        ce_loss, learnorm_loss, vector_error_loss = self.decoder.ce_blocks(
            seq,
            target_ids,
            chunk_size=self.loss_chunk_size,
            return_parts=True,
            vector_error_lambda=vector_error_lambda,
        )
        # PCSC: small BYOL-style regression loss on the slab predictor
        # outputs. Without this term the predictor would only feel the
        # CE gradient through the wg modulation, which is too indirect
        # for it to actually fit `actual_first`. Coefficient stays small
        # (default 0.01) because the per-slab MSE is on an O(1)-scale
        # ray feature and CE on the LM head is on a similar scale.
        predictor_loss = self._predictor_row_loss.to(ce_loss.dtype)
        dsl_exec_loss = self._dsl_exec_row_loss.to(ce_loss.dtype)
        dsl_carry_loss = self._dsl_carry_row_loss.to(ce_loss.dtype)
        dsl_src_loss = self._dsl_src_row_loss.to(ce_loss.dtype)
        total_loss = (
            ce_loss
            + float(learnorm_lambda) * learnorm_loss
            + float(predictor_lambda) * predictor_loss
            + float(vector_error_lambda) * vector_error_loss
            + float(dsl_exec_lambda) * dsl_exec_loss
            + float(dsl_carry_lambda) * dsl_carry_loss
            + float(dsl_source_lambda) * dsl_src_loss
        )
        if return_loss_parts:
            return {
                'loss': total_loss,
                'ce_loss': ce_loss,
                'learnorm_loss': learnorm_loss,
                'predictor_loss': predictor_loss,
                'vector_error_loss': vector_error_loss,
                'dsl_exec_loss': dsl_exec_loss,
                'dsl_carry_loss': dsl_carry_loss,
                'dsl_src_loss': dsl_src_loss,
            }
        return total_loss


def _load_yaml_config(path: str) -> Dict[str, object]:
    """Load a YAML training config into a flat dict.

    The YAML file may use either the dotted CLI form or its bare names:
    e.g. both ``batch_size: 8`` and ``batch-size: 8`` are accepted (we
    normalize ``-`` to ``_``). Nested mappings are flattened with a dot
    separator (``data.dir`` -> ``dataset_dir`` style is the caller's job;
    we leave key names verbatim apart from the dash/underscore rule).
    """
    try:
        import yaml  # PyYAML
    except ImportError as exc:
        raise RuntimeError(
            f"--config={path!r} requires PyYAML. Install via `pip install pyyaml`."
        ) from exc
    with open(path, 'r', encoding='utf-8') as f:
        raw = yaml.safe_load(f)
    if raw is None:
        return {}
    if not isinstance(raw, dict):
        raise ValueError(f'YAML config {path!r} must be a top-level mapping')

    flat: Dict[str, object] = {}

    def _walk(prefix: str, node: object) -> None:
        if isinstance(node, dict):
            for k, v in node.items():
                key = str(k).replace('-', '_')
                full = f'{prefix}.{key}' if prefix else key
                if isinstance(v, dict):
                    _walk(full, v)
                else:
                    flat[full if prefix else key] = v
            return
        flat[prefix] = node

    _walk('', raw)
    # Drop any nested-style keys that didn't collapse to a top-level name —
    # only top-level keys map to argparse dest names. (We still surface
    # them in the warning so the user can tell.)
    return {k: v for k, v in flat.items() if '.' not in k}


def main():
    p = argparse.ArgumentParser()
    # --config must be the first parsed argument so its values seed the
    # parser defaults BEFORE the rest are parsed. Anything passed on the
    # CLI still overrides YAML (argparse default precedence).
    p.add_argument('--config', type=str, default='',
                   help='Path to a YAML training config. Every CLI flag (without '
                        'the leading --) is a valid YAML key, so the YAML can fully '
                        'replace command-line arguments. CLI flags still take '
                        'precedence over YAML when both are present.')
    p.add_argument('--tokenizer_dir', type=str, default='tools/gwen2_tokenizer')
    p.add_argument('--dataset_dir', type=str, default='datasets/smol-smoltalk_train')
    p.add_argument('--pt', type=str, default='',
                   help='Path/spec for PT corpus. Implies --dataset_type string '
                        'and replaces --dataset_dir. Do not pass --dataset_dir '
                        'explicitly together with --pt.')
    p.add_argument('--dataset_split', type=str, default='train',
                   help='Split name when loading Hugging Face datasets.')
    p.add_argument('--dataset_text_column', type=str, default='auto',
                   help='Text column to wrap into messages for plain-text datasets.')
    p.add_argument('--val_dataset', type=str, default='',
                   help='Optional explicit validation dataset path/spec. When unset, '
                        'training will auto-discover validation* alongside --dataset_dir, '
                        'or create a held-out validation_auto.parquet for local dataset dirs.')
    p.add_argument('--val_fraction', type=float, default=0.001,
                   help='Fraction of a local training dataset to hold out automatically '
                        'for validation when --val_dataset is unset and no validation* '
                        'artifact exists already. Must be in (0,1). Default 0.001.')
    p.add_argument('--val_seed', type=int, default=1337,
                   help='Deterministic seed for the automatic held-out validation split.')
    p.add_argument('--dataset_type', type=str, default='messages',
                   choices=['messages', 'string'],
                   help='messages: chat-template packing with boundaries. '
                        'string: plain PT token stream, no user/messages semantics.')
    p.add_argument(
        '--string_tokenize_otf',
        action=argparse.BooleanOptionalAction,
        default=True,
        help='For --dataset_type string with text datasets: tokenize on-the-fly in the '
             'packing loop (default: true). Use --no-string_tokenize_otf to pretokenize '
             'the full dataset via map() before training.',
    )
    p.add_argument(
        '--modes',
        type=str,
        default='pt',
        help='Comma-separated task modes. Supported: pt,nat,blind,rev,fim. '
             'Default: pt. Example: --modes pt,nat,fim',
    )
    p.add_argument('--ptmode', action='store_true',
                   help='Enable PT objective mode (alias for adding "pt" to --modes).')
    p.add_argument('--blindmode', action='store_true',
                   help='Enable blind objective mode (alias for adding "blind" to --modes).')
    p.add_argument('--revmode', action='store_true',
                   help='Enable reverse objective mode (alias for adding "rev" to --modes).')
    p.add_argument('--fimmode', action='store_true',
                   help='Enable FIM objective mode (alias for adding "fim" to --modes).')
    p.add_argument('--nat_weight', type=float, default=1.5,
                   help='Loss weight for the NAT task relative to PT (default 1.5). '
                        'Total gradient contribution: nat_w/(1+nat_w). '
                        'Raise above 1.0 to counteract the model coasting on the '
                        'easier PT objective when objectives are co-directional.')
    p.add_argument('--natmode', action='store_true',
                   help='Enable NAT (blind-tail prefix-completion) objective mode '
                        '(alias for adding "nat" to --modes). Closes the train/inference '
                        'gap of one-shot generate(): each packed sample gets a random '
                        'suffix replaced by PAD on input while target stays the real '
                        'tokens, training the regime where the model emits from ray '
                        'state alone. Equal weight with pt; no warmup ramp.')
    p.add_argument('--ckpt_dir', type=str, default='ckpt')

    p.add_argument('--device', type=str, default='cuda:0')
    p.add_argument('--devices', type=str, default='',
                   help='DDP launcher mode. Set to "cudas" to use all visible CUDA '
                        'GPUs in full data-parallel mode (one process per GPU).')
    p.add_argument('--seq_len', type=int, default=1024,
                   help='Native chunk width (== L_native). Volumetric grid '
                        '(g, h, w) is derived as the near-cube factorization '
                        'of seq_len, no manual grid knobs. Rows in training '
                        'are N*seq_len tokens where N is sampled per-batch '
                        'from 1..seq_len_mul (triangle-weighted). Pick '
                        'seq_len as a composite number (power of 2 works '
                        'cleanest): it factorizes exactly.')
    p.add_argument('--seq_len_mul', type=int, default=1,
                   help='Max chunks-per-row multiplier. row_len = N * seq_len '
                        'where N ~ Triangle{1..seq_len_mul} (p(N) ~ N). Mul=1 '
                        'means single-chunk training (carry unused). Mul>=2 '
                        'activates carry and Container routing. Worst-case '
                        'VRAM peaks at N=mul; dry-run allocates at that peak '
                        'on init so OOM surfaces immediately.')
    p.add_argument('--batch_size', type=int, default=8)
    p.add_argument('--grad_accum_steps', type=int, default=1,
                   help='Number of microbatches to accumulate before each optimizer step. '
                        'Effective global batch = batch_size * grad_accum_steps. '
                        'Memory stays close to one microbatch because DDP all-reduce is '
                        'suppressed on intermediate microsteps and grads accumulate in-place.')
    p.add_argument('--decoder_dim', type=int, default=32)
    p.add_argument('--factor_dim', type=int, default=0,
                   help='ALBERT-style factorized embedding rank; 0 disables.')
    p.add_argument('--core_dim', type=int, default=960)
    p.add_argument('--virtual_dim', type=float, default=0.0,
                   help='Virtual manifold expansion in percent over core_dim. '
                        '0 disables. Example: --virtual_dim 14 means 1024 -> ~1168.')
    p.add_argument('--num_slabs', type=int, default=3)
    p.add_argument('--tau_fast', type=str, default='auto',
                   help='TXSSM fast-branch half-life in tokens. "auto" => seq_len/4.')
    p.add_argument('--tau_slow', type=str, default='auto',
                   help='TXSSM slow-branch half-life in tokens. "auto" => seq_len.')
    p.add_argument('--core_devices', type=str, default='',
                   help='Comma-separated devices for slabs (e.g. "cuda:0,cuda:1,cuda:2"). '
                        'Empty = use --device for all slabs. Slabs are round-robin distributed.')
    p.add_argument('--max_steps', type=int, default=1000,
                   help='Number of optimization steps to run. On a fresh '
                        'run this is the absolute total. With --resume it '
                        'is ADDITIONAL: total_steps = ckpt_step + max_steps '
                        '(the LR schedule ends at the new total).')
    p.add_argument('--lr', type=float, default=1e-4)
    p.add_argument('--lr_min', type=float, default=-1.0,
                   help='Final LR reached by the late decay tail. Default = lr * 0.1.')
    p.add_argument('--dtype', type=str, default='fp32', choices=list(DTYPE_NAMES),
                   help='Compute dtype inside the RMU core forward (bound/siloid/scan '
                        'are saturating and safe in low precision). Params, optimizer '
                        'states and vocab log-sum-exp stay fp32. fp8_e4m3/fp8_e5m2 are '
                        'reserved names; see dtypes.py for the pending pack/unpack path.')
    p.add_argument('--param_dtype', type=str, default='auto', choices=['auto', 'fp32', 'fp16', 'bf16'],
                   help='Model parameter storage dtype. auto => follow --dtype for fp16/bf16, '
                        'else fp32.')
    p.add_argument('--dtype_acts', type=str, default='fp32', choices=['fp32', 'fp16', 'bf16'],
                   help='Force-compute dtype for bounded activations (siloid, bound). '
                        'autocast bf16 does NOT convert log1p/sqrt/cos (fp32 whitelist), so '
                        'setting this to bf16/fp16 disables autocast inside those activations '
                        'and keeps the tensors in half precision — actually halving activation '
                        'memory. Safe because the output envelope is bounded independent of |x|.')
    p.add_argument('--grads_dtype', type=str, default='fp32', choices=['fp32', 'fp16', 'bf16'],
                   help='Storage dtype for accumulated grads between backward passes in the '
                        'multi-task mixer. fp32 = default PyTorch behavior. bf16/fp16 = after '
                        'each task backward, param.grad is drained into a half-prec buffer and '
                        'upcast to fp32 only right before optimizer.step(). Halves the grad '
                        'footprint during the 3-task accumulation window.')
    p.add_argument('--optim_mem_device', type=str, default='',
                   help='Unified memory-offload target for training-state tensors. '
                        'When set, this is used as the default target for: '
                        '(1) low-precision grad accumulator buffers, '
                        '(2) optimizer master/state offload if --optim_device is not set, '
                        '(3) autograd saved-tensor offload when the target is cpu. '
                        'Examples: "cpu", "cuda:1". Empty = keep everything local unless '
                        'other explicit offload flags are set.')
    p.add_argument('--grads_device', type=str, default='',
                   help=argparse.SUPPRESS)
    p.add_argument('--optim', type=str, default='atom', choices=['adamw', 'atom'],
                     help='Optimizer to use: adamw (torch.optim.AdamW) or atom (Atom). Default: atom.')
    p.add_argument('--optim_device', type=str, default='',
                   help='Run optimizer master weights + states on this '
                        'device instead of the compute GPU. Biggest single VRAM win: '
                        '~12*P bytes removed from compute GPU for a P-param model. '
                        'Pair with --grads_device to keep the whole grad + optim pipeline '
                        'off the compute GPU. Works with both adamw and atom optimizers.')
    p.add_argument('--optim_dtype', type=str, default='fp32', choices=['fp32', 'fp16', 'bf16'],
                   help='Master parameter dtype for offloaded optimizer copies. '
                        'Primarily useful with --optim_device cpu to reduce RAM.')
    p.add_argument('--warmup_steps', type=int, default=50,
                   help='Linear LR warmup 0 -> lr over this many steps; then long hold and late smooth decay to lr_min.')
    p.add_argument('--gqk_lr_scale', type=float, default=0.05)
    p.add_argument('--gqk_warmup_steps', type=int, default=100)
    p.add_argument('--emb_proj_lr_scale', type=float, default=0.3,
                   help='LR multiplier for the decoder factor->hidden lift '
                        '(emb_proj). Sensitivity probe shows it is ~3-5x '
                        'more loss-sensitive than slab interior matrices '
                        'because it sits on the forward path AND on the '
                        'tied readout, so an unscaled LR there dominates '
                        'effective updates and starves the rest. Default '
                        '0.3 ≈ 1/3 of base LR; set 1.0 to disable.')
    p.add_argument('--decoder_norm_lr_scale', type=float, default=1.0,
                   help='LR multiplier for decoder.norm (LeaRNorm). Set <1 '
                        'if the LeaRNorm gain/beta drift dominates loss '
                        'descent. Default 1.0 (no scaling).')
    p.add_argument('--weight_decay', type=float, default=0.01)
    p.add_argument('--learnorm_lambda', type=float, default=0.05,
                   help='Aux loss weight for decoder LeaRNorm quality regularizer. '
                        'Total loss = CE + learnorm_lambda * learnorm_loss.')
    p.add_argument('--predictor_lambda', type=float, default=0.01,
                   help='PCSC: aux MSE weight for per-slab state-predictor '
                        '(pred vs actual_first, BYOL target detached). The '
                        'predictor also receives gradient from CE through the '
                        'write-gate modulation path; this aux term is what '
                        'forces it to actually predict instead of collapsing '
                        'to a constant that only minimizes the modulation '
                        'response. 0 disables (predictor still runs but is '
                        'driven only by CE via modulation, which usually '
                        'collapses).')
    p.add_argument('--vector_error_lambda', type=float, default=0.0,
                   help='Weight for embedding-space vector error objective. '
                        'Adds dense directional correction E[target]-E_p[pred] '
                        'on top of CE. 0 disables.')
    p.add_argument('--predictor_lr_scale', type=float, default=1.0,
                   help='LR multiplier for per-slab PCSC predictor groups '
                        '(state_predictor.* + surprise_gain).')
    p.add_argument('--grad_clip', type=float, default=1.0)
    p.add_argument('--save_every', type=int, default=100)
    p.add_argument('--log_every', type=int, default=10)
    p.add_argument('--update_stats_every', type=int, default=0,
                   help='Cadence (in optimizer steps) at which the per-group '
                        'L1-update ratio (`upd=…ppm` fields) is computed. '
                        'Computing the ratio requires snapshotting EVERY model '
                        'parameter before the step (~param-dtype × num_params '
                        'extra alloc + a full elementwise pass) — for a 372M '
                        'model that is ~700MB and ~30ms per step, dominating '
                        'wall time when log_every is small. 0 = follow '
                        'log_every (legacy behavior). >0 = independent '
                        'cadence; on non-snapshot logs the `upd` fields are '
                        'omitted from the rendered table.')
    p.add_argument('--loss_chunk_size', type=int, default=1024)
    p.add_argument('--checkpoint_core', type=int, default=1)
    p.add_argument('--num_workers', type=int, default=2)
    p.add_argument('--container_k', type=int, default=0,
                   help='Internal Container controller state dim (per batch row). 0 = '
                        'disabled (default, pre-Container behavior, bit-identical). '
                        'Recommended 512..2048 when enabled. Container acts as an '
                        'inter-chunk router: updates once per chunk, stopgrad at '
                        'chunk boundary (no cross-chunk BPTT), influences the NEXT '
                        'chunks per-slab write-amplitude and scan timescale biases. '
                        'Requires multi-chunk training rows, i.e. seq_len_mul > 1, '
                        'to see any routing signal.')
    p.add_argument('--skip_slabs', action='store_true',
                   help='Allow Container to hard-skip middle slabs per batch/chunk. '
                        'First and last slab stay always active.')
    p.add_argument('--skip_units', type=float, default=0.2,
                   help='Maximum fraction of per-slab units Container can hard-skip. '
                        'Range [0,1). Default 0.2.')
    p.add_argument('--container_lr_scale', type=float, default=1.0,
                   help='LR multiplier for all Container trainable groups.')
    p.add_argument('--cont_layers', type=int, default=1,
                   help='Container routing trunk depth. 1 = direct heads from state; >1 adds MLP trunk.')
    p.add_argument('--use_dsl', action=argparse.BooleanOptionalAction, default=False,
                   help='Physically instantiate and execute Thought DSL. '
                        'When false, DSL modules are not built even if dsl_lanes is non-zero.')
    p.add_argument('--dsl_lanes', type=int, default=0,
                   help='Thought DSL command lanes per slab. 0 disables the internal '
                        'typed tensor-program language.')
    p.add_argument('--dsl_rank', type=int, default=64,
                   help='Low-rank width used by Thought DSL queries/answers.')
    p.add_argument('--dsl_max_sources', type=str, default='16',
                   help='Maximum number of causal trace streams visible to each DSL '
                        'interpreter. Pass an integer or the string "all" to expose every '
                        'receiver published in the chunk-level registry; "all" resolves '
                        'to num_slabs * 8 at startup.')
    p.add_argument('--dsl_max_sources_per_exec', type=int, default=0,
                   help='Top-K hard cap on simultaneously firing receivers per (batch, '
                        'lane) per DSL step. 0 disables the cap (limited only by '
                        'dsl_max_sources). With a large source vocabulary this is the '
                        'knob that bounds peak DSL VRAM.')
    p.add_argument('--dsl_source_lambda', type=float, default=0.0,
                   help='L0-style regularizer on Hard Concrete source-open '
                        'probabilities. Encourages each DSL step to keep only '
                        'a few receivers open. 0 disables.')
    p.add_argument('--dsl_steps', type=int, default=2,
                   help='Number of recurrent internal program steps executed per slab.')
    p.add_argument('--dsl_log', action=argparse.BooleanOptionalAction, default=False,
                   help='Print a separate Thought DSL program/answer box under the main log.')
    p.add_argument('--dsl_lr_scale', type=float, default=0.5,
                   help='LR multiplier for per-slab Thought DSL parameters.')
    p.add_argument('--dsl_exec_lambda', type=float, default=0.0,
                   help='Tiny usage penalty on Thought DSL execute probabilities. '
                        'Encourages sparse/noop programs. 0 disables.')
    p.add_argument('--dsl_carry_lambda', type=float, default=0.0,
                   help='Tiny usage penalty on Thought DSL answer-carry probabilities. '
                        'Encourages carrying only useful receiver outputs. 0 disables.')
    p.add_argument('--xmem_size', type=float, default=0.0,
                   help='Total XMem long-term memory budget (MB, bf16) per batch '
                        'item for the shared bank. 0 = disabled (default). The N-dim memory tensor has '
                        'shape [B, mem_dim, A^N] where mem_dim defaults to '
                        'max(32, core_dim/8), N=6 axes (twice VAttn spatial). The axis '
                        'size A is solved from the budget: A = floor((MB·5e5/mem_dim)^(1/N)). '
                        'Tape footprint during training is approximately '
                        'xmem_size · batch_size · seq_len_mul MB; with defaults at '
                        'xmem_size=10 MB, batch=44, mul=4 that is ~1.8 GB extra '
                        'activation tape. Increase only when there is '
                        'headroom — the memory geometry scales exponentially in N, so '
                        'A drop below 4 around xmem_size~0.5 MB.')
    p.add_argument('--resume', type=str, default='',
                   help='Path to a checkpoint (.pt) to resume from. Loads model+optimizer '
                        'state and continues from the saved step counter. When set, the '
                        'model and optimizer state take precedence over their current '
                        'initialization; architectural shapes (rmu_cfg/dec_cfg) must still '
                        'match the CLI-provided ones.')
    p.add_argument('--dtype_audit_steps', type=int, default=0,
                   help='If >0, collect per-module output dtype stats for first N steps '
                        'and print top fp32 output producers (rank0).')
    # Two-pass parsing: first pull out --config, load YAML, and push its
    # values into parser defaults; then parse the full CLI so explicit
    # flags override YAML. This is the standard argparse+YAML pattern
    # that preserves --help and unknown-arg detection.
    pre_args, _ = p.parse_known_args()
    if pre_args.config:
        yaml_defaults = _load_yaml_config(pre_args.config)
        unknown = [k for k in yaml_defaults if k not in {a.dest for a in p._actions}]
        if unknown:
            raise ValueError(
                f"YAML config {pre_args.config!r} contains keys not matching any "
                f"CLI flag: {sorted(unknown)}"
            )
        p.set_defaults(**yaml_defaults)
    args = p.parse_args()
    # Resolve --dsl_max_sources: accept either a positive integer (as a
    # string) or the literal "all". "all" means "expose every published
    # receiver" and resolves to a compile-time cap that fits the actual
    # per-slab publication density. Each slab publishes ~16 receivers
    # (4 state streams + 4-5 bank executors + 4 vattn projections +
    # optional xmem + 2 trailing cml streams visible only to later
    # slabs); see `@inference.py:_process_chunk` publish() calls.
    # The `command` head and `source_slot` table are sized to this cap.
    _raw_msrc = str(getattr(args, 'dsl_max_sources', '16')).strip().lower()
    if _raw_msrc in ('all', 'auto', '*'):
        args.dsl_max_sources = max(16, int(args.num_slabs) * 16)
    else:
        try:
            args.dsl_max_sources = int(_raw_msrc)
        except ValueError as _exc:
            raise ValueError(
                f"--dsl_max_sources must be a positive integer or 'all', got {_raw_msrc!r}"
            ) from _exc
    if int(args.dsl_max_sources) <= 0:
        raise ValueError(f"--dsl_max_sources must be > 0, got {args.dsl_max_sources}")
    # Backward-compat alias: grads_device -> optim_mem_device.
    _legacy_grads_dev = str(getattr(args, 'grads_device', '') or '').strip()
    _optim_mem_dev = str(getattr(args, 'optim_mem_device', '') or '').strip()
    if _legacy_grads_dev and _optim_mem_dev and _legacy_grads_dev != _optim_mem_dev:
        raise ValueError(
            f'--grads_device={_legacy_grads_dev} conflicts with '
            f'--optim_mem_device={_optim_mem_dev}; use only optim_mem_device'
        )
    if not _optim_mem_dev and _legacy_grads_dev:
        args.optim_mem_device = _legacy_grads_dev
    args.grads_device = str(getattr(args, 'optim_mem_device', '') or '')
    # PT alias mode: reuse the existing string data-path.
    if args.pt:
        if '--dataset_dir' in sys.argv:
            raise ValueError('--pt excludes explicit --dataset_dir; pass only one data source')
        if '--dataset_type' in sys.argv and args.dataset_type != 'string':
            raise ValueError('--pt requires --dataset_type string (or omit --dataset_type)')
        args.dataset_dir = args.pt
        args.dataset_type = 'string'
    modes = _parse_modes(args.modes)
    if args.ptmode:
        modes.add('pt')
    if args.blindmode:
        modes.add('blind')
    if args.revmode:
        modes.add('rev')
    if args.fimmode:
        modes.add('fim')
    if args.natmode:
        modes.add('nat')
    if not modes:
        raise ValueError('no active training modes; enable at least one mode')
    # PT mode now supports both string and messages dataset types.
    # For messages, only assistant content is used (see _assistant_only_pt_ids in packing.py).
    # UX alias: allow --device cudas as a shorthand for full DDP launch.
    if args.device == 'cudas':
        args.devices = 'cudas'
        args.device = 'cuda:0'
    if args.devices and args.devices != 'cudas':
        raise ValueError(f'--devices supports only "cudas" (got {args.devices})')
    if not (0.0 <= args.skip_units < 1.0):
        raise ValueError(f'--skip_units must be in [0,1), got {args.skip_units}')
    if args.val_dataset and args.val_fraction <= 0:
        raise ValueError(f'--val_fraction must be > 0 when provided; got {args.val_fraction}')
    if not args.val_dataset and not (0.0 < args.val_fraction < 1.0):
        raise ValueError(f'--val_fraction must be in (0,1) for auto validation split; got {args.val_fraction}')

    if (args.device.startswith('cuda') or args.devices == 'cudas') and not torch.cuda.is_available():
        raise RuntimeError('CUDA unavailable')
    if args.devices == 'cudas' and 'WORLD_SIZE' not in os.environ:
        n = torch.cuda.device_count()
        if n < 2:
            raise RuntimeError('--devices cudas requested but fewer than 2 CUDA GPUs visible')
        cmd = ['torchrun', '--standalone', '--nproc_per_node', str(n), __file__] + sys.argv[1:]
        env = os.environ.copy()
        if not env.get('OMP_NUM_THREADS'):
            env['OMP_NUM_THREADS'] = str(_auto_omp_threads(n))
        print('[ddp] launching:', ' '.join(cmd), flush=True)
        if 'OMP_NUM_THREADS' in env and not os.environ.get('OMP_NUM_THREADS'):
            print(f'[ddp] auto OMP_NUM_THREADS={env["OMP_NUM_THREADS"]}', flush=True)
        raise SystemExit(subprocess.call(cmd, env=env))

    world_size = int(os.environ.get('WORLD_SIZE', '1'))
    distributed = world_size > 1
    rank = int(os.environ.get('RANK', '0'))
    local_rank = int(os.environ.get('LOCAL_RANK', '0'))
    is_main = (rank == 0)
    if world_size > 1 and not os.environ.get('OMP_NUM_THREADS'):
        os.environ['OMP_NUM_THREADS'] = str(_auto_omp_threads(world_size))
    if distributed:
        dist.init_process_group(backend='nccl')
        torch.cuda.set_device(local_rank)
        args.device = f'cuda:{local_rank}'
        if args.batch_size < world_size:
            raise ValueError(
                f'--batch_size={args.batch_size} is global in DDP mode and must be >= world_size={world_size}'
            )
        if args.batch_size % world_size != 0:
            raise ValueError(
                f'--batch_size={args.batch_size} must be divisible by world_size={world_size} in DDP mode'
            )
        global_batch_size = args.batch_size
        args.batch_size = args.batch_size // world_size
    else:
        global_batch_size = args.batch_size
    if is_main and distributed:
        print(f'[ddp] world_size={world_size} backend=nccl')
        print(f'[ddp] batch_size(global)={global_batch_size} -> batch_size(local)={args.batch_size}')
        print(
            f'[ddp] grad_accum_steps={args.grad_accum_steps} -> '
            f'effective_global_batch={global_batch_size * args.grad_accum_steps}'
        )
    elif is_main:
        print(
            f'[train] batch_size={args.batch_size} grad_accum_steps={args.grad_accum_steps} '
            f'-> effective_global_batch={args.batch_size * args.grad_accum_steps}'
        )
    if is_main and args.optim_mem_device:
        acts_note = ''
        if args.optim_mem_device.strip().lower() == 'cpu':
            acts_note = ' + acts(save_on_cpu)' if _activation_offload_enabled(args) else ' + acts(save_on_cpu disabled: checkpoint_core=1)'
        print(f'[train] optim_mem_device={args.optim_mem_device}{acts_note}')
    if distributed:
        print(f'[ddp] rank={rank} local_rank={local_rank} device={args.device}', flush=True)

    compute_dtype = parse_compute_dtype(args.dtype)
    if args.param_dtype == 'auto':
        if args.optim == 'atom' and not (args.optim_device and args.optim_device.strip().lower() == 'cpu'):
            # Atom keeps optimizer moments in parameter dtype on CUDA path.
            # Auto-lowering params to bf16/fp16 destabilizes updates.
            param_dtype = torch.float32
            param_dtype_name = 'fp32(auto-for-atom)'
        elif compute_dtype in (torch.float16, torch.bfloat16):
            param_dtype = compute_dtype
            param_dtype_name = args.dtype
        else:
            param_dtype = torch.float32
            param_dtype_name = 'fp32'
    else:
        param_dtype = {
            'fp32': torch.float32,
            'fp16': torch.float16,
            'bf16': torch.bfloat16,
        }[args.param_dtype]
        param_dtype_name = args.param_dtype
    if is_low_precision(compute_dtype):
        if is_main:
            print(f'[train] compute dtype = {args.dtype} (param dtype = {param_dtype_name})')

    # Bounded-activation override: force siloid/bound to run at the
    # requested dtype (bypasses autocast's fp32 promotion for log1p/cos
    # /sqrt). Applies module-wide via activations.set_acts_dtype.
    import activations as _acts
    _acts_dt = parse_compute_dtype(args.dtype_acts)
    _acts.set_acts_dtype(None if _acts_dt is torch.float32 else _acts_dt)
    if _acts_dt is not torch.float32:
        if is_main:
            print(f'[train] bounded activations forced to {args.dtype_acts} '
                  '(autocast disabled inside siloid/bound)')

    torch.backends.cudnn.benchmark = True
    device = torch.device(args.device)

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir, use_fast=True)
    if tokenizer.pad_token_id is None:
        raise RuntimeError(
            f'tokenizer at {args.tokenizer_dir} has no pad_token_id; '
            'please configure a dedicated pad token in tokenizer files.'
        )
    vocab_size = len(tokenizer)

    ds_spec = args.dataset_dir.strip()
    ds_path = Path(ds_spec)
    if ds_path.exists():
        try:
            ds = load_from_disk(str(ds_path))
            if is_main:
                print(f'[data] source=load_from_disk path={ds_path}')
        except Exception:
            if args.dataset_type == 'messages':
                if (not distributed) or rank == 0:
                    arrow_dir = _ensure_messages(
                        ds_spec, text_column=args.dataset_text_column, is_main=is_main,
                    )
                else:
                    arrow_dir = (ds_path / '.arrow_messages') if ds_path.is_dir() else ds_path.with_suffix(ds_path.suffix + '.arrow_messages')
                if distributed:
                    _dist_barrier(local_rank)
                ds = load_from_disk(str(arrow_dir))
                if is_main:
                    print(f'[data] source=arrow_messages path={arrow_dir}')
            else:
                # Check for parquet files first
                parquet_files = sorted(
                    f for f in ds_path.glob('*.parquet')
                    if 'validation' not in f.name.lower()
                ) if ds_path.is_dir() else []
                if parquet_files:
                    ds = load_dataset('parquet', data_files=[str(f) for f in parquet_files], split='train')
                    if is_main:
                        print(f'[data] source=parquet files={len(parquet_files)} path={ds_path}')
                else:
                    ds, txt_path = _load_text(ds_spec)
                    if is_main:
                        print(f'[data] source=plain_text path={txt_path}')
    else:
        hf_id = _get_hf_id(ds_spec) or ds_spec
        ds = load_dataset(hf_id, split=args.dataset_split)
        if is_main:
            print(f'[data] source=hf dataset={hf_id} split={args.dataset_split}')
    string_text_col: Optional[str] = None
    if args.dataset_type == 'messages':
        ds = _to_messages(ds, text_column=args.dataset_text_column)
    else:
        if args.string_tokenize_otf:
            if 'input_ids' in ds.column_names:
                pass
            else:
                col = args.dataset_text_column
                if col == 'auto':
                    if 'text' in ds.column_names:
                        col = 'text'
                    else:
                        str_cols = [
                            c for c, feat in ds.features.items()
                            if getattr(feat, 'dtype', None) == 'string'
                        ]
                        col = str_cols[0] if str_cols else ''
                if not col or col not in ds.column_names:
                    raise ValueError(
                        f'dataset has no input_ids/text column for string mode. '
                        f'columns={list(ds.column_names)}; pass --dataset_text_column <name> '
                        f'or disable OTF with --no-string_tokenize_otf.'
                    )
                string_text_col = str(col)
                keep_cols = [string_text_col]
                ds = ds.remove_columns([c for c in ds.column_names if c not in keep_cols])
        else:
            ds = _to_str_ids(ds, tokenizer, text_column=args.dataset_text_column)
    val_ds = None
    if args.val_dataset:
        val_spec = args.val_dataset.strip()
        val_path = Path(val_spec)
        if val_path.exists():
            try:
                val_ds = load_from_disk(str(val_path))
            except Exception:
                if val_path.is_dir():
                    parquet_files = sorted(val_path.glob('*.parquet'))
                    if parquet_files:
                        val_ds = load_dataset('parquet', data_files=[str(f) for f in parquet_files], split='train')
                    else:
                        val_ds, _ = _load_text(val_spec)
                else:
                    val_ds, _ = _load_text(val_spec)
        else:
            hf_val_id = _get_hf_id(val_spec) or val_spec
            val_ds = load_dataset(hf_val_id, split=args.dataset_split)
        if is_main and val_ds is not None:
            print(f'[data] validation explicit source={val_spec} size={len(val_ds)}')
    elif ds_path.is_dir():
        val_ds = _load_validation_dataset(ds_path, args.dataset_type, args.dataset_text_column)
        if val_ds is None:
            if distributed:
                if is_main:
                    ds, val_ds, _ = _maybe_create_validation_split(
                        ds, ds_path, args.val_fraction, args.val_seed, is_main,
                    )
                _dist_barrier(local_rank)
                if not is_main:
                    ds = _exclude_validation_rows(ds, ds_path)
                    val_ds = _load_validation_dataset(ds_path, args.dataset_type, args.dataset_text_column)
            else:
                ds, val_ds, _ = _maybe_create_validation_split(
                    ds, ds_path, args.val_fraction, args.val_seed, is_main,
                )
        else:
            ds = _exclude_validation_rows(ds, ds_path)
        if val_ds is not None and is_main:
            print(f'[data] validation loaded columns={list(val_ds.column_names)} size={len(val_ds)}')
    
    if is_main:
        print(f'[data] type={args.dataset_type} columns={list(ds.column_names)} size={len(ds)}')
        if args.dataset_type == 'string':
            if 'input_ids' in ds.column_names:
                print('[data] string tokenization=pretokenized input_ids')
            else:
                print(f'[data] string tokenization=otf text_column={string_text_col}')
        print(f'[modes] active={",".join(sorted(modes))}')

    # Grid auto-factorization: seq_len -> near-cube (g, h, w). No manual
    # grid knobs. seq_len defines L_native; row_len at training = N*seq_len
    # with N sampled per-batch via packing.triangle_sample_n.
    if args.seq_len_mul < 1:
        raise ValueError(f'--seq_len_mul must be >= 1, got {args.seq_len_mul}')
    if args.grad_accum_steps < 1:
        raise ValueError(f'--grad_accum_steps must be >= 1, got {args.grad_accum_steps}')
    grid_g, grid_h, grid_w = factorize_near_cube(args.seq_len)
    def _resolve_tau(v: str, default_val: float) -> float:
        s = str(v).strip().lower()
        if s in ('auto', ''):
            return float(default_val)
        t = float(s)
        if t <= 0:
            raise ValueError(f'tau must be > 0, got {v}')
        return float(t)
    tau_fast = _resolve_tau(args.tau_fast, max(1.0, float(args.seq_len) / 4.0))
    tau_slow = _resolve_tau(args.tau_slow, max(1.0, float(args.seq_len)))
    row_len_max = args.seq_len * args.seq_len_mul
    if is_main:
        print(f'[grid] seq_len={args.seq_len} -> (g,h,w)=({grid_g},{grid_h},{grid_w})  '
              f'row_len_max={row_len_max} (seq_len_mul={args.seq_len_mul})')

    v_pct = max(0.0, float(args.virtual_dim))
    v_dim = int(round(args.core_dim * (1.0 + v_pct / 100.0)))
    if v_dim < args.core_dim:
        v_dim = args.core_dim
    rmu_cfg = CoreConfig(
        channels=args.core_dim,
        ray_dim=args.core_dim,
        virtual_dim=v_dim,
        grid_g=grid_g,
        grid_h=grid_h,
        grid_w=grid_w,
        vocab_size=vocab_size,
        bos_token_id=tokenizer.bos_token_id if tokenizer.bos_token_id is not None else 1,
        eos_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 2,
        pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0,
        chunk_size=32,
        device=args.device,
        dtype=param_dtype,
        xmem_size=float(args.xmem_size),
        use_dsl=bool(args.use_dsl),
        dsl_lanes=int(args.dsl_lanes),
        dsl_rank=int(args.dsl_rank),
        dsl_max_sources=int(args.dsl_max_sources),
        dsl_max_per_exec=int(args.dsl_max_sources_per_exec),
        dsl_steps=int(args.dsl_steps),
        tau_fast=float(tau_fast),
        tau_slow=float(tau_slow),
    )
    if is_main and rmu_cfg.xmem_size > 0.0:
        geom = rmu_cfg.xmem_geometry()
        if geom is not None:
            mem_dim_x, axis_x, num_ax_x = geom
            tape_mb = rmu_cfg.xmem_size * args.batch_size * args.seq_len_mul
            print(
                f'[xmem] enabled: mem_dim={mem_dim_x} axes={num_ax_x} A={axis_x} '
                f'-> {axis_x ** num_ax_x} positions/bank '
                f'({rmu_cfg.xmem_size:.2f} MB per batch item total; '
                f'~{tape_mb:.0f} MB activation tape across the row)'
            )

    dec_cfg = DecoderConfig(
        vocab_size=vocab_size,
        hidden_dim=args.decoder_dim,
        grid_g=grid_g,
        grid_h=grid_h,
        grid_w=grid_w,
        bos_token_id=rmu_cfg.bos_token_id,
        eos_token_id=rmu_cfg.eos_token_id,
        pad_token_id=rmu_cfg.pad_token_id,
        dtype=param_dtype,
        vocab_chunk_size=args.loss_chunk_size,
        use_chunked_loss=True,
        factor_dim=args.factor_dim,
    )

    # Round-robin slab placement across --core_devices (pipeline parallel).
    if args.core_devices.strip():
        core_devs = [d.strip() for d in args.core_devices.split(',') if d.strip()]
        slab_devices = tuple(core_devs[i * len(core_devs) // args.num_slabs] for i in range(args.num_slabs))
    else:
        slab_devices = tuple(args.device for _ in range(args.num_slabs))

    model_base = ViperLanguageModel(
        rmu_cfg,
        dec_cfg,
        loss_chunk_size=args.loss_chunk_size,
        checkpoint_core=bool(args.checkpoint_core),
        num_slabs=args.num_slabs,
        io_device=args.device,
        slab_devices=slab_devices,
        container_k=args.container_k,
        container_layers=args.cont_layers,
        container_skip_slabs=bool(args.skip_slabs),
        container_skip_units=float(args.skip_units),
    )
    if param_dtype != torch.float32:
        model_base.to(dtype=param_dtype)
    # Place only IO-side modules on io device; core slabs are already pinned
    # to their per-slab devices by the pipeline constructor.
    model_base.core_in_proj.to(device)
    model_base.core_out_proj.to(device)
    model_base.decoder.to(device)
    if args.container_k > 0:
        model_base.container.to(device)
        if is_main:
            print(f'[container] enabled: K={args.container_k} on {device} '
                  f'(reservoir state_proj frozen, routing heads trainable; '
                  f'cont_layers={args.cont_layers}; '
                  f'skip_slabs={"on" if args.skip_slabs else "off"}; '
                  f'skip_units<={args.skip_units:.0%})')
    if is_main and bool(args.use_dsl) and int(args.dsl_lanes) > 0:
        per_exec = (
            int(args.dsl_max_sources_per_exec)
            if int(args.dsl_max_sources_per_exec) > 0 else int(args.dsl_max_sources)
        )
        print(f'[thought_dsl] enabled: lanes={args.dsl_lanes} rank={args.dsl_rank} '
              f'steps={args.dsl_steps} max_sources={args.dsl_max_sources} '
              f'per_exec={per_exec} lr_scale={args.dsl_lr_scale} '
              f'exec_lam={args.dsl_exec_lambda:g} carry_lam={args.dsl_carry_lambda:g} '
              f'src_lam={args.dsl_source_lambda:g} log={bool(args.dsl_log)}')
    # core itself stays on its per-slab devices; do NOT call model.to(device).
    if is_main:
        print(f'slab_devices={slab_devices}')
    dtype_audit = None
    dtype_audit_hooks = []
    if int(args.dtype_audit_steps) > 0 and is_main:
        dtype_audit = {}

        def _audit_walk_tensors(x, mod_key: str):
            if torch.is_tensor(x):
                dt = str(x.dtype).replace('torch.', '')
                ent = dtype_audit.setdefault(mod_key, {})
                ent[dt] = ent.get(dt, 0) + int(x.numel())
                return
            if isinstance(x, (list, tuple)):
                for y in x:
                    _audit_walk_tensors(y, mod_key)
                return
            if isinstance(x, dict):
                for y in x.values():
                    _audit_walk_tensors(y, mod_key)

        for n, m in model_base.named_modules():
            if n == '':
                continue
            mod_key = f'{m.__class__.__name__}:{n}'
            dtype_audit_hooks.append(m.register_forward_hook(
                lambda _m, _i, o, _k=mod_key: _audit_walk_tensors(o, _k)
            ))
    # SlabAdamW: build per-slab x per-subgroup param groups automatically
    # from model.named_parameters(). No hardcoded slab count; scales with
    # --num_slabs. gqk pair is routed into its own slow-warmup group.
    current_extra_lr_scales = {
        'decoder_emb_proj': float(args.emb_proj_lr_scale),
        'decoder_norm': float(args.decoder_norm_lr_scale),
        'container_alpha': float(args.container_lr_scale),
        'container_slab_gate': float(args.container_lr_scale),
        'container_neuron_gate': float(args.container_lr_scale),
        'container_tau_sa1': float(args.container_lr_scale),
        'container_tau_sa2': float(args.container_lr_scale),
        'container_tau_sa3': float(args.container_lr_scale),
        'container_other': float(args.container_lr_scale),
        'predictor': float(args.predictor_lr_scale),
        'dsl': float(args.dsl_lr_scale),
    }
    optim_groups = build_groups(
        model_base,
        base_lr=args.lr,
        weight_decay=args.weight_decay,
        gqk_lr_scale=args.gqk_lr_scale,
        extra_lr_scales=current_extra_lr_scales,
    )
    for g in optim_groups:
        n = str(g.get('name', ''))
        if n.startswith('shared/container_layer_') or n.startswith('shared/container_layer_skip_') or n.startswith('shared/container_layer_in'):
            scl = float(args.container_lr_scale)
            g['base_lr'] = args.lr * scl
            g['lr'] = g['base_lr']
            g['lr_scale'] = scl
    # Pre-warmup: gqk groups start at 0 until gqk_warmup_steps elapsed.
    if args.gqk_warmup_steps > 0:
        for g in optim_groups:
            if g.get('is_gqk'):
                g['lr'] = 0.0
    n_slab_groups = sum(1 for g in optim_groups if not g['name'].startswith('shared'))
    n_shared = sum(1 for g in optim_groups if g['name'].startswith('shared'))
    if is_main:
        print(f'[slab_adamw] groups: {len(optim_groups)} total '
              f'({n_slab_groups} slab subgroups + {n_shared} shared)')
        for g in optim_groups:
            n_p = sum(p.numel() for p in g['params'])
            print(f"  group={g['name']:<28s} params={len(g['params']):3d} numel={n_p:>10d} "
                  f"base_lr={g['base_lr']:.2e} wd={g['weight_decay']:.4g}")
            for pname in g.get('param_names', ()):
                print(f'    - {pname}')
    optim_master_dtype = {
        'fp32': torch.float32,
        'fp16': torch.float16,
        'bf16': torch.bfloat16,
    }[args.optim_dtype]
    if args.optim_device:
        if args.optim == 'atom':
            optimizer = OffloadedOptim(
                optim_groups,
                optim_device=args.optim_device,
                optim_class=ATOM,
                master_dtype=optim_master_dtype,
            )
            if is_main:
                print(f'[train] optimizer offloaded: master={args.optim_dtype} + Atom states on {args.optim_device}')
        else:
            optimizer = OffloadedOptim(
                optim_groups,
                optim_device=args.optim_device,
                optim_class=torch.optim.AdamW,
                master_dtype=optim_master_dtype,
            )
            if is_main:
                print(f'[train] optimizer offloaded: master={args.optim_dtype} + AdamW states on {args.optim_device}')
    elif args.optim == 'atom':
        optimizer = ATOM(optim_groups)
        if is_main:
            print(f'[train] optimizer: Atom (MUCH better AdamW drop-in replacement)')
    else:
        optimizer = torch.optim.AdamW(optim_groups)
        if is_main:
            print(f'[train] optimizer: AdamW (torch.optim.AdamW)')

    # Optional low-prec grad accumulator for the mixer's 3-backward window.
    grad_accum = None
    if args.grads_dtype in ('fp16', 'bf16'):
        grad_accum_dtype = torch.float16 if args.grads_dtype == 'fp16' else torch.bfloat16
        accum_device = args.optim_mem_device.strip() or None
        if accum_device and accum_device.startswith('cuda'):
            if not torch.cuda.is_available():
                raise RuntimeError(f'--optim_mem_device={accum_device} requested but CUDA unavailable')
            # Validate the device ordinal exists on this box. Otherwise the
            # offload hook fails deep inside autograd with the unhelpful
            # "CUDA error: invalid device ordinal" — fail early instead.
            try:
                dev_idx = torch.device(accum_device).index
            except (RuntimeError, ValueError) as e:
                raise RuntimeError(f'--optim_mem_device={accum_device} is not a valid device spec: {e}')
            n_dev = torch.cuda.device_count()
            if dev_idx is not None and dev_idx >= n_dev:
                raise RuntimeError(
                    f'--optim_mem_device={accum_device} requested ordinal {dev_idx} but only '
                    f'{n_dev} CUDA device(s) are visible. Set optim_mem_device: "" in your '
                    f'config to disable cross-device grad offload, or pick a valid ordinal.'
                )
        grad_accum = LowPrecGradAccumulator(
            model_base.parameters(), grad_accum_dtype, accum_device=accum_device,
        )
        if is_main:
            if args.optim == 'atom':
                print(f'[train] grads_dtype = {args.grads_dtype} '
                      f'device={accum_device or "(same as param)"} '
                      f'(Atom mixed-grad path: no fp32 param.grad materialization)')
            else:
                print(f'[train] grads_dtype = {args.grads_dtype} '
                      f'device={accum_device or "(same as param)"} '
                      f'(upcast to fp32 on param.device at optimizer.step)')

    _grad_bytes = sum(p.numel() * p.element_size()
                      for p in model_base.parameters() if p.requires_grad)
    _bucket_mb = max(25, int(_grad_bytes / 1024 / 1024) + 16)
    model = DDP(
        model_base,
        device_ids=[local_rank],
        output_device=local_rank,
        find_unused_parameters=True,
        bucket_cap_mb=_bucket_mb,
    ) if distributed else model_base

    total_params = count_params(model_base)
    core_params = count_params(model_base.core) + count_params(model_base.core_in_proj) + count_params(model_base.core_out_proj)
    dec_params = count_params(model_base.decoder)
    state_dict_params = sum(
        int(t.numel())
        for t in model_base.state_dict().values()
        if torch.is_tensor(t)
    )

    if is_main:
        gpu_name = torch.cuda.get_device_name(local_rank if distributed else 0) if args.device.startswith("cuda") else "cpu"
        print(f'GPU: {gpu_name}')
        print(f'params_total={total_params}')
        print(f'params_core={core_params}')
        print(f'params_decoder={dec_params}')
        if state_dict_params != total_params:
            print(
                f'params_state_dict={state_dict_params} '
                f'(includes aliased/tied entries)'
            )
    mem = mem_mb(args.batch_size, args.seq_len, args.decoder_dim, args.core_dim, total_params, train=True)
    if is_main:
        print(
            "mem_est_mb "
            f"params={mem['params_mb']:.1f} grads={mem['grads_mb']:.1f} "
            f"adam={mem['adam_mb']:.1f} acts~={mem['acts_mb_est']:.1f} total~={mem['total_mb_est']:.1f}"
        )

    ckpt_dir = Path(args.ckpt_dir)
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    if is_main:
        with open(ckpt_dir / 'run_config.json', 'w', encoding='utf-8') as f:
            json.dump(
                {
                    'rmu_cfg': {**asdict(rmu_cfg), 'dtype': str(rmu_cfg.dtype)},
                    'dec_cfg': {**asdict(dec_cfg), 'dtype': str(dec_cfg.dtype)},
                    'args': vars(args),
                },
                f,
                ensure_ascii=False,
                indent=2,
            )

    # DataLoader yields lists of raw `messages`; packing (tokenize +
    # boundary-aware mask construction) is done in-loop on CPU, right before
    # the GPU forward. We request a larger per-batch fetch than batch_size
    # because the packer may skip short samples and/or consume multiple
    # samples to fill a single row. The overflow is buffered across
    # iterations so we never stall the GPU.
    _PACK_FETCH_MULT = max(4, args.seq_len_mul * 4)

    # Pre-read resume step + data position (if any) so DataLoader/sampler and
    # the TrackedDataStream can be seeded from them before streaming starts.
    # Full resume (state dict load) happens further below.
    start_step = 0
    start_epoch = 0
    resume_dl_epoch = 0
    resume_skip_chars = 0
    if args.resume:
        _resume_path_preview = Path(args.resume)
        if _resume_path_preview.is_file():
            try:
                _peek = torch.load(_resume_path_preview, map_location='cpu', weights_only=False)
                start_step = int(_peek.get('step', 0))
                start_epoch = int(_peek.get('epoch', 0))
                _dp = _peek.get('data_pos')
                if _dp and int(_dp.get('world_size', 0)) == world_size:
                    resume_dl_epoch = int(_dp.get('dl_epoch', 0))
                    _prc = _dp.get('per_rank_chars') or [0] * world_size
                    if rank < len(_prc):
                        resume_skip_chars = int(_prc[rank])
                    if is_main:
                        print(f'resume: data_pos dl_epoch={resume_dl_epoch} '
                              f'per_rank_chars={_prc}')
                elif _dp and is_main:
                    print(f'resume: data_pos world_size mismatch '
                          f'(ckpt={_dp.get("world_size")} vs now={world_size}); '
                          f'falling back to fresh shuffle from start_step')
                elif is_main:
                    print(f'resume: no data_pos in ckpt; fresh shuffle seeded by start_step')
                del _peek
            except Exception as _e:
                if is_main:
                    print(f'resume: peek failed ({_e}); falling back')

    sampler = DistributedSampler(ds, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False) \
        if distributed else None
    if args.dataset_type == 'messages':
        collate_fn = (lambda batch: [x['messages'] for x in batch])
    elif 'input_ids' in ds.column_names:
        collate_fn = (lambda batch: [x['input_ids'] for x in batch])
    else:
        assert string_text_col is not None
        collate_fn = (lambda batch: [x[string_text_col] for x in batch])

    # Deterministic shuffle generator for non-DDP path, seeded from start_step
    # so resume produces a fresh data order (avoids revisiting batches from
    # the first `start_step` steps of the original run). DDP path uses the
    # same idea via `sampler.set_epoch(start_step)` below.
    _dl_gen = None
    if sampler is None:
        _dl_gen = torch.Generator()
        _dl_gen.manual_seed(0xC0FFEE ^ int(start_step))

    dl = DataLoader(
        ds,
        batch_size=args.batch_size * _PACK_FETCH_MULT,
        shuffle=(sampler is None),
        sampler=sampler,
        generator=_dl_gen,
        drop_last=False,
        num_workers=args.num_workers,
        pin_memory=False,  # packing happens on CPU -> keep non-pinned
        collate_fn=collate_fn,
    )

    import math
    lr_min = args.lr * 0.1 if args.lr_min < 0 else args.lr_min

    # --- Resume from checkpoint ---------------------------------------------
    # Loads weights and optimizer state if --resume is provided. The step
    # counter is restored so the LR schedule / NAT+FIM curriculum progress
    # pick up exactly where the previous run left off. If `--max_steps`
    # equals the saved step the training loop will exit immediately; raise
    # `--max_steps` beyond it to continue.
    # Note: start_step/start_epoch are pre-initialized above from the
    # checkpoint payload so DataLoader/sampler seeds match the resumed step.
    if args.resume:
        resume_path = Path(args.resume)
        if not resume_path.is_file():
            raise FileNotFoundError(f'--resume target not found: {resume_path}')
        if is_main:
            print(f'resume_from={resume_path}')
        # `weights_only=False` because we stored config dicts alongside the
        # state dict and optimizer state (trusted local file only).
        payload = torch.load(resume_path, map_location='cpu', weights_only=False)
        # Sanity check: key architecture fields must match, otherwise
        # weights could silently mis-shape.
        saved_rmu = payload.get('rmu_cfg', {})
        saved_dec = payload.get('dec_cfg', {})
        live_rmu = {**asdict(rmu_cfg), 'dtype': str(rmu_cfg.dtype)}
        live_dec = {**asdict(dec_cfg), 'dtype': str(dec_cfg.dtype)}
        for name, saved_cfg, live_cfg in (
            ('rmu_cfg', saved_rmu, live_rmu),
            ('dec_cfg', saved_dec, live_dec),
        ):
            for k in ('vocab_size', 'grid_g', 'grid_h', 'grid_w', 'channels',
                      'ray_dim', 'hidden_dim', 'factor_dim'):
                if k in saved_cfg and k in live_cfg and saved_cfg[k] != live_cfg[k]:
                    raise ValueError(
                        f'resume: {name}.{k} mismatch (ckpt={saved_cfg[k]} '
                        f'vs current={live_cfg[k]}); cannot load weights across '
                        f'architectural changes.'
                    )
        # Upgrade legacy fused readout (3D→D) to split readout_h + readout_c.
        import re as _re
        for k in list(payload['model_state'].keys()):
            _m = _re.match(r'^(core\.slabs\.slab\d+\.bank\.params\.slab\d+)_readout_([wb])$', k)
            if not _m:
                continue
            _pfx, _sfx = _m.group(1), _m.group(2)
            if f'{_pfx}_readout_h_{_sfx}' in payload['model_state']:
                continue
            _old = payload['model_state'].pop(k)
            if _sfx == 'w':
                _D = _old.shape[1]
                payload['model_state'][f'{_pfx}_readout_h_{_sfx}'] = _old[:2*_D, :]
                payload['model_state'][f'{_pfx}_readout_c_{_sfx}'] = _old[2*_D:, :]
            else:
                payload['model_state'][f'{_pfx}_readout_h_{_sfx}'] = _old.clone()
                payload['model_state'][f'{_pfx}_readout_c_{_sfx}'] = _old.clone()
        # Shape-adaptive resume for DSL: when DSL dimensions changed
        # (steps/lanes/rank/max_sources), copy the overlapping slice from
        # checkpoint tensor into current tensor and keep the rest as freshly
        # initialized params. This allows "expanding" DSL capacity without
        # throwing away all learned DSL weights.
        _ckpt_sd = payload['model_state']
        _live_sd = model_base.state_dict()
        _patched_sd = dict(_ckpt_sd)
        _dsl_adapted = []
        for _k, _v_live in _live_sd.items():
            if '.dsl.' not in _k:
                continue
            _v_ckpt = _patched_sd.get(_k, None)
            if _v_ckpt is None or not torch.is_tensor(_v_ckpt) or not torch.is_tensor(_v_live):
                continue
            if tuple(_v_ckpt.shape) == tuple(_v_live.shape):
                continue
            _new = _v_live.detach().clone()
            _slices = tuple(slice(0, min(int(a), int(b))) for a, b in zip(_v_ckpt.shape, _v_live.shape))
            if len(_slices) == _new.dim():
                _new[_slices] = _v_ckpt[_slices].to(dtype=_new.dtype)
                _patched_sd[_k] = _new
                _dsl_adapted.append((_k, tuple(_v_ckpt.shape), tuple(_v_live.shape)))
        if is_main and _dsl_adapted:
            print(f'resume: DSL shape-adapted tensors={len(_dsl_adapted)}')
            for _k, _old, _new in _dsl_adapted[:10]:
                print(f'  - {_k}: ckpt{_old} -> live{_new} (overlap copied)')
            if len(_dsl_adapted) > 10:
                print(f'  ... and {len(_dsl_adapted)-10} more')

        # Load weights onto the already-placed submodules. Using a CPU copy
        # then letting module.load_state_dict handle placement per tensor via
        # the already-initialized module devices.
        missing, unexpected = model_base.load_state_dict(_patched_sd, strict=False)
        if is_main:
            if missing:
                print(f'resume: missing_keys={len(missing)} (first 5): {missing[:5]}')
            if unexpected:
                print(f'resume: unexpected_keys={len(unexpected)} (first 5): {unexpected[:5]}')
        try:
            optimizer.load_state_dict(payload['optimizer_state'])
        except ValueError as e:
            if is_main:
                print(f"resume: Could not load optimizer state (likely due to added layers). Starting with fresh optimizer. Error: {e}")
            if hasattr(optimizer, '_step'):
                optimizer._step = int(payload.get('step', 0))
                
        start_step = int(payload.get('step', 0))
        start_epoch = int(payload.get('epoch', 0))
        # On resume, --max_steps means ADDITIONAL steps to run beyond the
        # checkpoint, so total_steps = start_step + args.max_steps. The LR
        # schedule, FIM warmup window, and loop termination all read
        # args.max_steps, so rebasing it here lets every downstream consumer
        # treat the resumed run as a clean continuation that ends at the
        # new total. Without this rebase the user would need to pass
        # --max_steps = (already_done + extra), which is error-prone.
        _extra_steps = int(args.max_steps)
        args.max_steps = start_step + _extra_steps
        # Re-base group LRs to the CURRENT CLI/YAML scales, not the stale
        # per-group lr_scale values serialized inside the checkpoint.
        # Otherwise a resume silently keeps old dampening for container /
        # decoder groups even when the new run_config changed them.
        for g in optimizer.param_groups:
            name = str(g.get('name', ''))
            sub = name.split('/', 1)[1] if '/' in name else name
            if g.get('is_gqk'):
                scale = float(args.gqk_lr_scale)
            elif name.startswith('shared/container_layer_') or name.startswith('shared/container_layer_skip_') or name.startswith('shared/container_layer_in'):
                scale = float(args.container_lr_scale)
            else:
                scale = float(current_extra_lr_scales.get(sub, 1.0))
            g['lr_scale'] = scale
            g['base_lr'] = args.lr * scale
            g['lr'] = g['base_lr']
        if is_main:
            print(f'resume: start_step={start_step} start_epoch={start_epoch} '
                  f'last_loss={payload.get("loss", float("nan")):.4f}')
            print(f'resume: --max_steps={_extra_steps} treated as ADDITIONAL '
                  f'-> total_steps={args.max_steps} (will run '
                  f'{args.max_steps - start_step} more steps)')

    model.train()

    # ---------------- Packing + buffered message stream ---------------------
    # `TrackedDataStream` wraps the DataLoader to track consumed chars/tokens
    # so resume can skip to the exact position without re-tokenizing. Also
    # manages `sampler.set_epoch` per DataLoader-restart: DO NOT call
    # sampler.set_epoch elsewhere.
    msg_iter = TrackedDataStream(
        dl, sampler=sampler,
        start_epoch=resume_dl_epoch,
        skip_chars=resume_skip_chars,
    )
    if resume_skip_chars > 0 and is_main:
        print(f'resume: stream will skip {resume_skip_chars} chars/tokens '
              f'on rank={rank} (no tokenization during skip)')
    token_tail: List[int] = []

    def _fetch_packed(row_len: int) -> PackedBatch:
        """Pull messages and pack into one batch at the requested row_len."""
        nonlocal token_tail
        if args.dataset_type == 'messages':
            pb = pack_messages_to_batch(
                msg_iter, tokenizer, row_len=row_len,
                batch_size=args.batch_size, pad_id=rmu_cfg.pad_token_id,
                eos_id=rmu_cfg.eos_token_id,
            )
            return pb.to(device, non_blocking=True)

        # PT mode: pure token stream, no chat/message boundaries.
        total = args.batch_size * row_len
        stream: List[int] = token_tail
        token_tail = []
        while len(stream) < total:
            ids = next(msg_iter)
            if isinstance(ids, torch.Tensor):
                ids = ids.tolist()
            elif isinstance(ids, str):
                ids = tokenizer(ids, add_special_tokens=False)["input_ids"]
            if not ids:
                continue
            stream.extend(int(t) for t in ids)
        take = stream[:total]
        token_tail = stream[total:]
        x = torch.tensor(take, dtype=torch.long).view(args.batch_size, row_len)
        valid = torch.ones_like(x, dtype=torch.bool)
        reset = torch.zeros_like(x, dtype=torch.bool)
        reset[:, 0] = True
        write = torch.ones_like(x, dtype=torch.bool)
        # BOS write included: scan isolation is provided by reset_mask, not
        # by excluding position 0 from write (consistent with packing.py).
        pb = PackedBatch(
            input_ids=x,
            target_ids=x.clone(),
            reset_mask=reset,
            valid_mask=valid,
            write_mask=write,
            user_mask=torch.zeros_like(x, dtype=torch.bool),
            sample_offsets=[[0, row_len] for _ in range(args.batch_size)],
        )
        return pb.to(device, non_blocking=True)

    # ---------------- Validation function -------------------------
    # Per-mode task weights for val mixing — kept identical to the
    # training task_inputs weights so the val number tracks exactly the
    # objective the optimizer minimizes (no metric drift between train
    # and val). When a mode is silently a no-op on the current dataset
    # (blind/rev on string corpora), its loss is excluded from the mix.
    _VAL_TASK_W: Dict[str, float] = {
        'pt': 1.0, 'nat': float(args.nat_weight), 'blind': 1.0, 'rev': 0.2, 'fim': 1.0,
    }

    @torch.no_grad()
    def _xmem_retention(model) -> Optional[float]:
        """Mean per-position retention = mean(∏_a σ(decay_logits[a]))."""
        dl = getattr(model.core, 'xmem_decay_logits', None)
        if dl is None:
            return None
        per_axis = torch.sigmoid(dl.float())           # [N, A]
        per_pos = per_axis.prod(dim=0).mean()           # scalar
        return float(per_pos.item())

    def _run_validation(val_dataset) -> Optional[float]:
        """Run validation across ALL active training modes and return a
        single task-weighted mean loss. Mirrors the training task mixer:
        for each batch, every active mode produces its (input, target) view
        of the same packed batch (no batch drift), runs one forward, and
        contributes `task_w` to the weighted sum. The final scalar is the
        same objective the optimizer is minimizing, evaluated on held-out
        data."""
        if val_dataset is None:
            return None
        model.eval()
        val_stream: List[int] = []
        val_iter = msg_stream(DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0))

        # Per-mode accumulator: list of per-batch losses; averaged below.
        per_mode_losses: Dict[str, List[float]] = {m: [] for m in modes}

        # Process up to 10 batches for validation (to avoid long validation time)
        max_val_batches = 10
        with torch.no_grad():
            for _ in range(max_val_batches):
                try:
                    # Get tokens from validation dataset
                    total = args.batch_size * row_len_max
                    while len(val_stream) < total:
                        ids = next(val_iter)
                        if isinstance(ids, torch.Tensor):
                            ids = ids.tolist()
                        elif isinstance(ids, str):
                            ids = tokenizer(ids, add_special_tokens=False)["input_ids"]
                        if not ids:
                            continue
                        val_stream.extend(int(t) for t in ids)
                    if len(val_stream) < total:
                        break
                    take = val_stream[:total]
                    val_stream = val_stream[total:]
                    x = torch.tensor(take, dtype=torch.long).view(args.batch_size, row_len_max)
                    valid = torch.ones_like(x, dtype=torch.bool)
                    reset = torch.zeros_like(x, dtype=torch.bool)
                    reset[:, 0] = True
                    write = torch.ones_like(x, dtype=torch.bool)
                    val_pb = PackedBatch(
                        input_ids=x,
                        target_ids=x.clone(),
                        reset_mask=reset,
                        valid_mask=valid,
                        write_mask=write,
                        user_mask=torch.zeros_like(x, dtype=torch.bool),
                        sample_offsets=[[0, row_len_max] for _ in range(args.batch_size)],
                    ).to(device, non_blocking=True)

                    # Build per-mode (input, target) pairs that mirror the
                    # training task_inputs construction exactly.
                    val_tasks: List[Tuple[str, torch.Tensor, torch.Tensor]] = []
                    if 'pt' in modes:
                        val_tasks.append(('pt', val_pb.input_ids, val_pb.target_ids))
                    if 'nat' in modes:
                        n_in, n_tgt = apply_nat_to_packed(
                            val_pb, pad_id=rmu_cfg.pad_token_id,
                            step=args.max_steps, max_steps=args.max_steps,
                            warmup_frac=0.0, peak_prob=1.0,
                            min_blind_frac=0.15, max_blind_frac=0.9,
                        )
                        val_tasks.append(('nat', n_in, n_tgt))
                    if 'fim' in modes:
                        f_in, f_tgt = apply_fim_to_packed(
                            val_pb, pad_id=rmu_cfg.pad_token_id, vocab_size=vocab_size,
                            step=args.max_steps, max_steps=args.max_steps,
                            warmup_frac=0.0,
                            min_span_frac=FIM_MIN_SPAN_FRAC, max_span_frac=FIM_MAX_SPAN_FRAC,
                        )
                        val_tasks.append(('fim', f_in, f_tgt))
                    # blind/rev require user_mask; on string corpora the
                    # supervised set is empty -> skip cleanly to avoid a
                    # zero-loss entry that would bias the weighted mean.
                    if ('blind' in modes or 'rev' in modes) and bool(val_pb.user_mask.any().item()):
                        b_allowed = _segments_with_user_mask(val_pb.user_mask, val_pb.reset_mask)
                        if 'blind' in modes:
                            sup = val_pb.target_ids.ne(rmu_cfg.pad_token_id) & b_allowed
                            b_in = val_pb.input_ids.clone()
                            b_in.masked_fill_(sup, rmu_cfg.pad_token_id)
                            b_tgt = val_pb.target_ids.clone()
                            b_tgt.masked_fill_(~sup, rmu_cfg.pad_token_id)
                            val_tasks.append(('blind', b_in, b_tgt))
                        if 'rev' in modes:
                            sup = val_pb.target_ids.ne(rmu_cfg.pad_token_id) & b_allowed
                            r_in = val_pb.target_ids.clone()
                            r_in.masked_fill_(sup, rmu_cfg.pad_token_id)
                            r_tgt = val_pb.input_ids.clone()
                            r_tgt.masked_fill_(~sup, rmu_cfg.pad_token_id)
                            val_tasks.append(('rev', r_in, r_tgt))

                    for tname, x_in, x_tgt in val_tasks:
                        with autocast_ctx(args.device, compute_dtype):
                            out_val = model(
                                x_in, target_ids=x_tgt,
                                valid_mask=val_pb.valid_mask,
                                write_mask=val_pb.write_mask,
                                reset_mask=val_pb.reset_mask,
                                learnorm_lambda=args.learnorm_lambda,
                                predictor_lambda=args.predictor_lambda,
                                vector_error_lambda=args.vector_error_lambda,
                                dsl_exec_lambda=args.dsl_exec_lambda,
                                dsl_carry_lambda=args.dsl_carry_lambda,
                                dsl_source_lambda=args.dsl_source_lambda,
                                return_loss_parts=True,
                            )
                        per_mode_losses[tname].append(float(out_val['loss'].item()))
                except StopIteration:
                    break
        model.train()

        # Weighted aggregation: identical formula to training's
        # `total_task = loss * task_w / total_task_weight` accumulated over
        # tasks, but applied to per-mode mean-of-batches.
        active = [(m, sum(ls) / len(ls)) for m, ls in per_mode_losses.items() if ls]
        if not active:
            return None
        total_w = sum(_VAL_TASK_W.get(m, 1.0) for m, _ in active)
        if total_w <= 0.0:
            return None
        weighted = sum(_VAL_TASK_W.get(m, 1.0) * v for m, v in active) / total_w
        return float(weighted)

    def _build_val_cache() -> Optional[np.ndarray]:
        """Pre-tokenize val_ds once into a flat int32 mmap file.

        Only rank 0 writes; all ranks load after barrier. Cache key is the
        total token count so it auto-invalidates on config changes.
        Returns read-only np.memmap [n_tokens] or None."""
        if val_ds is None:
            return None
        n_tokens = 10 * args.batch_size * row_len_max
        cache_path = ckpt_dir / f'val_cache_{n_tokens}.bin'
        if is_main:
            need_build = True
            if cache_path.exists():
                try:
                    mm = np.memmap(cache_path, dtype=np.int32, mode='r')
                    need_build = mm.shape[0] != n_tokens
                    del mm
                except Exception:
                    pass
            if need_build:
                print('[val_cache] pre-tokenizing val split...', flush=True)
                stream: List[int] = []
                _vit = msg_stream(DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=0))
                try:
                    while len(stream) < n_tokens:
                        ids = next(_vit)
                        if isinstance(ids, torch.Tensor):
                            ids = ids.tolist()
                        elif isinstance(ids, str):
                            ids = tokenizer(ids, add_special_tokens=False)['input_ids']
                        if ids:
                            stream.extend(int(t) for t in ids)
                except StopIteration:
                    pass
                if len(stream) < n_tokens:
                    print('[val_cache] val_ds too small, validation disabled', flush=True)
                    if distributed:
                        _dist_barrier(local_rank)
                    return None
                mm = np.memmap(cache_path, dtype=np.int32, mode='w+', shape=(n_tokens,))
                mm[:] = stream[:n_tokens]
                mm.flush()
                del mm
                print(f'[val_cache] {n_tokens:,} tokens -> {cache_path.name}', flush=True)
        if distributed:
            _dist_barrier(local_rank)
        if not cache_path.exists():
            return None
        return np.memmap(cache_path, dtype=np.int32, mode='r')

    def _run_validation_mmap(val_cache: np.ndarray) -> Optional[float]:
        """GPU-bound validation from pre-tokenized mmap: no tokenizer, no DataLoader."""
        model.eval()
        step_toks = args.batch_size * row_len_max
        n_batches = len(val_cache) // step_toks
        per_mode_losses: Dict[str, List[float]] = {m: [] for m in modes}
        with torch.no_grad():
            for i in range(n_batches):
                raw = np.array(val_cache[i * step_toks:(i + 1) * step_toks], dtype=np.int64)
                x = torch.from_numpy(raw).view(args.batch_size, row_len_max).to(device, non_blocking=True)
                valid = torch.ones(args.batch_size, row_len_max, dtype=torch.bool, device=device)
                reset = torch.zeros_like(valid)
                reset[:, 0] = True
                write = torch.ones_like(valid)
                val_pb = PackedBatch(
                    input_ids=x,
                    target_ids=x.clone(),
                    reset_mask=reset,
                    valid_mask=valid,
                    write_mask=write,
                    user_mask=torch.zeros_like(valid),
                    sample_offsets=[[0, row_len_max] for _ in range(args.batch_size)],
                ).to(device, non_blocking=True)
                val_tasks: List[Tuple[str, torch.Tensor, torch.Tensor]] = []
                if 'pt' in modes:
                    val_tasks.append(('pt', val_pb.input_ids, val_pb.target_ids))
                if 'nat' in modes:
                    n_in, n_tgt = apply_nat_to_packed(
                        val_pb, pad_id=rmu_cfg.pad_token_id,
                        step=args.max_steps, max_steps=args.max_steps,
                        warmup_frac=0.0, peak_prob=1.0,
                        min_blind_frac=0.15, max_blind_frac=0.9,
                    )
                    val_tasks.append(('nat', n_in, n_tgt))
                for tname, x_in, x_tgt in val_tasks:
                    with autocast_ctx(args.device, compute_dtype):
                        out_val = model(
                            x_in, target_ids=x_tgt,
                            valid_mask=val_pb.valid_mask,
                            write_mask=val_pb.write_mask,
                            reset_mask=val_pb.reset_mask,
                            learnorm_lambda=args.learnorm_lambda,
                            predictor_lambda=args.predictor_lambda,
                            vector_error_lambda=args.vector_error_lambda,
                            dsl_exec_lambda=args.dsl_exec_lambda,
                            dsl_carry_lambda=args.dsl_carry_lambda,
                            dsl_source_lambda=args.dsl_source_lambda,
                            return_loss_parts=True,
                        )
                    per_mode_losses[tname].append(float(out_val['loss'].item()))
        model.train()
        active = [(m, sum(ls) / len(ls)) for m, ls in per_mode_losses.items() if ls]
        if not active:
            return None
        total_w = sum(_VAL_TASK_W.get(m, 1.0) for m, _ in active)
        if total_w <= 0.0:
            return None
        return float(sum(_VAL_TASK_W.get(m, 1.0) * v for m, v in active) / total_w)

    val_cache: Optional[np.ndarray] = _build_val_cache()

    # ---------------- Dry-run at worst-case row_len -------------------------
    # Allocates peak activation footprint on step 0. If OOM, fails now with a
    # clear message instead of ~45 minutes into training. Peak-reservation
    # stays in the allocator so subsequent smaller-N batches reuse blocks
    # without cudaMalloc churn.
    # Note: sampler.set_epoch is managed by TrackedDataStream; do not call here.
    if is_main:
        print(f'[dry_run] forward+backward at row_len={row_len_max} '
              f'B={args.batch_size} (worst-case VRAM anchor)...')
    torch.cuda.reset_peak_memory_stats(device) if args.device.startswith('cuda') else None
    dry_pb = _fetch_packed(row_len_max)
    with (_activation_offload_ctx(args.optim_mem_device) if _activation_offload_enabled(args) else contextlib.nullcontext()):
        with autocast_ctx(args.device, compute_dtype):
            dry_loss = model(
                dry_pb.input_ids, target_ids=dry_pb.target_ids,
                valid_mask=dry_pb.valid_mask, write_mask=dry_pb.write_mask,
                reset_mask=dry_pb.reset_mask,
                learnorm_lambda=args.learnorm_lambda,
                predictor_lambda=args.predictor_lambda,
                vector_error_lambda=args.vector_error_lambda,
                dsl_exec_lambda=args.dsl_exec_lambda,
                dsl_carry_lambda=args.dsl_carry_lambda,
                dsl_source_lambda=args.dsl_source_lambda,
            )
    if not torch.isfinite(dry_loss):
        raise RuntimeError(f'[dry_run] non-finite loss on init: {dry_loss.item()}')
    dry_loss.backward()
    if args.device.startswith('cuda'):
        peak_alloc = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
        peak_resv = torch.cuda.max_memory_reserved(device) / (1024 ** 3)
        total_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
        headroom_pct = 100.0 * (1.0 - peak_resv / total_gb)
        if is_main:
            print(f'[dry_run] ok loss={dry_loss.item():.4f}  peak_alloc={peak_alloc:.2f}GB '
                  f'peak_reserved={peak_resv:.2f}GB / {total_gb:.2f}GB  '
                  f'headroom={headroom_pct:.1f}%')
            if headroom_pct < 10.0:
                print(f'[dry_run] WARNING: headroom <10% -> fragmentation + growing Adam '
                      f'state may OOM late. Reduce batch_size/seq_len_mul/core_dim.')
    optimizer.zero_grad(set_to_none=True)
    if grad_accum is not None:
        grad_accum.reset()
    del dry_pb, dry_loss

    t0 = time.time()
    t_log = t0
    step = start_step
    epoch = start_epoch
    grad_tokens_total = 0
    proc_tokens_total = 0
    proc_tokens_log = 0

    import random as _py_random
    _rng = _py_random.Random(0xC0FFEE ^ start_step)

    # ---- per-batch metric EMAs --------------------------------------------
    # Per-batch losses (pt/ce/lnorm and per-task) are extremely noisy on
    # heterogeneous corpora (SlimPajama: code/wiki/cc/arxiv mixed). The
    # instantaneous log line can swing ±1 nat between consecutive steps
    # while the underlying trend is monotonic. EMA with alpha=0.05 gives
    # an effective window of ~40 steps (half-life ≈ 14) so the smoothed
    # value tracks real progress without lagging behind LR transitions.
    # Stored separately per metric key; new keys are auto-initialized to
    # the first observed value (no zero-bias warmup transient).
    _LOSS_EMA_ALPHA = 0.05
    _loss_ema: Dict[str, float] = {}

    def _ema_update(key: str, value: float) -> float:
        prev = _loss_ema.get(key)
        new = float(value) if prev is None else (1.0 - _LOSS_EMA_ALPHA) * prev + _LOSS_EMA_ALPHA * float(value)
        _loss_ema[key] = new
        return new

    while step < args.max_steps:
        epoch += 1
        step += 1
        if step > args.max_steps:
            break
        # Note: sampler.set_epoch is managed by TrackedDataStream per
        # DataLoader-restart; do not call it here (was buggy per-step call).

        # ---------- Variable-N chunk schedule (triangle-biased) ------------
        N = triangle_sample_n(args.seq_len_mul, rng=_rng)
        row_len = N * args.seq_len

        # ---------- LR schedule -------------------------------------------
        base_lr_now = lr_at(step, args.lr, lr_min, args.warmup_steps, args.max_steps)
        schedule_scale = base_lr_now / max(args.lr, 1e-12)
        for g in optimizer.param_groups:
            if g.get('is_gqk') and step <= args.gqk_warmup_steps:
                g['lr'] = 0.0
            else:
                g['lr'] = g['base_lr'] * schedule_scale
        optimizer.zero_grad(set_to_none=True)

        # ---------- Per-sub-sample NAT / FIM curricula --------------------
        # During warmup each curriculum is a no-op (returns input_ids as-is),
        # so running them would burn 3x compute for identical loss. Skip them
        # cleanly until their respective warmup_frac window opens; the task
        # count drops to 1 (std only) and loss/n_tasks adapts automatically.
        _fim_warm = int(args.max_steps * FIM_WARMUP_FRAC)
        # Task mixer weights:
        #   pt    : main supervised objective (former "std")
        #   blind : full-weight reconstruction objective
        #   rev   : weaker intent-reversal objective
        #   fim   : full-weight corruption-recovery objective
        per_task_losses: Dict[str, float] = {}
        per_task_ce_losses: Dict[str, float] = {}
        per_task_learnorm_losses: Dict[str, float] = {}
        per_task_predictor_losses: Dict[str, float] = {}
        per_task_vector_error_losses: Dict[str, float] = {}
        per_task_dsl_src_losses: Dict[str, float] = {}
        grad_tokens_step = 0
        proc_tokens_step = 0
        skip_step = False
        accum_steps = max(1, int(args.grad_accum_steps))
        task_count_last = 0
        for micro_idx in range(accum_steps):
            # ---------- Fetch packed batch at this row_len -----------------
            packed = _fetch_packed(row_len)
            task_inputs = []
            if 'pt' in modes:
                task_inputs.append(('pt', packed.input_ids, packed.target_ids, 1.0))

            if 'nat' in modes:
                nat_in, nat_tgt = apply_nat_to_packed(
                    packed, pad_id=rmu_cfg.pad_token_id,
                    step=args.max_steps, max_steps=args.max_steps,
                    warmup_frac=0.0, peak_prob=1.0,
                    min_blind_frac=0.15, max_blind_frac=0.9,
                )
                task_inputs.append(('nat', nat_in, nat_tgt, args.nat_weight))

            blind_allowed = None
            if 'blind' in modes or 'rev' in modes:
                blind_allowed = _segments_with_user_mask(packed.user_mask, packed.reset_mask)

            if 'blind' in modes:
                blind_supervised = packed.target_ids.ne(rmu_cfg.pad_token_id) & blind_allowed
                blind_in = packed.input_ids.clone()
                blind_in.masked_fill_(blind_supervised, rmu_cfg.pad_token_id)
                blind_tgt = packed.target_ids.clone()
                blind_tgt.masked_fill_(~blind_supervised, rmu_cfg.pad_token_id)
                task_inputs.append(('blind', blind_in, blind_tgt, 1.0))

            if 'rev' in modes:
                rev_supervised = packed.user_mask & blind_allowed
                rev_in = packed.input_ids.clone()
                rev_in.masked_fill_(rev_supervised, rmu_cfg.pad_token_id)
                rev_tgt = packed.input_ids.clone()
                rev_tgt.masked_fill_(~rev_supervised, rmu_cfg.pad_token_id)
                task_inputs.append(('rev', rev_in, rev_tgt, 0.2))

            if 'fim' in modes and step >= _fim_warm:
                fim_in, fim_tgt = apply_fim_to_packed(
                    packed, pad_id=rmu_cfg.pad_token_id, vocab_size=vocab_size,
                    step=step, max_steps=args.max_steps,
                    warmup_frac=FIM_WARMUP_FRAC, min_span_frac=FIM_MIN_SPAN_FRAC, max_span_frac=FIM_MAX_SPAN_FRAC,
                )
                task_inputs.append(('fim', fim_in, fim_tgt, 1.0))
            if not task_inputs:
                raise RuntimeError('no active task inputs for this step; check --modes')
            task_count_last = len(task_inputs)
            total_task_weight = float(sum(w for _, _, _, w in task_inputs))
            sync_ctx = (
                model.no_sync()
                if distributed and micro_idx + 1 < accum_steps
                else contextlib.nullcontext()
            )
            with sync_ctx:
                for task_name, x_in, x_tgt, task_w in task_inputs:
                    grad_tokens_step += int(x_tgt[:, 1:].ne(rmu_cfg.pad_token_id).sum().item())
                    proc_tokens_step += int(x_in.numel())
                    with (_activation_offload_ctx(args.optim_mem_device) if _activation_offload_enabled(args) else contextlib.nullcontext()):
                        with autocast_ctx(args.device, compute_dtype):
                            out_task = model(
                                x_in, target_ids=x_tgt,
                                valid_mask=packed.valid_mask,
                                write_mask=packed.write_mask,
                                reset_mask=packed.reset_mask,
                                learnorm_lambda=args.learnorm_lambda,
                                predictor_lambda=args.predictor_lambda,
                                vector_error_lambda=args.vector_error_lambda,
                                dsl_exec_lambda=args.dsl_exec_lambda,
                                dsl_carry_lambda=args.dsl_carry_lambda,
                                dsl_source_lambda=args.dsl_source_lambda,
                                return_loss_parts=True,
                            )
                            total_task = out_task['loss'] * (
                                task_w / max(total_task_weight, 1e-12) / float(accum_steps)
                            )
                    if not torch.isfinite(total_task):
                        if is_main:
                            print(
                                f'step={step} micro={micro_idx + 1}/{accum_steps} '
                                f'task={task_name} loss_non_finite={total_task.item()} (skip step)'
                            )
                        skip_step = True
                        break
                    total_task.backward()
                    per_task_losses[task_name] = per_task_losses.get(task_name, 0.0) + (
                        out_task['loss'].item() / float(accum_steps)
                    )
                    per_task_ce_losses[task_name] = per_task_ce_losses.get(task_name, 0.0) + (
                        out_task['ce_loss'].item() / float(accum_steps)
                    )
                    per_task_learnorm_losses[task_name] = per_task_learnorm_losses.get(task_name, 0.0) + (
                        out_task['learnorm_loss'].item() / float(accum_steps)
                    )
                    per_task_predictor_losses[task_name] = per_task_predictor_losses.get(task_name, 0.0) + (
                        out_task['predictor_loss'].item() / float(accum_steps)
                    )
                    per_task_vector_error_losses[task_name] = per_task_vector_error_losses.get(task_name, 0.0) + (
                        out_task['vector_error_loss'].item() / float(accum_steps)
                    )
                    per_task_dsl_src_losses[task_name] = per_task_dsl_src_losses.get(task_name, 0.0) + (
                        out_task['dsl_src_loss'].item() / float(accum_steps)
                    )
                    # Drop transient graph refs held on the module and in the
                    # local loss dict as soon as this task backward finishes.
                    for _attr in (
                        '_predictor_step_loss',
                        '_dsl_exec_step_loss',
                        '_dsl_carry_step_loss',
                        '_dsl_src_step_loss',
                        '_predictor_row_loss',
                        '_dsl_exec_row_loss',
                        '_dsl_carry_row_loss',
                        '_dsl_src_row_loss',
                    ):
                        if hasattr(model_base, _attr):
                            setattr(model_base, _attr, None)
                    del out_task, total_task
            if skip_step:
                break
        # -- end per-task loop --
        if skip_step:
            optimizer.zero_grad(set_to_none=True)
            if grad_accum is not None:
                grad_accum.reset()
            continue

        if isinstance(optimizer, (OffloadedOptim, OffloadedAdamW)):
            # Pull grads onto master params on optim_device (drains the
            # grad_accum buffer if present, else reads compute .grad).
            optimizer.ingest_grads(grad_accum=grad_accum)
            grad_norm = optimizer.clip_grad_norm(args.grad_clip) \
                if args.grad_clip > 0 else None
        else:
            if grad_accum is not None:
                if args.optim == 'atom':
                    optimizer.ingest_grads(grad_accum)
                else:
                    grad_accum.flush_to_param_grads(out_dtype=torch.float32)
            if args.grad_clip > 0:
                if args.optim == 'atom':
                    grad_norm = optimizer.clip_grad_norm(args.grad_clip)
                else:
                    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            else:
                grad_norm = None

        if grad_norm is not None and not torch.isfinite(grad_norm):
            if is_main:
                print(f'step={step} grad_norm_non_finite={grad_norm.item()} (skip step)')
            optimizer.zero_grad(set_to_none=True)
            if grad_accum is not None:
                grad_accum.reset()
            continue
        if distributed:
            _tok_t = torch.tensor([grad_tokens_step], device=device, dtype=torch.long)
            dist.all_reduce(_tok_t, op=dist.ReduceOp.SUM)
            grad_tokens_step = int(_tok_t.item())
            _proc_t = torch.tensor([proc_tokens_step], device=device, dtype=torch.long)
            dist.all_reduce(_proc_t, op=dist.ReduceOp.SUM)
            proc_tokens_step = int(_proc_t.item())
        grad_tokens_total += grad_tokens_step
        proc_tokens_total += proc_tokens_step
        proc_tokens_log += proc_tokens_step
        will_log_now = bool(is_main and (step % args.log_every == 0 or step == 1))
        # Snapshot is only needed when we're going to render `upd=…ppm`.
        # When `update_stats_every > 0` it gates the snapshot independently
        # of `log_every`; otherwise legacy behavior (snapshot on every log).
        if args.update_stats_every > 0:
            will_snapshot = bool(
                is_main and (step % args.update_stats_every == 0 or step == 1)
            )
        else:
            will_snapshot = will_log_now
        pre_step_snapshot = _snapshot(optimizer) if will_snapshot else None
        optimizer.step()
        actual_step_stats = slab_updates(pre_step_snapshot) if pre_step_snapshot is not None else None

        mixed_loss = (
            sum(per_task_losses[name] * task_w for name, _, _, task_w in task_inputs)
            / max(total_task_weight, 1e-12)
        )
        mixed_ce_loss = (
            sum(per_task_ce_losses[name] * task_w for name, _, _, task_w in task_inputs)
            / max(total_task_weight, 1e-12)
        )
        mixed_learnorm_loss = (
            sum(per_task_learnorm_losses[name] * task_w for name, _, _, task_w in task_inputs)
            / max(total_task_weight, 1e-12)
        )
        mixed_predictor_loss = (
            sum(per_task_predictor_losses[name] * task_w for name, _, _, task_w in task_inputs)
            / max(total_task_weight, 1e-12)
        )
        mixed_vector_error_loss = (
            sum(per_task_vector_error_losses[name] * task_w for name, _, _, task_w in task_inputs)
            / max(total_task_weight, 1e-12)
        )
        mixed_dsl_src_loss = (
            sum(per_task_dsl_src_losses[name] * task_w for name, _, _, task_w in task_inputs)
            / max(total_task_weight, 1e-12)
        )
        # Update per-batch EMAs every step (cheap; only float arithmetic).
        # Logged values below pull the smoothed view, instantaneous values
        # remain available via the unsuffixed keys.
        _ema_update('mixed', mixed_loss)
        _ema_update('ce', mixed_ce_loss)
        _ema_update('lnorm', mixed_learnorm_loss)
        _ema_update('pred', mixed_predictor_loss)
        _ema_update('vec', mixed_vector_error_loss)
        _ema_update('dsl_src', mixed_dsl_src_loss)
        for _tname, _tval in per_task_losses.items():
            _ema_update(_tname, _tval)
        ckpt_saved_now = (step % args.save_every == 0 or step == args.max_steps)
        if ckpt_saved_now:
            # Build data_pos. all_gather must run on ALL ranks (if DDP).
            if distributed:
                _local = torch.tensor(
                    [int(msg_iter.chars_in_epoch)], dtype=torch.long, device=device,
                )
                _gathered = [torch.zeros_like(_local) for _ in range(world_size)]
                dist.all_gather(_gathered, _local)
                per_rank_chars = [int(g.item()) for g in _gathered]
            else:
                per_rank_chars = [int(msg_iter.chars_in_epoch)]
            data_pos = {
                'dl_epoch': int(msg_iter.dl_epoch),
                'per_rank_chars': per_rank_chars,
                'world_size': int(world_size),
            }
            if is_main:
                ckpt_path = ckpt_dir / f'step_{step:07d}.pt'
                _ckpt_loss = _loss_ema.get('mixed', mixed_loss)
                save_ckpt(ckpt_path, model_base, optimizer, step, epoch,
                          float(_ckpt_loss), data_pos=data_pos)
                save_ckpt(ckpt_dir / 'latest.pt', model_base, optimizer, step, epoch,
                          float(_ckpt_loss), data_pos=data_pos)

        # Run validation on ALL ranks simultaneously so both GPUs stay busy
        # (avoids rank-1 spin-waiting in DDP allreduce while rank-0 runs val).
        # Results are averaged across ranks via all_reduce before logging.
        _will_val = (val_cache is not None) and (step % args.log_every == 0 or step == 1)
        val_loss_raw: Optional[float] = None
        if _will_val:
            _val_local = _run_validation_mmap(val_cache)
            if _val_local is not None and distributed:
                _vt = torch.tensor([_val_local], dtype=torch.float32, device=device)
                dist.all_reduce(_vt, op=dist.ReduceOp.AVG)
                val_loss_raw = float(_vt.item())
            else:
                val_loss_raw = _val_local

        if is_main and (step % args.log_every == 0 or step == 1):
            now = time.time()
            elapsed_log = now - t_log
            tps = proc_tokens_log / max(elapsed_log, 1e-6)
            ss = slab_stats(optimizer)
            grp_stats = group_stats(optimizer)
            grp_upd_stats = group_updates(pre_step_snapshot) if pre_step_snapshot is not None else {}
            ln_group = 'shared/decoder_norm' if 'shared/decoder_norm' in grp_stats else \
                ('shared/misc' if 'shared/misc' in grp_stats else None)
            learnorm_stats = None
            if ln_group is not None:
                ent = grp_stats[ln_group]
                upd = grp_upd_stats.get(ln_group, {}).get('upd_w', 0.0)
                learnorm_stats = {
                    'params': ent['params'],
                    'numel': ent['numel'],
                    'g': ent['g'],
                    'w': ent['w'],
                    'upd_w': upd,
                }
            if val_loss_raw is not None:
                _ema_update('val', val_loss_raw)
                val_loss = _loss_ema['val']
            else:
                val_loss = None
            box = render_log(
                # Header `lss` is the EMA of the weighted mixed task loss:
                # the same objective accumulated in backward above.
                step=step, loss=_loss_ema.get('mixed', mixed_loss),
                lr=base_lr_now, N=N,
                row_len=row_len, tok_per_sec=tps,
                grad_tokens_total=proc_tokens_total,
                seq_len=args.seq_len, grid=(grid_g, grid_h, grid_w),
                stats=ss, grp_stats=grp_stats,
                actual_update_stats=grp_upd_stats,
                learnorm_stats=learnorm_stats,
                # Logged values are EMAs (alpha=0.05, ~40-step half-life).
                # Per-batch instantaneous losses on heterogeneous corpora
                # swing ±1 nat between consecutive steps and visually
                # mask the underlying monotonic trend; the smoothed view
                # tracks real progress without lagging behind LR moves.
                # Instantaneous per-task values remain in `per_task_losses`
                # for callers that want raw signal (e.g. NaN diagnostics).
                per_task_losses={
                    **{k: _loss_ema.get(k, v) for k, v in per_task_losses.items()},
                    'ce': _loss_ema.get('ce', mixed_ce_loss),
                    'lnorm': _loss_ema.get('lnorm', mixed_learnorm_loss),
                    'lam': float(args.learnorm_lambda),
                    'pred': _loss_ema.get('pred', mixed_predictor_loss),
                    'pred_lam': float(args.predictor_lambda),
                    'vec': _loss_ema.get('vec', mixed_vector_error_loss),
                    'vec_lam': float(args.vector_error_lambda),
                    'src': _loss_ema.get('dsl_src', mixed_dsl_src_loss),
                    'src_lam': float(args.dsl_source_lambda),
                },
                container_k=args.container_k,
                container_runtime_stats=getattr(model_base, '_container_last_stats', None),
                n_tasks=len(task_inputs),
                ckpt_saved=ckpt_saved_now, skipped=False,
                val_loss=val_loss,
                xmem_retention=_xmem_retention(model_base),
            )
            print(box, flush=True)
            if bool(args.dsl_log) and bool(args.use_dsl) and int(args.dsl_lanes) > 0:
                dsl_box = render_dsl_log(getattr(model_base, '_dsl_last_stats', None))
                if dsl_box:
                    print(dsl_box, flush=True)
            t_log = now
            proc_tokens_log = 0
        if dtype_audit is not None and step == int(args.dtype_audit_steps):
            for h in dtype_audit_hooks:
                h.remove()
            dtype_audit_hooks = []
            rows = []
            for mod_key, by_dt in dtype_audit.items():
                total = sum(by_dt.values())
                fp32 = by_dt.get('float32', 0)
                if total > 0 and fp32 > 0:
                    rows.append((fp32, total, mod_key, by_dt))
            rows.sort(reverse=True, key=lambda x: x[0])
            print('[dtype_audit] top fp32 output modules:', flush=True)
            for fp32, total, mod_key, by_dt in rows[:30]:
                share = 100.0 * fp32 / max(total, 1)
                print(f'  {mod_key:<64s} fp32={fp32:>12d} ({share:5.1f}%) dtypes={by_dt}', flush=True)

    if distributed:
        dist.destroy_process_group()


def _segments_with_user_mask(user_mask: torch.Tensor, reset_mask: torch.Tensor) -> torch.Tensor:
    """Return per-token mask of segments that contain at least one user token.

    user_mask/reset_mask are [B, L] bool tensors from packing. A segment is a
    packed sub-sample delimited by reset_mask=True at its first token.
    """
    if user_mask.shape != reset_mask.shape:
        raise ValueError(
            f'user_mask shape {tuple(user_mask.shape)} != reset_mask shape {tuple(reset_mask.shape)}'
        )
    bsz, seqlen = user_mask.shape
    out = torch.zeros_like(user_mask)
    for b in range(bsz):
        starts = torch.nonzero(reset_mask[b], as_tuple=False).flatten()
        if starts.numel() == 0:
            starts = torch.tensor([0], device=reset_mask.device, dtype=torch.long)
        starts = starts.tolist()
        for i, s in enumerate(starts):
            e = starts[i + 1] if i + 1 < len(starts) else seqlen
            if bool(user_mask[b, s:e].any().item()):
                out[b, s:e] = True
    return out


def _parse_modes(spec: str) -> set[str]:
    """Parse task-mode list from CLI.

    Supported modes:
      - pt    : main supervised objective (former "std")
      - blind : full blind reconstruction
      - rev   : reverse user reconstruction
      - fim   : corruption-recovery objective
      - nat   : blind-tail prefix-completion with PAD input suffix
    """
    allowed = {'pt', 'nat', 'blind', 'rev', 'fim'}
    raw = [x.strip().lower() for x in (spec or '').split(',') if x.strip()]
    modes = set(raw)
    bad = sorted(m for m in modes if m not in allowed)
    if bad:
        raise ValueError(f'unknown --modes entries: {bad}; allowed={sorted(allowed)}')
    if not modes:
        raise ValueError('--modes produced empty set; provide at least one mode')
    return modes


if __name__ == '__main__':
    main()