modeling_cubic.py

#!/usr/bin/env python3
"""One-file training pipeline for the faithful ~157M CubicV7-Hier LM.

Stages
------
0. Train a 32K byte-level BPE and prepare deterministic local caches.
1. 3B-token multilingual/code/math base pretraining.
2. Dialogue/instruction SFT (SmolTalk + Aya EN/RU).
3. Verified reasoning SFT (OpenR1-Math-220k, shortest correct trace).
4. Offline human-preference alignment with DPO
   (Anthropic HH-RLHF + corrected UltraFeedback).
5. Evaluate, generate samples, save safetensors and upload a complete folder
   to asilarkness/llm-150m.

This file deliberately contains the model, faithful Cubic depth memory, the
hierarchical late retrieval path, Muon+AdamW, data preparation, training,
resume, inference and Hub publishing.  No hidden local modules are required.

Useful environment overrides:
  CUBIC_MODE=train|chat             (default: train)
  CUBIC_ROOT=/marimo/cubic_hier_150m_gemma4
  CUBIC_PRETRAIN_TOKENS=3000000000
  CUBIC_SEQ_LEN=4096
  CUBIC_MICRO_BATCH=16
  CUBIC_ACCUM=1
  CUBIC_COMPILE=1
  CUBIC_COMPILE_MODE=default
  CUBIC_CHECKPOINT_EVERY=0
  CUBIC_HF_CHECKPOINT_UPLOAD=1
  CUBIC_HF_CHECKPOINT_REPO=asilarkness/next
  CUBIC_HF_CHECKPOINT_MINUTES=60
  CUBIC_REASONING=0|1              (chat mode; direct or <think> mode)
  CUBIC_HF_REPO=asilarkness/llm-150m
  CUBIC_UPLOAD=1
  HF_TOKEN=...                      (optional after `hf auth login`)

The default run is intentionally a real training run, not a smoke test.
Set CUBIC_PRETRAIN_TOKENS=2000000 and the dataset limits below for a short
end-to-end test before committing GPU-days.
"""

from __future__ import annotations

import copy
import gc
import hashlib
import importlib.util
import json
import math
import os
import random
import re
import shutil
import subprocess
import sys
import threading
import time
from array import array
from collections import defaultdict
from contextlib import nullcontext
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Iterator, Sequence


def ensure_dependencies() -> None:
    packages = {
        "torch": "torch",
        "datasets": "datasets>=3.2.0",
        "tokenizers": "tokenizers>=0.20.0",
        "huggingface_hub": "huggingface_hub[hf_xet]>=0.27.0",
        "safetensors": "safetensors>=0.4.5",
        "liger_kernel": "liger-kernel>=0.5.2",
    }
    missing = [spec for module, spec in packages.items() if importlib.util.find_spec(module) is None]
    if missing:
        print("Installing dependencies:", ", ".join(missing), flush=True)
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])


os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")
os.environ.setdefault("RAYON_NUM_THREADS", os.environ.get("CUBIC_TOKENIZER_THREADS", "4"))
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
ensure_dependencies()

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint


# -------------------------------------------------------------------------------------------------
# Configuration
# -------------------------------------------------------------------------------------------------


def env_int(name: str, default: int) -> int:
    return int(os.environ.get(name, str(default)).replace("_", ""))


def env_float(name: str, default: float) -> float:
    return float(os.environ.get(name, str(default)))


@dataclass(frozen=True)
class Config:
    vocab_size: int = 32_768
    dim: int = 800
    depth: int = 16
    n_heads: int = 10
    mlp_hidden: int = 2_176
    depth_rank: int = 200
    depth_heads: int = 5
    seq_len: int = env_int("CUBIC_SEQ_LEN", 4_096)
    max_context: int = env_int("CUBIC_MAX_CONTEXT", 8_192)
    group_chunk: int = env_int("CUBIC_GROUP_CHUNK", 128)
    group_keep: int = env_int("CUBIC_GROUP_KEEP", 4)
    initial_mix: float = 0.18
    recency_bias: float = 0.30
    rope_base: float = 500_000.0
    norm_eps: float = 1e-5

    @property
    def head_dim(self) -> int:
        return self.dim // self.n_heads

    @property
    def depth_head_dim(self) -> int:
        return self.depth_rank // self.depth_heads


CFG = Config()
EXPECTED_PARAMETER_COUNT = 157_065_144
if CFG.dim % CFG.n_heads or CFG.depth_rank % CFG.depth_heads:
    raise ValueError("dim/head counts are not divisible")
if CFG.seq_len > CFG.max_context:
    raise ValueError("CUBIC_SEQ_LEN must not exceed CUBIC_MAX_CONTEXT")

SEED = env_int("CUBIC_SEED", 20_260_722)
MODE = os.environ.get("CUBIC_MODE", "train").lower()
ARCH_VERSION = "cubic-hier-v7-gemma4-150m-v2"
# A new default root deliberately prevents an older v1 cache/checkpoint from
# being silently reused with the changed tokenizer mixture and architecture.
ROOT = Path(os.environ.get("CUBIC_ROOT", "/marimo/cubic_hier_150m_gemma4"))
WORK = ROOT / "work"
CACHE = ROOT / "cache"
CHECKPOINTS = ROOT / "checkpoints"
EXPORT = ROOT / "export"
TOKENIZER_PATH = WORK / "tokenizer.json"
LATEST = CHECKPOINTS / "latest.pt"
BEST = CHECKPOINTS / "best.pt"
FINAL = CHECKPOINTS / "final.pt"
DPO_REFERENCE = CHECKPOINTS / "dpo_reference.pt"
REPORT = CHECKPOINTS / "samples.jsonl"
HF_REPO = os.environ.get("CUBIC_HF_REPO", "asilarkness/llm-150m")

PRETRAIN_TOKENS = env_int("CUBIC_PRETRAIN_TOKENS", 3_000_000_000)
PRETRAIN_VAL_TOKENS = env_int("CUBIC_PRETRAIN_VAL_TOKENS", 8_000_000)
TOKENIZER_DOCS = env_int("CUBIC_TOKENIZER_DOCS", 200_000)
SMOLTALK_LIMIT = env_int("CUBIC_SMOLTALK_LIMIT", 500_000)
AYA_LIMIT = env_int("CUBIC_AYA_LIMIT", 200_000)
REASON_LIMIT = env_int("CUBIC_REASON_LIMIT", 100_000)
HH_LIMIT = env_int("CUBIC_HH_LIMIT", 150_000)
UF_LIMIT = env_int("CUBIC_UF_LIMIT", 100_000)

# RTX PRO 6000 96 GB: batch 16 with no activation recomputation is faster per
# token and uses less peak memory than batch 32 with half the blocks recomputed.
MICRO_BATCH = env_int("CUBIC_MICRO_BATCH", 16)
GRAD_ACCUM = env_int("CUBIC_ACCUM", 1)
# Same effective DPO batch (8) and exactly the same eight examples per update,
# but two launches instead of eight.
DPO_MICRO_BATCH = env_int("CUBIC_DPO_MICRO_BATCH", 4)
DPO_ACCUM = env_int("CUBIC_DPO_ACCUM", 2)
SFT_EPOCHS = env_float("CUBIC_SFT_EPOCHS", 1.0)
REASON_EPOCHS = env_float("CUBIC_REASON_EPOCHS", 1.5)
DPO_EPOCHS = env_float("CUBIC_DPO_EPOCHS", 1.0)
DPO_BETA = env_float("CUBIC_DPO_BETA", 0.10)
DPO_SFT_WEIGHT = env_float("CUBIC_DPO_SFT_WEIGHT", 0.05)
MTP_PRETRAIN_WEIGHT = env_float("CUBIC_MTP_PRETRAIN_WEIGHT", 0.30)
MTP_DECAY_WEIGHT = env_float("CUBIC_MTP_DECAY_WEIGHT", 0.10)
MTP_SFT_WEIGHT = env_float("CUBIC_MTP_SFT_WEIGHT", 0.10)

COMPILE = os.environ.get("CUBIC_COMPILE", "1") == "1"
COMPILE_MODE = os.environ.get("CUBIC_COMPILE_MODE", "default")
# Checkpoint every second block. This preserves forward/backward mathematics
# while spending the otherwise idle 96 GB VRAM to avoid half of recomputation.
# Use 1 for the old low-memory behavior, or 0 to disable recomputation entirely.
CHECKPOINT_EVERY = env_int("CUBIC_CHECKPOINT_EVERY", 0)
if CHECKPOINT_EVERY < 0:
    raise ValueError("CUBIC_CHECKPOINT_EVERY must be >= 0")
UPLOAD = os.environ.get("CUBIC_UPLOAD", "1") == "1"
CHECKPOINT_MINUTES = env_float("CUBIC_CHECKPOINT_MINUTES", 30.0)
HF_CHECKPOINT_UPLOAD = os.environ.get("CUBIC_HF_CHECKPOINT_UPLOAD", "1") == "1"
HF_CHECKPOINT_REPO = os.environ.get("CUBIC_HF_CHECKPOINT_REPO", "asilarkness/next")
HF_CHECKPOINT_MINUTES = env_float("CUBIC_HF_CHECKPOINT_MINUTES", 60.0)
HF_RESUME_DOWNLOAD = os.environ.get("CUBIC_HF_RESUME_DOWNLOAD", "1") == "1"
ALLOW_COLD_START = os.environ.get("CUBIC_ALLOW_COLD_START", "0") == "1"
LOG_EVERY = env_int("CUBIC_LOG_EVERY", 10)
EVAL_EVERY = env_int("CUBIC_EVAL_EVERY", 500)
EVAL_BATCHES = env_int("CUBIC_EVAL_BATCHES", 16)
EVAL_BATCH_SIZE = env_int("CUBIC_EVAL_BATCH_SIZE", 8)
GRAD_CLIP = 1.0

PRETRAIN_TOKENS_FILE = CACHE / "pretrain_train_u16.bin"
PRETRAIN_VAL_FILE = CACHE / "pretrain_val_u16.bin"
SFT_TOKENS_FILE = CACHE / "sft_tokens_u16.bin"
SFT_MASK_FILE = CACHE / "sft_mask_u8.bin"
REASON_TOKENS_FILE = CACHE / "reason_tokens_u16.bin"
REASON_MASK_FILE = CACHE / "reason_mask_u8.bin"
PREF_FILES = {
    "chosen_tokens": CACHE / "pref_chosen_tokens_u16.bin",
    "chosen_mask": CACHE / "pref_chosen_mask_u8.bin",
    "rejected_tokens": CACHE / "pref_rejected_tokens_u16.bin",
    "rejected_mask": CACHE / "pref_rejected_mask_u8.bin",
}
META_FILE = CACHE / "metadata.json"

SPECIAL_TOKENS = (
    "<pad>", "<bos>", "<eos>", "<unk>",
    "<|system|>", "<|user|>", "<|assistant|>", "<think>", "</think>",
)


def setup_runtime() -> str:
    for path in (ROOT, WORK, CACHE, CHECKPOINTS, EXPORT):
        path.mkdir(parents=True, exist_ok=True)
    random.seed(SEED)
    torch.manual_seed(SEED)
    if not torch.cuda.is_available():
        raise RuntimeError("This full run requires a CUDA GPU")
    torch.cuda.manual_seed_all(SEED)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.set_float32_matmul_precision("high")
    return "cuda"


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
AMP_DTYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16


def amp_context():
    return torch.autocast("cuda", dtype=AMP_DTYPE) if DEVICE == "cuda" else nullcontext()


def stable_hash(value: str) -> int:
    return int.from_bytes(hashlib.blake2b(value.encode("utf-8", "ignore"), digest_size=8).digest(), "little")


def atomic_json(data: Any, path: Path) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    os.replace(temporary, path)


# -------------------------------------------------------------------------------------------------
# Dataset streaming, tokenizer and compact caches
# -------------------------------------------------------------------------------------------------


def stream_pretrain(seed: int = SEED) -> Iterator[str]:
    """Deterministic quality mixture: web EN/RU, deduplicated code and math.

    The 72/13/10/5 mix keeps the bilingual identity of the original run but
    fixes its largest pretraining hole: virtually no code or mathematical
    language before reasoning SFT.  All four inputs are streamed, so the
    one-file pipeline still needs no manual corpus download.
    """
    from datasets import load_dataset

    sources = {
        "web_en": iter(load_dataset(
        "HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True,
        ).shuffle(seed=seed, buffer_size=10_000)),
        "web_ru": iter(load_dataset(
        "HuggingFaceFW/fineweb-2", name="rus_Cyrl", split="train", streaming=True,
        ).shuffle(seed=seed + 1, buffer_size=10_000)),
        "code": iter(load_dataset(
            "codeparrot/codeparrot-clean", split="train", streaming=True,
        ).shuffle(seed=seed + 2, buffer_size=10_000)),
        "math": iter(load_dataset(
            "HuggingFaceTB/finemath", name="finemath-4plus", split="train", streaming=True,
        ).shuffle(seed=seed + 3, buffer_size=10_000)),
    }
    probabilities = (("web_en", 0.72), ("web_ru", 0.13), ("code", 0.10), ("math", 0.05))
    rng = random.Random(seed)
    while True:
        draw = rng.random()
        cumulative = 0.0
        source_name = probabilities[-1][0]
        for name, probability in probabilities:
            cumulative += probability
            if draw < cumulative:
                source_name = name
                break
        try:
            record = next(sources[source_name])
        except StopIteration:
            # Every selected source is much larger than this 3B-token run. If
            # a Hub revision ever changes that, fail loudly instead of silently
            # changing the declared mixture.
            raise RuntimeError(f"pretraining source ended early: {source_name}")
        text = record.get("content", "") if source_name == "code" else record.get("text", "")
        if isinstance(text, str) and len(text) >= 200:
            yield text


def train_tokenizer():
    from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers

    if TOKENIZER_PATH.exists():
        print(f"Tokenizer already present: {TOKENIZER_PATH}")
        return Tokenizer.from_file(str(TOKENIZER_PATH))

    print(f"Training byte-level BPE: vocab={CFG.vocab_size:,}, documents={TOKENIZER_DOCS:,}")
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.normalizer = normalizers.NFC()
    # Gemma-style digit splitting improves arithmetic compositionality, while
    # ByteLevel preserves whitespace/code formatting and guarantees byte-level
    # coverage through its complete initial alphabet.
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.Digits(individual_digits=True),
        pre_tokenizers.ByteLevel(add_prefix_space=False),
    ])
    tokenizer.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=CFG.vocab_size,
        min_frequency=2,
        special_tokens=list(SPECIAL_TOKENS),
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=True,
    )

    def texts() -> Iterator[str]:
        for index, text in enumerate(stream_pretrain(SEED + 10)):
            if index >= TOKENIZER_DOCS:
                break
            yield text

    tokenizer.train_from_iterator(texts(), trainer=trainer, length=TOKENIZER_DOCS)
    if tokenizer.get_vocab_size() != CFG.vocab_size:
        raise RuntimeError(f"tokenizer produced {tokenizer.get_vocab_size()}, expected {CFG.vocab_size}")
    tokenizer.save(str(TOKENIZER_PATH))
    return tokenizer


class FlatU16Writer:
    def __init__(self, path: Path):
        self.path = path
        # An interrupted multi-hour corpus build resumes by appending. It may
        # repeat a small prefix of documents, but never discards gigabytes of
        # already-tokenized data and never publishes a half-written final file.
        self.count = path.stat().st_size // 2 if path.exists() else 0
        self.file = path.open("ab")
        self.buffer = array("H")

    def add(self, ids: Sequence[int], limit: int | None = None) -> int:
        if limit is not None:
            ids = ids[: max(0, limit - self.count)]
        self.buffer.extend(ids)
        self.count += len(ids)
        if len(self.buffer) >= 2_000_000:
            self.buffer.tofile(self.file)
            self.buffer = array("H")
        return len(ids)

    def flush(self, durable: bool = False) -> None:
        if self.buffer:
            self.buffer.tofile(self.file)
            self.buffer = array("H")
        self.file.flush()
        if durable:
            os.fsync(self.file.fileno())

    def close(self) -> None:
        self.flush()
        self.file.close()


def prepare_pretrain_cache(tokenizer) -> dict[str, int]:
    wanted_train_bytes = PRETRAIN_TOKENS * 2
    if (
        PRETRAIN_TOKENS_FILE.exists() and PRETRAIN_VAL_FILE.exists()
        and PRETRAIN_TOKENS_FILE.stat().st_size >= wanted_train_bytes
        and PRETRAIN_VAL_FILE.stat().st_size >= PRETRAIN_VAL_TOKENS * 2
    ):
        return {
            "pretrain_tokens": PRETRAIN_TOKENS_FILE.stat().st_size // 2,
            "pretrain_val_tokens": PRETRAIN_VAL_FILE.stat().st_size // 2,
        }

    print(f"Encoding {PRETRAIN_TOKENS / 1e9:.2f}B pretraining tokens ...")
    train_tmp = PRETRAIN_TOKENS_FILE.with_suffix(".tmp")
    val_tmp = PRETRAIN_VAL_FILE.with_suffix(".tmp")
    progress_path = CACHE / "pretrain_encoding_progress.json"
    encode_batch_size = env_int("CUBIC_ENCODING_BATCH", 256)
    train = FlatU16Writer(train_tmp)
    val_ready = (
        PRETRAIN_VAL_FILE.exists()
        and PRETRAIN_VAL_FILE.stat().st_size >= PRETRAIN_VAL_TOKENS * 2
    )
    val = None if val_ready else FlatU16Writer(val_tmp)
    val_count = PRETRAIN_VAL_TOKENS if val_ready else val.count
    resume_documents = 0
    if progress_path.exists():
        try:
            progress = json.loads(progress_path.read_text(encoding="utf-8"))
            saved_train = int(progress.get("train_tokens", -1))
            saved_val = int(progress.get("val_tokens", -1))
            # A hard SIGKILL may leave a few post-cursor bytes on disk. Roll
            # those uncommitted bytes back to the last atomic cursor instead
            # of abandoning the cursor and restarting the entire stream.
            if 0 <= saved_train <= train.count and 0 <= saved_val <= val_count:
                if saved_train != train.count:
                    train.file.truncate(saved_train * 2)
                    train.count = saved_train
                if val is not None and saved_val != val.count:
                    val.file.truncate(saved_val * 2)
                    val.count = saved_val
                    val_count = saved_val
                resume_documents = int(progress.get("documents", 0))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            resume_documents = 0
    eos = tokenizer.token_to_id("<eos>")
    started = time.perf_counter()
    initial_total = train.count + val_count
    print(
        f"  resume train={train.count / 1e9:.3f}B val={val_count / 1e6:.1f}M "
        f"documents={resume_documents:,} batch={encode_batch_size} "
        f"threads={os.environ['RAYON_NUM_THREADS']}",
        flush=True,
    )

    def save_progress(documents: int) -> None:
        # Cursor is recorded only after all corresponding token bytes are
        # durable, so a hard process/container restart cannot skip data.
        train.flush(durable=True)
        if val is not None:
            val.flush(durable=True)
        atomic_json({
            "documents": documents,
            "train_tokens": train.count,
            "val_tokens": val_count if val is None else val.count,
            "updated_at": time.time(),
        }, progress_path)

    try:
        stream = enumerate(stream_pretrain(SEED + 20))
        pending_texts: list[str] = []
        pending_last_document = resume_documents
        completed_documents = resume_documents
        for document_index, text in stream:
            if document_index < resume_documents:
                continue
            pending_texts.append(text)
            pending_last_document = document_index + 1
            if len(pending_texts) < encode_batch_size:
                continue

            encodings = tokenizer.encode_batch(pending_texts)
            for source_text, encoding in zip(pending_texts, encodings):
                ids = encoding.ids + [eos]
                # Once validation is complete, every document goes directly
                # into train; this avoids rebuilding/duplicating val on resume.
                use_val = (
                    val is not None and val.count < PRETRAIN_VAL_TOKENS
                    and stable_hash(source_text[:512]) % 101 == 0
                )
                if use_val:
                    val.add(ids, PRETRAIN_VAL_TOKENS)
                else:
                    train.add(ids, PRETRAIN_TOKENS)
            pending_texts.clear()
            completed_documents = pending_last_document

            if pending_last_document % (encode_batch_size * 20) == 0:
                save_progress(completed_documents)
                current_val = val_count if val is None else val.count
                new_tokens = max(0, train.count + current_val - initial_total)
                rate = new_tokens / max(1e-6, time.perf_counter() - started)
                print(
                    f"  documents={pending_last_document:,} "
                    f"train={train.count / 1e9:.3f}B val={current_val / 1e6:.1f}M "
                    f"{rate:,.0f} tok/s",
                    flush=True,
                )
            current_val = val_count if val is None else val.count
            if train.count >= PRETRAIN_TOKENS and current_val >= PRETRAIN_VAL_TOKENS:
                break
    finally:
        # KeyboardInterrupt, notebook stop and ordinary exceptions all leave a
        # durable cursor. At worst a partly processed batch is duplicated;
        # already encoded gigabytes are never discarded or skipped.
        save_progress(completed_documents)
        train.close()
        if val is not None:
            val.close()
    final_val_count = val_count if val is None else val.count
    if train.count < PRETRAIN_TOKENS or final_val_count < PRETRAIN_VAL_TOKENS:
        raise RuntimeError("pretraining stream ended before cache targets were reached")
    os.replace(train_tmp, PRETRAIN_TOKENS_FILE)
    if not val_ready:
        os.replace(val_tmp, PRETRAIN_VAL_FILE)
    elif val_tmp.exists():
        # Stale partial validation data is no longer needed after the complete
        # final validation file has been verified.
        val_tmp.unlink()
    progress_path.unlink(missing_ok=True)
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("Pretraining cache finalized; CPU encoding memory released. Starting GPU training.", flush=True)
    return {"pretrain_tokens": train.count, "pretrain_val_tokens": final_val_count}


class PackedRecords:
    """Packs variable-length supervised samples into fixed seq_len+1 records."""

    def __init__(self, token_path: Path, mask_path: Path, pad_id: int):
        self.token_tmp = token_path.with_suffix(".tmp")
        self.mask_tmp = mask_path.with_suffix(".tmp")
        self.token_path, self.mask_path = token_path, mask_path
        self.token_file = self.token_tmp.open("wb")
        self.mask_file = self.mask_tmp.open("wb")
        self.ids: list[int] = []
        self.mask: list[int] = []
        self.pad_id = pad_id
        self.record_len = CFG.seq_len + 1
        self.records = 0

    def _flush(self) -> None:
        if not self.ids:
            return
        padding = self.record_len - len(self.ids)
        array("H", self.ids + [self.pad_id] * padding).tofile(self.token_file)
        array("B", self.mask + [0] * padding).tofile(self.mask_file)
        self.ids.clear()
        self.mask.clear()
        self.records += 1

    def add(self, ids: Sequence[int], mask: Sequence[int]) -> None:
        if len(ids) != len(mask):
            raise ValueError("token/mask length mismatch")
        # Preserve the final answer if an example exceeds the context.
        if len(ids) > self.record_len:
            supervised = next((i for i, value in enumerate(mask) if value), len(ids) - 1)
            start = max(0, min(supervised - self.record_len // 3, len(ids) - self.record_len))
            ids, mask = ids[start:start + self.record_len], mask[start:start + self.record_len]
        if len(self.ids) + len(ids) > self.record_len:
            self._flush()
        self.ids.extend(ids)
        self.mask.extend(mask)

    def close(self) -> int:
        self._flush()
        self.token_file.close()
        self.mask_file.close()
        os.replace(self.token_tmp, self.token_path)
        os.replace(self.mask_tmp, self.mask_path)
        return self.records


def role_id(tokenizer, role: str) -> int:
    role = "assistant" if role in ("assistant", "gpt", "model") else "user" if role in ("user", "human") else "system"
    return tokenizer.token_to_id(f"<|{role}|>")


def normalize_messages(value: Any) -> list[dict[str, str]]:
    if not isinstance(value, list):
        return []
    result = []
    for item in value:
        if not isinstance(item, dict):
            continue
        role = str(item.get("role", item.get("from", ""))).lower()
        content = item.get("content", item.get("value", ""))
        if role and isinstance(content, str) and content.strip():
            result.append({"role": role, "content": content.strip()})
    return result


DIRECT_SYSTEM = (
    "Answer directly, accurately and concisely. Do not invent facts. "
    "If the request is ambiguous, state the assumption you use."
)
REASONING_SYSTEM = (
    "Solve the task carefully. Put the internal solution trace between "
    "<think> and </think>, then give a clear final answer."
)


def with_system_mode(messages: Sequence[dict[str, str]], reasoning: bool) -> list[dict[str, str]]:
    result = list(messages)
    if result and result[0]["role"] == "system":
        return result
    instruction = REASONING_SYSTEM if reasoning else DIRECT_SYSTEM
    return [{"role": "system", "content": instruction}, *result]


def encode_chat(tokenizer, messages: Sequence[dict[str, str]], supervise_assistant: bool = True) -> tuple[list[int], list[int]]:
    bos, eos = tokenizer.token_to_id("<bos>"), tokenizer.token_to_id("<eos>")
    ids, mask = [bos], [0]
    for message in messages:
        role = str(message["role"]).lower()
        content = str(message["content"])
        segment = [role_id(tokenizer, role), *tokenizer.encode("\n" + content + "\n").ids, eos]
        supervised = supervise_assistant and role in ("assistant", "gpt", "model")
        ids.extend(segment)
        mask.extend([int(supervised)] * len(segment))
    return ids, mask


def cached_record_count(token_path: Path) -> int:
    return token_path.stat().st_size // (2 * (CFG.seq_len + 1))


def prepare_dialogue_cache(tokenizer) -> int:
    if SFT_TOKENS_FILE.exists() and SFT_MASK_FILE.exists():
        return cached_record_count(SFT_TOKENS_FILE)
    from datasets import load_dataset

    print("Preparing dialogue SFT cache (SmolTalk + Aya EN/RU) ...")
    writer = PackedRecords(SFT_TOKENS_FILE, SFT_MASK_FILE, tokenizer.token_to_id("<pad>"))
    accepted = 0
    seen: set[int] = set()
    smol = load_dataset("HuggingFaceTB/smol-smoltalk", split="train", streaming=True)
    for record in smol:
        messages = normalize_messages(record.get("messages", record.get("conversations")))
        fingerprint = stable_hash(json.dumps(messages, ensure_ascii=False, sort_keys=True))
        if (
            messages and fingerprint not in seen
            and any(m["role"] in ("assistant", "gpt", "model") for m in messages)
        ):
            seen.add(fingerprint)
            writer.add(*encode_chat(tokenizer, with_system_mode(messages, reasoning=False)))
            accepted += 1
        if accepted >= SMOLTALK_LIMIT:
            break
    aya_count = 0
    aya = load_dataset("CohereLabs/aya_dataset", split="train", streaming=True)
    for record in aya:
        language = str(record.get("language", "")).lower()
        if language not in ("english", "russian", "eng", "rus", "en", "ru"):
            continue
        prompt, answer = record.get("inputs"), record.get("targets")
        if isinstance(prompt, str) and isinstance(answer, str) and prompt.strip() and answer.strip():
            fingerprint = stable_hash(prompt.strip() + "\0" + answer.strip())
            if fingerprint in seen:
                continue
            seen.add(fingerprint)
            writer.add(*encode_chat(tokenizer, with_system_mode([
                {"role": "user", "content": prompt}, {"role": "assistant", "content": answer},
            ], reasoning=False)))
            aya_count += 1
        if aya_count >= AYA_LIMIT:
            break
    records = writer.close()
    print(f"  dialogue examples={accepted + aya_count:,}, packed records={records:,}")
    return records


def shortest_correct_reasoning(record: dict[str, Any]) -> str:
    generations = record.get("generations") or []
    if isinstance(generations, str):
        generations = [generations]
    math_flags = record.get("correctness_math_verify") or []
    judge_flags = record.get("correctness_llama") or []
    complete_flags = record.get("is_reasoning_complete") or []
    candidates = []
    for index, generation in enumerate(generations):
        if not isinstance(generation, str) or not generation.strip():
            continue
        verified = (
            (index < len(math_flags) and bool(math_flags[index]))
            or (index < len(judge_flags) and bool(judge_flags[index]))
        )
        complete = not complete_flags or index >= len(complete_flags) or bool(complete_flags[index])
        if verified and complete:
            candidates.append(generation.strip())
    if candidates:
        return min(candidates, key=len)
    solution = record.get("solution", "")
    return solution.strip() if isinstance(solution, str) else ""


def prepare_reasoning_cache(tokenizer) -> int:
    if REASON_TOKENS_FILE.exists() and REASON_MASK_FILE.exists():
        return cached_record_count(REASON_TOKENS_FILE)
    from datasets import load_dataset

    print("Preparing verified reasoning SFT cache (OpenR1-Math-220k) ...")
    writer = PackedRecords(REASON_TOKENS_FILE, REASON_MASK_FILE, tokenizer.token_to_id("<pad>"))
    accepted = 0
    seen: set[int] = set()
    dataset = load_dataset("open-r1/OpenR1-Math-220k", name="default", split="train", streaming=True)
    for record in dataset:
        problem = record.get("problem", "")
        answer = shortest_correct_reasoning(record)
        if not isinstance(problem, str) or not problem.strip() or not answer:
            continue
        fingerprint = stable_hash(problem.strip() + "\0" + answer)
        if fingerprint in seen:
            continue
        seen.add(fingerprint)
        if "<think>" not in answer:
            answer = f"<think>\n{answer}\n</think>"
        writer.add(*encode_chat(tokenizer, with_system_mode([
            {"role": "user", "content": problem.strip()},
            {"role": "assistant", "content": answer},
        ], reasoning=True)))
        accepted += 1
        if accepted >= REASON_LIMIT:
            break
    records = writer.close()
    print(f"  reasoning examples={accepted:,}, packed records={records:,}")
    return records


def parse_hh_transcript(text: str) -> list[dict[str, str]]:
    """Converts Anthropic's '\n\nHuman:' transcript format to messages."""
    if not isinstance(text, str):
        return []
    pieces = re.split(r"\n\n(Human|Assistant):\s*", text)
    messages: list[dict[str, str]] = []
    # pieces: optional prefix, role, content, role, content, ...
    for index in range(1, len(pieces) - 1, 2):
        role = "user" if pieces[index] == "Human" else "assistant"
        content = pieces[index + 1].strip()
        if content:
            messages.append({"role": role, "content": content})
    return messages


def preference_pair_from_messages(
    chosen: Sequence[dict[str, str]], rejected: Sequence[dict[str, str]],
) -> tuple[list[dict[str, str]], str, str] | None:
    if not chosen or not rejected:
        return None
    if chosen[-1]["role"] != "assistant" or rejected[-1]["role"] != "assistant":
        return None
    prompt_chosen, prompt_rejected = list(chosen[:-1]), list(rejected[:-1])
    # The prompt should be identical. A conservative common-prefix fallback
    # protects against harmless metadata differences without inventing context.
    common: list[dict[str, str]] = []
    for left, right in zip(prompt_chosen, prompt_rejected):
        if left != right:
            break
        common.append(left)
    if not common:
        return None
    return common, chosen[-1]["content"], rejected[-1]["content"]


def encode_preference_side(tokenizer, prompt: Sequence[dict[str, str]], response: str) -> tuple[list[int], list[int]]:
    prompt_ids, _ = encode_chat(
        tokenizer, with_system_mode(prompt, reasoning=False), supervise_assistant=False,
    )
    eos = tokenizer.token_to_id("<eos>")
    response_ids = [
        tokenizer.token_to_id("<|assistant|>"),
        *tokenizer.encode("\n" + response.strip() + "\n").ids,
        eos,
    ]
    ids = prompt_ids + response_ids
    mask = [0] * len(prompt_ids) + [1] * len(response_ids)
    limit = CFG.seq_len + 1
    if len(ids) > limit:
        # Keep the entire preference signal whenever possible and trim the
        # oldest prompt turns first.
        first_response = len(prompt_ids)
        start = max(0, min(first_response - limit // 4, len(ids) - limit))
        ids, mask = ids[start:start + limit], mask[start:start + limit]
    return ids, mask


class PreferenceWriter:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.record_len = CFG.seq_len + 1
        self.pad = tokenizer.token_to_id("<pad>")
        self.temporary = {name: path.with_suffix(".tmp") for name, path in PREF_FILES.items()}
        self.handles = {name: path.open("wb") for name, path in self.temporary.items()}
        self.records = 0

    def _fixed(self, values: Sequence[int], pad: int) -> list[int]:
        return list(values[:self.record_len]) + [pad] * max(0, self.record_len - len(values))

    def add(self, prompt: Sequence[dict[str, str]], chosen: str, rejected: str) -> None:
        chosen_ids, chosen_mask = encode_preference_side(self.tokenizer, prompt, chosen)
        rejected_ids, rejected_mask = encode_preference_side(self.tokenizer, prompt, rejected)
        if not any(chosen_mask) or not any(rejected_mask):
            return
        array("H", self._fixed(chosen_ids, self.pad)).tofile(self.handles["chosen_tokens"])
        array("B", self._fixed(chosen_mask, 0)).tofile(self.handles["chosen_mask"])
        array("H", self._fixed(rejected_ids, self.pad)).tofile(self.handles["rejected_tokens"])
        array("B", self._fixed(rejected_mask, 0)).tofile(self.handles["rejected_mask"])
        self.records += 1

    def close(self) -> int:
        for handle in self.handles.values():
            handle.close()
        for name, final in PREF_FILES.items():
            os.replace(self.temporary[name], final)
        return self.records


def prepare_preference_cache(tokenizer) -> int:
    if all(path.exists() for path in PREF_FILES.values()):
        return cached_record_count(PREF_FILES["chosen_tokens"])
    from datasets import load_dataset

    print("Preparing DPO pairs (Anthropic HH-RLHF + corrected UltraFeedback) ...")
    writer = PreferenceWriter(tokenizer)
    seen: set[int] = set()
    hh_count = 0
    for record in load_dataset("Anthropic/hh-rlhf", split="train", streaming=True):
        pair = preference_pair_from_messages(
            parse_hh_transcript(record.get("chosen", "")),
            parse_hh_transcript(record.get("rejected", "")),
        )
        if pair:
            fingerprint = stable_hash(json.dumps(pair, ensure_ascii=False, sort_keys=True))
            if fingerprint in seen:
                continue
            seen.add(fingerprint)
            writer.add(*pair)
            hh_count += 1
        if hh_count >= HH_LIMIT:
            break
    uf_count = 0
    for record in load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="train_prefs", streaming=True):
        pair = preference_pair_from_messages(
            normalize_messages(record.get("chosen")), normalize_messages(record.get("rejected")),
        )
        if pair:
            fingerprint = stable_hash(json.dumps(pair, ensure_ascii=False, sort_keys=True))
            if fingerprint in seen:
                continue
            seen.add(fingerprint)
            writer.add(*pair)
            uf_count += 1
        if uf_count >= UF_LIMIT:
            break
    records = writer.close()
    print(f"  HH pairs={hh_count:,}, UltraFeedback pairs={uf_count:,}, total={records:,}")
    return records


def prepare_everything(tokenizer) -> dict[str, int]:
    meta = prepare_pretrain_cache(tokenizer)
    meta["sft_records"] = prepare_dialogue_cache(tokenizer)
    meta["reason_records"] = prepare_reasoning_cache(tokenizer)
    meta["preference_records"] = prepare_preference_cache(tokenizer)
    meta["record_len"] = CFG.seq_len + 1
    atomic_json(meta, META_FILE)
    return meta


# -------------------------------------------------------------------------------------------------
# Faithful CubicV7 + hierarchical late retrieval
# -------------------------------------------------------------------------------------------------


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        # Keep activations in bf16/fp16 under autocast; a float32 norm weight
        # must not accidentally promote the full sequence tensor to float32.
        return normalized.to(x.dtype) * self.weight.to(x.dtype)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_context: int, base: float):
        super().__init__()
        inverse = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
        positions = torch.arange(max_context, dtype=torch.float32)
        frequencies = torch.outer(positions, inverse)
        self.register_buffer("cos", frequencies.cos(), persistent=False)
        self.register_buffer("sin", frequencies.sin(), persistent=False)

    def forward(self, length: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.cos[:length], self.sin[:length]


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    # x: [batch, heads, sequence, head_dim]
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    cos = cos[:x.shape[-2]].to(device=x.device, dtype=x.dtype).view(1, 1, x.shape[-2], -1)
    sin = sin[:x.shape[-2]].to(device=x.device, dtype=x.dtype).view(1, 1, x.shape[-2], -1)
    return torch.stack((x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1).flatten(-2)


class SwiGLU(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.gate = nn.Linear(config.dim, config.mlp_hidden, bias=False)
        self.value = nn.Linear(config.dim, config.mlp_hidden, bias=False)
        self.down = nn.Linear(config.mlp_hidden, config.dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(F.silu(self.gate(x)) * self.value(x))


class DepthCompressor(nn.Module):
    """Shared compressor. The returned history is deliberately NOT detached."""

    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.norm = RMSNorm(config.dim, config.norm_eps)
        self.kv = nn.Linear(config.dim, 2 * config.depth_rank, bias=False)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        batch, sequence, _ = x.shape
        kv = self.kv(self.norm(x)).view(
            batch, sequence, 2, self.config.depth_heads, self.config.depth_head_dim,
        )
        return kv[:, :, 0], kv[:, :, 1]


class CubicHierAttention(nn.Module):
    """Faithful V7 sequence attention and separate gated depth residual.

    Two penultimate layers use the validated cubic_hier approximation: cosine
    score against key-group means, causal top-group routing under no_grad,
    then one fused SDPA over the selected token groups. The last layer is full
    global causal attention so it can integrate all evidence selected below.
    Only the hard routing decision is non-differentiable; surviving Q/K/V
    scores keep their gradients.
    """

    def __init__(self, config: Config, layer_index: int):
        super().__init__()
        self.config = config
        self.layer_index = layer_index
        # Gemma 4's useful invariant transferred without changing Cubic's
        # depth path: retrieval happens before, never instead of, the final
        # global integration layer.
        self.hierarchical = layer_index in (config.depth - 3, config.depth - 2)
        self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False)
        self.proj = nn.Linear(config.dim, config.dim, bias=False)
        # Per-head QK RMSNorm prevents attention-logit growth. It is applied
        # before RoPE and is independent of the faithful depth-memory channel.
        self.q_norm = RMSNorm(config.head_dim, config.norm_eps)
        self.k_norm = RMSNorm(config.head_dim, config.norm_eps)
        if self.hierarchical:
            self.log_temperature = nn.Parameter(torch.full((config.n_heads,), math.log(10.0)))
            causal_future = torch.ones(config.max_context, config.max_context, dtype=torch.bool).triu(1)
            self.register_buffer("causal_future", causal_future, persistent=False)
            n_groups = math.ceil(config.max_context / config.group_chunk)
            group_ids = torch.arange(n_groups)
            query_groups = torch.arange(config.max_context) // config.group_chunk
            self.register_buffer(
                "group_previous", query_groups.unsqueeze(1) > group_ids.unsqueeze(0), persistent=False,
            )
            self.register_buffer(
                "group_current", query_groups.unsqueeze(1) == group_ids.unsqueeze(0), persistent=False,
            )

        self.has_depth = layer_index > 0
        self.has_depth_choice = layer_index > 1
        if self.has_depth:
            self.depth_up = nn.Linear(config.depth_rank, config.dim, bias=False)
            if self.has_depth_choice:
                self.q_depth = nn.Linear(config.dim, config.depth_rank, bias=False)
                self.layer_bias = nn.Parameter(
                    torch.linspace(-config.recency_bias, config.recency_bias, layer_index),
                )
            self.depth_mix_logit = nn.Parameter(torch.full((config.dim,), math.atanh(config.initial_mix)))
            self.depth_content_gate = nn.Linear(config.dim, 1)
            self.depth_norm = RMSNorm(config.dim, config.norm_eps)
            self.agreement_scale = nn.Parameter(torch.zeros(1))
            self.agreement_bias = nn.Parameter(torch.zeros(1))
            nn.init.zeros_(self.depth_content_gate.weight)
            nn.init.zeros_(self.depth_content_gate.bias)

    def sequence_attention(self, x: torch.Tensor, rope: RotaryEmbedding) -> torch.Tensor:
        config = self.config
        batch, sequence, _ = x.shape
        q, k, v = self.qkv(x).split(config.dim, dim=-1)
        q = q.view(batch, sequence, config.n_heads, config.head_dim).transpose(1, 2)
        k = k.view(batch, sequence, config.n_heads, config.head_dim).transpose(1, 2)
        v = v.view(batch, sequence, config.n_heads, config.head_dim).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        cos, sin = rope(sequence)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)

        if not self.hierarchical:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
        else:
            q = F.normalize(q, dim=-1)
            k = F.normalize(k, dim=-1)
            temperature = self.log_temperature.clamp(math.log(2.0), math.log(32.0)).exp()
            q = q * temperature.view(1, config.n_heads, 1, 1)
            future = self.causal_future[:sequence, :sequence].view(1, 1, sequence, sequence)
            with torch.no_grad():
                padding = (-sequence) % config.group_chunk
                grouped_k = F.pad(k, (0, 0, 0, padding)) if padding else k
                n_groups = grouped_k.shape[-2] // config.group_chunk
                grouped_k = grouped_k.view(
                    batch, config.n_heads, n_groups, config.group_chunk, config.head_dim,
                ).mean(dim=3)
                coarse = torch.matmul(q, grouped_k.transpose(-2, -1))
                # Causal routing: summaries are used only for fully completed
                # groups. The current group is forced into the mask but never
                # scored through a mean containing future keys.
                previous = self.group_previous[:sequence, :n_groups].view(1, 1, sequence, n_groups)
                current = self.group_current[:sequence, :n_groups].view(1, 1, sequence, n_groups)
                coarse = coarse.masked_fill(~previous, float("-inf"))
                if config.group_keep > 1:
                    keep_previous = min(config.group_keep - 1, n_groups)
                    threshold = coarse.topk(keep_previous, dim=-1).values[..., -1:]
                    group_keep = ((coarse >= threshold) & previous) | current
                else:
                    group_keep = current
                token_keep = group_keep.repeat_interleave(config.group_chunk, dim=-1)[..., :sequence]
                keep_mask = token_keep & (~future)
            out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=keep_mask, scale=1.0, dropout_p=0.0,
            )
        return out.transpose(1, 2).contiguous().view(batch, sequence, config.dim)

    def depth_residual(
        self, x: torch.Tensor, history_k: Sequence[torch.Tensor], history_v: Sequence[torch.Tensor],
    ) -> torch.Tensor:
        config = self.config
        batch, sequence, _ = x.shape
        if self.has_depth_choice:
            layers = len(history_k)
            q = self.q_depth(x).view(batch * sequence, config.depth_heads, 1, config.depth_head_dim)
            k = torch.stack(tuple(history_k), dim=2).permute(0, 1, 3, 2, 4).reshape(
                batch * sequence, config.depth_heads, layers, config.depth_head_dim,
            )
            v = torch.stack(tuple(history_v), dim=2).permute(0, 1, 3, 2, 4).reshape(
                batch * sequence, config.depth_heads, layers, config.depth_head_dim,
            )
            # The depth axis is tiny (at most 15 states), so a manual softmax
            # is both cheap and more robust than dispatching SDPA here. On
            # PyTorch 2.11/CUDA 13, compiled efficient-attention backward can
            # produce a 65536-vs-65535 shape error for batch=16, seq=4096 when
            # its broadcast bias has zero strides. Sequence attention still
            # uses fused FlashAttention; only this small cross-layer selector
            # avoids the faulty backend. The formula is exactly standard
            # scaled dot-product attention plus the learned recency bias.
            scale = config.depth_head_dim ** -0.5
            scores = torch.matmul(q, k.transpose(-2, -1)) * scale
            scores = scores + self.layer_bias.view(1, 1, 1, layers).to(scores.dtype)
            weights = torch.softmax(scores.float(), dim=-1).to(v.dtype)
            depth = torch.matmul(weights, v)
            depth = depth.reshape(batch, sequence, config.depth_rank)
        else:
            depth = history_v[0].reshape(batch, sequence, config.depth_rank)
        depth = self.depth_up(depth)
        agreement = (
            F.normalize(x, dim=-1) * F.normalize(self.depth_norm(depth), dim=-1)
        ).sum(dim=-1, keepdim=True)
        agreement_gate = 2.0 * torch.sigmoid(self.agreement_scale * agreement + self.agreement_bias)
        content_gate = 2.0 * torch.sigmoid(self.depth_content_gate(x))
        return agreement_gate * content_gate * self.depth_mix_logit.tanh() * depth

    def forward(
        self, x: torch.Tensor, rope: RotaryEmbedding,
        history_k: Sequence[torch.Tensor], history_v: Sequence[torch.Tensor],
    ) -> torch.Tensor:
        # V7 invariant: only sequence attention passes through `proj`.
        out = self.proj(self.sequence_attention(x, rope))
        if self.has_depth:
            out = out + self.depth_residual(x, history_k, history_v)
        return out


class CubicBlock(nn.Module):
    def __init__(self, config: Config, layer_index: int):
        super().__init__()
        self.norm1 = RMSNorm(config.dim, config.norm_eps)
        self.attn = CubicHierAttention(config, layer_index)
        self.norm2 = RMSNorm(config.dim, config.norm_eps)
        self.mlp = SwiGLU(config)
        self.ls1 = nn.Parameter(torch.ones(config.dim))
        self.ls2 = nn.Parameter(torch.ones(config.dim))

    def forward(self, x, rope, history_k, history_v):
        x = x + self.ls1 * self.attn(self.norm1(x), rope, history_k, history_v)
        return x + self.ls2 * self.mlp(self.norm2(x))


class CubicHierLM(nn.Module):
    def __init__(self, config: Config = CFG):
        super().__init__()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.dim)
        self.rope = RotaryEmbedding(config.head_dim, config.max_context, config.rope_base)
        self.blocks = nn.ModuleList([CubicBlock(config, index) for index in range(config.depth)])
        self.depth_memory = DepthCompressor(config)
        self.final_norm = RMSNorm(config.dim, config.norm_eps)
        # One-token-ahead multi-token prediction head. Given h_t and the known
        # token x_(t+1), it predicts x_(t+2) through the tied LM embedding.
        # It is training-only: inference latency and checkpoint compatibility
        # remain straightforward, while each sequence supplies an extra dense
        # learning signal.
        self.mtp_hidden = nn.Linear(config.dim, config.dim, bias=False)
        self.mtp_token = nn.Linear(config.dim, config.dim, bias=False)
        self.mtp_norm = RMSNorm(config.dim, config.norm_eps)
        self.gradient_checkpointing = CHECKPOINT_EVERY > 0
        self.checkpoint_every = CHECKPOINT_EVERY
        self.apply(self._init_weights)
        residual_std = 0.02 / math.sqrt(2.0 * config.depth)
        for block in self.blocks:
            nn.init.normal_(block.attn.proj.weight, mean=0.0, std=residual_std)
            nn.init.normal_(block.mlp.down.weight, mean=0.0, std=residual_std)
            if block.attn.has_depth:
                # Restore the faithful V7 zero-initialized content selector;
                # the global initializer above must not randomize this gate.
                nn.init.zeros_(block.attn.depth_content_gate.weight)
                nn.init.zeros_(block.attn.depth_content_gate.bias)

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        if tokens.shape[1] > self.config.max_context:
            raise ValueError(f"sequence {tokens.shape[1]} exceeds max_context={self.config.max_context}")
        x = self.embed(tokens)
        history_k: list[torch.Tensor] = []
        history_v: list[torch.Tensor] = []
        for index, block in enumerate(self.blocks):
            if index < self.config.depth - 1:
                new_k, new_v = self.depth_memory(x)  # no detach: faithful V7
            should_checkpoint = (
                self.gradient_checkpointing
                and self.training
                and self.checkpoint_every > 0
                and index % self.checkpoint_every == 0
            )
            if should_checkpoint:
                # Non-reentrant checkpointing records tensors captured by this
                # closure, including the entire differentiable depth history.
                x = checkpoint(
                    lambda value, b=block, hk=tuple(history_k), hv=tuple(history_v): b(value, self.rope, hk, hv),
                    x, use_reentrant=False,
                )
            else:
                x = block(x, self.rope, history_k, history_v)
            if index < self.config.depth - 1:
                history_k.append(new_k)
                history_v.append(new_v)
        return self.final_norm(x)

    def logits(self, hidden: torch.Tensor) -> torch.Tensor:
        return F.linear(hidden, self.embed.weight)

    def mtp_conditioned_hidden(self, hidden: torch.Tensor, next_tokens: torch.Tensor) -> torch.Tensor:
        return self.mtp_norm(self.mtp_hidden(hidden) + self.mtp_token(self.embed(next_tokens)))


# -------------------------------------------------------------------------------------------------
# Faithful hybrid Muon + AdamW
# -------------------------------------------------------------------------------------------------


MUON_LR = 0.020
DEPTH_MUON_LR = 0.030
MUON_MOMENTUM = 0.95
MUON_NS_STEPS = env_int("CUBIC_MUON_NS", 4)
MUON_WEIGHT_DECAY = 0.01
MUON_BUCKET_CHUNK = 8
AUX_LR = 3e-4
GATE_LR = 5e-3
AUX_WEIGHT_DECAY = 0.01
MIN_MUON_NUMEL = 4_096


@torch.no_grad()
def zeropower_batched(updates: torch.Tensor, steps: int = MUON_NS_STEPS, eps: float = 1e-7) -> torch.Tensor:
    a, b, c = 3.4445, -4.7750, 2.0315
    x = updates.to(torch.bfloat16)
    x = x / (x.norm(dim=(-2, -1), keepdim=True) + eps)
    for _ in range(steps):
        gram = x @ x.mT
        x = a * x + (b * gram + c * (gram @ gram)) @ x
    return x


class CubicMuon(torch.optim.Optimizer):
    def __init__(self, muon_groups, adam_groups, lr_multiplier: float):
        super().__init__(muon_groups, dict(
            lr=MUON_LR, momentum=MUON_MOMENTUM, ns_steps=MUON_NS_STEPS,
            weight_decay=MUON_WEIGHT_DECAY,
        ))
        self.lr_multiplier = lr_multiplier
        self.muon_base_lrs = [float(group["lr"]) * lr_multiplier for group in self.param_groups]
        self.muon_depth = [bool(group.get("is_depth", False)) for group in self.param_groups]
        self.layouts: list[list[list[tuple[torch.Tensor, bool, float]]]] = []
        for group in self.param_groups:
            buckets: dict[tuple[int, int], list[tuple[torch.Tensor, bool, float]]] = defaultdict(list)
            for parameter in group["params"]:
                rows, columns = parameter.shape
                transpose = rows > columns
                canonical = (min(rows, columns), max(rows, columns))
                scale = math.sqrt(max(1.0, rows / columns))
                buckets[canonical].append((parameter, transpose, scale))
            self.layouts.append(list(buckets.values()))
        for group in adam_groups:
            group["lr"] *= lr_multiplier
        try:
            self.adamw = torch.optim.AdamW(adam_groups, betas=(0.9, 0.95), eps=1e-8, fused=True)
        except (TypeError, RuntimeError):
            self.adamw = torch.optim.AdamW(adam_groups, betas=(0.9, 0.95), eps=1e-8)
        self.adam_base_lrs = [float(group["lr"]) for group in self.adamw.param_groups]
        self.adam_gate = [bool(group.get("is_gate", False)) for group in self.adamw.param_groups]

    def set_lr_scale(self, scale: float, gate_scale: float, depth_boost: float) -> None:
        for group, base_lr, is_depth in zip(self.param_groups, self.muon_base_lrs, self.muon_depth):
            group["lr"] = base_lr * scale * (depth_boost if is_depth else 1.0)
        for group, base_lr, is_gate in zip(self.adamw.param_groups, self.adam_base_lrs, self.adam_gate):
            group["lr"] = base_lr * (gate_scale if is_gate else scale)

    @torch.no_grad()
    def step(self, closure=None):
        for group_index, group in enumerate(self.param_groups):
            lr, beta = group["lr"], group["momentum"]
            active = [parameter for parameter in group["params"] if parameter.grad is not None]
            if group["weight_decay"] and active:
                torch._foreach_mul_(active, 1.0 - lr * group["weight_decay"])
            for bucket in self.layouts[group_index]:
                for start in range(0, len(bucket), MUON_BUCKET_CHUNK):
                    prepared, metadata = [], []
                    for parameter, transpose, scale in bucket[start:start + MUON_BUCKET_CHUNK]:
                        if parameter.grad is None:
                            continue
                        state = self.state[parameter]
                        if "momentum_buffer" not in state:
                            state["momentum_buffer"] = torch.zeros_like(parameter)
                        momentum = state["momentum_buffer"]
                        momentum.lerp_(parameter.grad, 1.0 - beta)
                        # Nesterov update used by the validated V7 optimizer.
                        update = torch.lerp(parameter.grad, momentum, beta)
                        prepared.append(update.mT if transpose else update)
                        metadata.append((parameter, transpose, scale))
                    if prepared:
                        orthogonal = zeropower_batched(torch.stack(prepared), group["ns_steps"])
                        for update, (parameter, transpose, scale) in zip(orthogonal.unbind(0), metadata):
                            parameter.add_(update.mT if transpose else update, alpha=-lr * scale)
        self.adamw.step()

    def zero_grad(self, set_to_none: bool = True) -> None:
        super().zero_grad(set_to_none=set_to_none)
        self.adamw.zero_grad(set_to_none=set_to_none)


def build_optimizer(model: nn.Module, lr_multiplier: float) -> CubicMuon:
    muon, depth, auxiliary_decay, auxiliary_no_decay, gates = [], [], [], [], []
    gate_names = (
        # This matches the validated V7 split exactly. Agreement scalars and
        # cosine temperatures remain conservative AdamW auxiliaries (3e-4),
        # not fast 5e-3 gates.
        "depth_mix_logit", "depth_content_gate", "layer_bias", "ls1", "ls2",
    )
    depth_names = ("q_depth", "depth_up", "depth_memory")
    forbidden_muon = ("embed", "norm")
    for name, parameter in model.named_parameters():
        is_gate = any(fragment in name for fragment in gate_names)
        use_muon = (
            parameter.ndim == 2 and parameter.numel() >= MIN_MUON_NUMEL
            and not any(fragment in name for fragment in forbidden_muon) and not is_gate
        )
        if use_muon:
            (depth if any(fragment in name for fragment in depth_names) else muon).append(parameter)
        elif is_gate:
            gates.append(parameter)
        elif parameter.ndim < 2 or any(fragment in name for fragment in ("embed", "norm")):
            # Embeddings, norm scales and scalar controls are more stable
            # without AdamW decay in small language models.
            auxiliary_no_decay.append(parameter)
        else:
            auxiliary_decay.append(parameter)
    print(
        f"optimizer: Muon={len(muon)} + depth={len(depth)}, "
        f"AdamW decay={len(auxiliary_decay)} no_decay={len(auxiliary_no_decay)} "
        f"+ gates={len(gates)}, phase_lr={lr_multiplier:.3f}x",
    )
    muon_groups = [
        dict(params=muon, lr=MUON_LR, momentum=MUON_MOMENTUM, ns_steps=MUON_NS_STEPS,
             weight_decay=MUON_WEIGHT_DECAY, is_depth=False),
        dict(params=depth, lr=DEPTH_MUON_LR, momentum=MUON_MOMENTUM, ns_steps=MUON_NS_STEPS,
             weight_decay=MUON_WEIGHT_DECAY, is_depth=True),
    ]
    adam_groups = [
        dict(params=auxiliary_decay, lr=AUX_LR, weight_decay=AUX_WEIGHT_DECAY, is_gate=False),
        dict(params=auxiliary_no_decay, lr=AUX_LR, weight_decay=0.0, is_gate=False),
        dict(params=gates, lr=GATE_LR, weight_decay=0.0, is_gate=True),
    ]
    return CubicMuon(
        [group for group in muon_groups if group["params"]],
        [group for group in adam_groups if group["params"]], lr_multiplier,
    )


def cosine_schedule(step: int, total_steps: int, warmup: int, floor: float = 0.10) -> float:
    if step < warmup:
        return (step + 1) / max(1, warmup)
    progress = min(1.0, max(0.0, (step - warmup) / max(1, total_steps - warmup - 1)))
    return floor + (1.0 - floor) * 0.5 * (1.0 + math.cos(math.pi * progress))


def wsd_schedule(
    step: int, total_steps: int, warmup: int, decay_fraction: float = 0.10, floor: float = 0.10,
) -> float:
    """Warmup-stable-decay schedule for long base pretraining.

    Unlike a full-run cosine, it does not starve the middle of a short 3B-token
    run of learning rate. The last 10% remains a clean annealing stage.
    """
    if step < warmup:
        return (step + 1) / max(1, warmup)
    decay_steps = max(1, int(total_steps * decay_fraction))
    decay_start = max(warmup, total_steps - decay_steps)
    if step < decay_start:
        return 1.0
    progress = min(1.0, max(0.0, (step - decay_start) / max(1, total_steps - decay_start - 1)))
    return floor + (1.0 - floor) * 0.5 * (1.0 + math.cos(math.pi * progress))


def gate_schedule(step: int, base_scale: float) -> float:
    opened = min(1.0, (step + 1) / 200.0)
    return opened * (0.25 + 0.75 * base_scale)


def mtp_loss_weight(phase_name: str, step: int, total_steps: int) -> float:
    if phase_name != "pretrain":
        return MTP_SFT_WEIGHT
    decay_start = max(0, total_steps - max(1, int(total_steps * 0.10)))
    return MTP_DECAY_WEIGHT if step >= decay_start else MTP_PRETRAIN_WEIGHT


def depth_boost(step: int) -> float:
    return 1.0 + 0.75 * math.exp(-step / 900.0)


# -------------------------------------------------------------------------------------------------
# Memory-mapped samplers and fused losses
# -------------------------------------------------------------------------------------------------


def map_u16(path: Path) -> torch.Tensor:
    return torch.from_file(str(path), shared=False, size=path.stat().st_size // 2, dtype=torch.uint16)


def map_u8(path: Path) -> torch.Tensor:
    return torch.from_file(str(path), shared=False, size=path.stat().st_size, dtype=torch.uint8)


class PretrainSampler:
    def __init__(self, path: Path, sequence: int, seed: int):
        self.data = map_u16(path)
        self.sequence = sequence
        self.seed = seed
        self.windows = (len(self.data) - 1) // sequence
        if self.windows <= 0:
            raise RuntimeError(f"pretrain cache {path} is too short")
        self.a = coprime_multiplier(self.windows, seed)
        self.b = seed % self.windows

    def batch(self, micro_index: int, batch_size: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        rows = []
        for row in range(batch_size):
            global_window = micro_index * batch_size + row
            window = (self.a * global_window + self.b) % self.windows
            start = window * self.sequence
            rows.append(self.data[start:start + self.sequence + 1].to(torch.long))
        packed_cpu = torch.stack(rows).pin_memory()
        packed = packed_cpu.to(DEVICE, non_blocking=True)
        return packed[:, :-1], packed[:, 1:], torch.ones_like(packed[:, 1:], dtype=torch.bool)


def coprime_multiplier(records: int, seed: int) -> int:
    candidate = (seed | 1) % max(2, records)
    candidate = max(1, candidate)
    while math.gcd(candidate, records) != 1:
        candidate += 2
    return candidate


class RecordSampler:
    def __init__(self, token_path: Path, mask_path: Path, seed: int):
        self.record_len = CFG.seq_len + 1
        self.tokens = map_u16(token_path).view(-1, self.record_len)
        self.masks = map_u8(mask_path).view(-1, self.record_len)
        self.records = self.tokens.shape[0]
        self.a = coprime_multiplier(self.records, seed)
        self.b = seed % self.records

    def indices(self, micro_index: int, batch_size: int) -> list[int]:
        base = micro_index * batch_size
        return [(self.a * (base + row) + self.b) % self.records for row in range(batch_size)]

    def batch(self, micro_index: int, batch_size: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        index = torch.tensor(self.indices(micro_index, batch_size), dtype=torch.long)
        packed_cpu = self.tokens[index].to(torch.long).pin_memory()
        masks_cpu = self.masks[index, 1:].to(torch.bool).pin_memory()
        packed = packed_cpu.to(DEVICE, non_blocking=True)
        masks = masks_cpu.to(DEVICE, non_blocking=True)
        return packed[:, :-1], packed[:, 1:].to(DEVICE, non_blocking=True), masks


class PreferenceSampler:
    def __init__(self, seed: int):
        self.record_len = CFG.seq_len + 1
        self.chosen = map_u16(PREF_FILES["chosen_tokens"]).view(-1, self.record_len)
        self.chosen_mask = map_u8(PREF_FILES["chosen_mask"]).view(-1, self.record_len)
        self.rejected = map_u16(PREF_FILES["rejected_tokens"]).view(-1, self.record_len)
        self.rejected_mask = map_u8(PREF_FILES["rejected_mask"]).view(-1, self.record_len)
        self.records = self.chosen.shape[0]
        self.a = coprime_multiplier(self.records, seed)
        self.b = seed % self.records

    def batch(self, micro_index: int, batch_size: int):
        base = micro_index * batch_size
        indices = [(self.a * (base + row) + self.b) % self.records for row in range(batch_size)]
        index = torch.tensor(indices, dtype=torch.long)
        chosen = self.chosen[index].to(torch.long).pin_memory().to(DEVICE, non_blocking=True)
        rejected = self.rejected[index].to(torch.long).pin_memory().to(DEVICE, non_blocking=True)
        chosen_mask = self.chosen_mask[index, 1:].to(torch.bool).pin_memory().to(DEVICE, non_blocking=True)
        rejected_mask = self.rejected_mask[index, 1:].to(torch.bool).pin_memory().to(DEVICE, non_blocking=True)
        return chosen, chosen_mask, rejected, rejected_mask


class InterleavedSampler:
    """Low-cost replay that reduces forgetting between curriculum stages."""

    def __init__(self, primary, replay, replay_every: int):
        self.primary, self.replay, self.replay_every = primary, replay, replay_every

    def batch(self, micro_index: int, batch_size: int):
        if micro_index % self.replay_every == self.replay_every - 1:
            return self.replay.batch(micro_index, batch_size)
        return self.primary.batch(micro_index, batch_size)


class CudaBatchPrefetcher:
    """Double-buffer deterministic batches on a dedicated CUDA copy stream."""

    def __init__(self, sampler, batch_size: int, first_micro_index: int):
        self.sampler = sampler
        self.batch_size = batch_size
        self.stream = torch.cuda.Stream()
        self.next_batch = None
        self._preload(first_micro_index)

    def _preload(self, micro_index: int) -> None:
        with torch.cuda.stream(self.stream):
            self.next_batch = self.sampler.batch(micro_index, self.batch_size)

    def get(self, next_micro_index: int):
        current_stream = torch.cuda.current_stream()
        current_stream.wait_stream(self.stream)
        batch = self.next_batch
        if batch is None:
            raise RuntimeError("CUDA prefetcher has no prepared batch")
        for tensor in batch:
            tensor.record_stream(current_stream)
        self._preload(next_micro_index)
        return batch


class FusedTokenLoss:
    def __init__(self):
        try:
            from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
            self.liger = LigerFusedLinearCrossEntropyLoss()
            print("loss kernel: Liger fused linear cross-entropy")
        except Exception as error:
            print(f"loss kernel: chunked PyTorch fallback ({type(error).__name__}: {error})")
            self.liger = None

    def __call__(self, model: CubicHierLM, hidden: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        selected_hidden = hidden.reshape(-1, hidden.shape[-1])[mask.reshape(-1)]
        selected_targets = targets.reshape(-1)[mask.reshape(-1)]
        if not selected_targets.numel():
            raise RuntimeError("batch has no supervised tokens")
        if self.liger is not None:
            return self.liger(model.embed.weight, selected_hidden, selected_targets)
        total = torch.zeros((), device=hidden.device, dtype=torch.float32)
        chunk = 2_048
        for start in range(0, selected_targets.numel(), chunk):
            end = min(selected_targets.numel(), start + chunk)
            logits = F.linear(selected_hidden[start:end], model.embed.weight).float()
            total = total + F.cross_entropy(logits, selected_targets[start:end], reduction="sum")
        return total / selected_targets.numel()


def response_logps(model: CubicHierLM, hidden: torch.Tensor, tokens: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Length-normalized response log-probability for each item in a DPO batch."""
    targets = tokens[:, 1:]
    values = []
    for row in range(tokens.shape[0]):
        selected_hidden = hidden[row][mask[row]]
        selected_targets = targets[row][mask[row]]
        if not selected_targets.numel():
            raise RuntimeError("DPO row has no response tokens")
        pieces = []
        for start in range(0, selected_targets.numel(), 512):
            end = min(selected_targets.numel(), start + 512)
            logits = F.linear(selected_hidden[start:end], model.embed.weight).float()
            pieces.append(logits.gather(1, selected_targets[start:end, None]).squeeze(1) - logits.logsumexp(-1))
        values.append(torch.cat(pieces).mean())
    return torch.stack(values)


# -------------------------------------------------------------------------------------------------
# Resume-safe staged training
# -------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class Phase:
    name: str
    steps: int
    micro_batch: int
    accumulation: int
    lr_multiplier: float
    warmup: int


def phase_plan(meta: dict[str, int]) -> list[Phase]:
    effective_pretrain = MICRO_BATCH * GRAD_ACCUM * CFG.seq_len
    pretrain_steps = math.ceil(PRETRAIN_TOKENS / effective_pretrain)
    sft_steps = math.ceil(meta["sft_records"] * SFT_EPOCHS / (MICRO_BATCH * GRAD_ACCUM))
    reason_steps = math.ceil(meta["reason_records"] * REASON_EPOCHS / (MICRO_BATCH * GRAD_ACCUM))
    dpo_steps = math.ceil(meta["preference_records"] * DPO_EPOCHS / (DPO_MICRO_BATCH * DPO_ACCUM))
    limit = env_int("CUBIC_MAX_STEPS", 0)

    def limited(value: int) -> int:
        return min(value, limit) if limit > 0 else value

    return [
        Phase("pretrain", limited(pretrain_steps), MICRO_BATCH, GRAD_ACCUM, 1.00, min(2_000, max(100, pretrain_steps // 50))),
        Phase("dialogue_sft", limited(sft_steps), MICRO_BATCH, GRAD_ACCUM, 0.30, min(500, max(50, sft_steps // 20))),
        Phase("reasoning_sft", limited(reason_steps), MICRO_BATCH, GRAD_ACCUM, 0.20, min(300, max(30, reason_steps // 20))),
        Phase("dpo", limited(dpo_steps), DPO_MICRO_BATCH, DPO_ACCUM, 0.05, min(200, max(20, dpo_steps // 20))),
    ]


def cpu_model_state(model: nn.Module) -> dict[str, torch.Tensor]:
    return {name: value.detach().cpu() for name, value in model.state_dict().items()}


def first_nonfinite_tensor(state: Any, prefix: str = "") -> str | None:
    if isinstance(state, torch.Tensor):
        if (state.is_floating_point() or state.is_complex()) and not bool(
            torch.isfinite(state).all()
        ):
            return prefix or "<tensor>"
        return None
    if isinstance(state, dict):
        items = state.items()
    elif isinstance(state, (list, tuple)):
        items = enumerate(state)
    else:
        return None
    for name, value in items:
        child = f"{prefix}.{name}" if prefix else str(name)
        bad = first_nonfinite_tensor(value, child)
        if bad is not None:
            return bad
    return None


@torch.no_grad()
def first_nonfinite_parameter(model: nn.Module) -> str | None:
    for name, parameter in model.named_parameters():
        if not bool(torch.isfinite(parameter).all()):
            return name
    return None


def atomic_torch_save(payload: dict[str, Any], path: Path) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, temporary)
    os.replace(temporary, path)


class HourlyCheckpointPublisher:
    """Upload immutable hourly resume snapshots without stalling the GPU."""

    def __init__(self) -> None:
        self.enabled = HF_CHECKPOINT_UPLOAD
        self.interval = HF_CHECKPOINT_MINUTES * 60.0
        self.last_started = time.monotonic()
        self.thread: threading.Thread | None = None
        self.lock = threading.Lock()
        self.snapshot_root = CHECKPOINTS / ".hf_hourly_upload"

    @staticmethod
    def _snapshot_copy(source: Path, destination: Path) -> None:
        destination.parent.mkdir(parents=True, exist_ok=True)
        try:
            subprocess.run([
                "cp", "--reflink=auto", "--sparse=always",
                "--preserve=timestamps", "--", str(source), str(destination),
            ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except (FileNotFoundError, subprocess.CalledProcessError):
            shutil.copy2(source, destination)

    def _upload(self, folder: Path, phase_index: int, step: int) -> None:
        try:
            from huggingface_hub import HfApi

            api = HfApi(token=os.environ.get("HF_TOKEN") or None)
            api.create_repo(
                HF_CHECKPOINT_REPO, repo_type="model", private=False, exist_ok=True,
            )
            api.upload_folder(
                repo_id=HF_CHECKPOINT_REPO,
                repo_type="model",
                folder_path=str(folder),
                commit_message=(
                    f"Hourly Cubic resume checkpoint: phase {phase_index}, step {step}"
                ),
            )
            print(
                f"\nHF CHECKPOINT UPLOADED: https://huggingface.co/"
                f"{HF_CHECKPOINT_REPO} phase={phase_index} step={step}",
                flush=True,
            )
        except Exception as exc:
            # Training must never die because the network or Hub is unavailable.
            print(
                f"\nWARNING: hourly Hugging Face checkpoint upload failed: {exc!r}. "
                "Training continues; the next hourly attempt will retry.",
                flush=True,
            )
        finally:
            shutil.rmtree(folder, ignore_errors=True)

    def maybe_publish(
        self, checkpoint_path: Path, phase_index: int, step: int, force: bool = False,
    ) -> bool:
        if not self.enabled or not checkpoint_path.is_file():
            return False
        with self.lock:
            if self.thread is not None and self.thread.is_alive():
                if force:
                    self.thread.join()
                else:
                    print("Hourly HF upload is still active; skipping overlapping upload.")
                    return False
            now = time.monotonic()
            if not force and now - self.last_started < self.interval:
                return False
            self.last_started = now

            stamp = f"phase-{phase_index:02d}-step-{step:08d}-{int(time.time())}"
            folder = self.snapshot_root / stamp
            checkpoint_copy = folder / "checkpoints" / "latest.pt"
            self._snapshot_copy(checkpoint_path, checkpoint_copy)
            self._snapshot_copy(TOKENIZER_PATH, folder / "work" / "tokenizer.json")
            script_path = Path(__file__).resolve()
            if script_path.is_file():
                self._snapshot_copy(script_path, folder / "train_and_continue.py")
            status = {
                "format": ARCH_VERSION,
                "phase_index": phase_index,
                "next_step": step,
                "checkpoint": "checkpoints/latest.pt",
                "uploaded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
            }
            atomic_json(status, folder / "hourly_checkpoint.json")
            self.thread = threading.Thread(
                target=self._upload,
                args=(folder, phase_index, step),
                name="cubic-hf-checkpoint-upload",
                daemon=True,
            )
            self.thread.start()
            print(
                f"Hourly HF checkpoint upload started in background: "
                f"{HF_CHECKPOINT_REPO} phase={phase_index} step={step}",
                flush=True,
            )
            return True

    def finish(self, checkpoint_path: Path, phase_index: int, step: int) -> None:
        if not self.enabled:
            return
        existing = self.thread
        if existing is not None and existing.is_alive():
            print("Waiting for active Hugging Face checkpoint upload ...", flush=True)
            existing.join()
        self.maybe_publish(checkpoint_path, phase_index, step, force=True)
        if self.thread is not None:
            self.thread.join()


HOURLY_CHECKPOINT_PUBLISHER = HourlyCheckpointPublisher()


def save_checkpoint(
    model: CubicHierLM,
    optimizer: CubicMuon | None,
    phase_index: int,
    step: int,
    metrics: dict[str, Any],
    path: Path = LATEST,
) -> None:
    bad_parameter = first_nonfinite_parameter(model)
    if bad_parameter is not None:
        # Never replace a resumable checkpoint with poisoned weights.
        raise FloatingPointError(
            f"Refusing to save non-finite checkpoint; bad parameter: {bad_parameter}. "
            f"The previous checkpoint at {path} remains intact."
        )
    payload: dict[str, Any] = {
        "format": ARCH_VERSION,
        "config": asdict(CFG),
        "model": cpu_model_state(model),
        "phase_index": phase_index,
        "step": step,
        "metrics": metrics,
        "torch_rng": torch.get_rng_state(),
        "cuda_rng": torch.cuda.get_rng_state_all(),
        "saved_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }
    if optimizer is not None:
        payload["muon_optimizer"] = optimizer.state_dict()
        payload["adamw_optimizer"] = optimizer.adamw.state_dict()
    bad_state = first_nonfinite_tensor(payload)
    if bad_state is not None:
        raise FloatingPointError(
            f"Refusing to save non-finite training state: {bad_state}. "
            f"The previous checkpoint at {path} remains intact."
        )
    atomic_torch_save(payload, path)
    print(f"checkpoint={path} phase_index={phase_index} next_step={step}")
    if path.resolve() == LATEST.resolve():
        HOURLY_CHECKPOINT_PUBLISHER.maybe_publish(path, phase_index, step)


def fetch_hub_resume() -> Path | None:
    """Restore latest.pt and tokenizer.json from the hourly resume repository.

    A destroyed container loses /marimo but not the hourly upload, so this is
    the only path back to an eight-hour run.  The copy is written next to the
    local checkpoints under a distinct name and never overwrites latest.pt.
    """
    if not HF_RESUME_DOWNLOAD:
        return None
    if os.environ.get("CUBIC_RESUME", "auto").strip().lower() in {"none", "off", "0", "false"}:
        return None
    try:
        from huggingface_hub import snapshot_download

        snapshot = Path(snapshot_download(
            repo_id=HF_CHECKPOINT_REPO,
            repo_type="model",
            token=os.environ.get("HF_TOKEN") or None,
            allow_patterns=[
                "checkpoints/latest.pt", "work/tokenizer.json", "hourly_checkpoint.json",
            ],
        ))
    except Exception as exc:
        print(f"Hub resume unavailable ({HF_CHECKPOINT_REPO}): {exc!r}", flush=True)
        return None

    status: dict[str, Any] = {}
    status_path = snapshot / "hourly_checkpoint.json"
    if status_path.is_file():
        try:
            status = json.loads(status_path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            status = {}
    if status.get("format") not in (None, ARCH_VERSION):
        print(
            f"Hub resume ignored: format={status.get('format')!r}, expected {ARCH_VERSION!r}",
            flush=True,
        )
        return None

    remote_tokenizer = snapshot / "work" / "tokenizer.json"
    if remote_tokenizer.is_file() and not TOKENIZER_PATH.exists():
        TOKENIZER_PATH.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(remote_tokenizer, TOKENIZER_PATH)
        print(f"Restored tokenizer from {HF_CHECKPOINT_REPO}", flush=True)

    remote_checkpoint = snapshot / "checkpoints" / "latest.pt"
    if not remote_checkpoint.is_file():
        print(f"No checkpoints/latest.pt in {HF_CHECKPOINT_REPO}", flush=True)
        return None
    CHECKPOINTS.mkdir(parents=True, exist_ok=True)
    local = CHECKPOINTS / "hub_latest.pt"
    if not local.is_file() or local.stat().st_size != remote_checkpoint.stat().st_size:
        shutil.copy2(remote_checkpoint, local)
    print(
        f"Hub resume ready: {local} "
        f"(next_step={status.get('next_step', '?')}, uploaded_at={status.get('uploaded_at', '?')})",
        flush=True,
    )
    return local


def find_resume_checkpoint() -> Path | None:
    """Resolve an explicit checkpoint or find the newest resumable checkpoint.

    CUBIC_RESUME may point to a .pt file or to a directory.  Automatic mode
    searches the active root first and then the common /marimo locations, so
    changing a launch directory or CUBIC_ROOT cannot silently restart training.
    Set CUBIC_RESUME=none to deliberately start from scratch.
    """
    requested = os.environ.get("CUBIC_RESUME", "auto").strip()
    if requested.lower() in {"none", "off", "0", "false"}:
        print("Checkpoint resume explicitly disabled (CUBIC_RESUME=none).")
        return None

    if requested.lower() != "auto":
        path = Path(requested).expanduser()
        if path.is_dir():
            candidates = list(path.glob("*.pt")) + list(path.glob("checkpoints/*.pt"))
            if not candidates:
                raise FileNotFoundError(f"No .pt checkpoints under CUBIC_RESUME={path}")
            return max(candidates, key=lambda item: item.stat().st_mtime)
        if not path.is_file():
            raise FileNotFoundError(f"CUBIC_RESUME checkpoint does not exist: {path}")
        return path

    candidates: list[Path] = []
    search_dirs = [CHECKPOINTS, ROOT, Path("/marimo/checkpoints"), Path("/marimo")]
    seen: set[Path] = set()
    for folder in search_dirs:
        if not folder.exists():
            continue
        # Avoid a costly recursive walk through datasets: only known checkpoint
        # directories are recursive; /marimo itself is searched one level deep.
        patterns = ("*.pt", "checkpoints/*.pt") if folder == Path("/marimo") else ("*.pt",)
        for pattern in patterns:
            for path in folder.glob(pattern):
                resolved = path.resolve()
                if resolved not in seen and not path.name.endswith(".tmp"):
                    seen.add(resolved)
                    candidates.append(path)
    if not candidates:
        hub = fetch_hub_resume()
        if hub is not None:
            return hub
        print(f"No checkpoint found; starting from step 0. Expected: {LATEST}")
        return None
    # Prefer latest.pt when present in the active root; otherwise newest file.
    return LATEST if LATEST.is_file() else max(candidates, key=lambda item: item.stat().st_mtime)


def validated_checkpoint(path: Path) -> dict[str, Any]:
    payload = torch.load(path, map_location="cpu", weights_only=False)
    if payload.get("format") != ARCH_VERSION:
        raise RuntimeError(f"format={payload.get('format')!r}, expected={ARCH_VERSION!r}")
    bad_tensor = first_nonfinite_tensor(payload.get("model", {}))
    if bad_tensor is not None:
        raise FloatingPointError(f"non-finite model tensor {bad_tensor}")
    for optimizer_key in ("muon_optimizer", "adamw_optimizer"):
        if optimizer_key in payload:
            bad_tensor = first_nonfinite_tensor(payload[optimizer_key], optimizer_key)
            if bad_tensor is not None:
                raise FloatingPointError(f"non-finite optimizer tensor {bad_tensor}")
    return payload


def load_checkpoint(model: CubicHierLM) -> tuple[int, int, dict[str, Any], dict[str, Any] | None]:
    checkpoint_path = find_resume_checkpoint()
    if checkpoint_path is None:
        return 0, 0, {}, None
    candidates = [checkpoint_path]
    if BEST.is_file() and BEST.resolve() != checkpoint_path.resolve():
        candidates.append(BEST)
    payload = None
    loaded_path = None
    failures = []
    hub_tried = False
    position = 0
    while position < len(candidates):
        candidate = candidates[position]
        position += 1
        print(f"Checking resume checkpoint {candidate} ...", flush=True)
        try:
            payload, loaded_path = validated_checkpoint(candidate), candidate
            break
        except Exception as exc:
            failures.append(f"{candidate}: {exc!r}")
            print(f"  rejected checkpoint: {exc!r}", flush=True)
        if position >= len(candidates) and not hub_tried:
            # Only reach for the network once every local copy has failed.
            hub_tried = True
            hub_candidate = fetch_hub_resume()
            known = {item.resolve() for item in candidates}
            if hub_candidate is not None and hub_candidate.resolve() not in known:
                print("Local checkpoints all failed; trying the Hub copy ...", flush=True)
                candidates.append(hub_candidate)
    if payload is None or loaded_path is None:
        message = (
            "No finite compatible checkpoint is available:\n  "
            + "\n  ".join(failures)
        )
        if not ALLOW_COLD_START:
            raise RuntimeError(
                message
                + "\nRefusing to discard existing training progress by silently restarting "
                  "from step 0. Inspect the files above, or set CUBIC_ALLOW_COLD_START=1 "
                  "to deliberately begin a fresh run."
            )
        print(message, flush=True)
        print("CUBIC_ALLOW_COLD_START=1: starting from step 0.", flush=True)
        return 0, 0, {}, None
    print(f"Resuming from verified checkpoint {loaded_path} ...", flush=True)
    saved_config = payload.get("config", {})
    structural_fields = ("vocab_size", "dim", "depth", "n_heads", "mlp_hidden", "depth_rank", "depth_heads")
    mismatches = {
        key: (saved_config.get(key), getattr(CFG, key))
        for key in structural_fields
        if key in saved_config and saved_config[key] != getattr(CFG, key)
    }
    if mismatches:
        raise RuntimeError(f"checkpoint architecture is incompatible: {mismatches}")
    model.load_state_dict(payload["model"], strict=True)
    if "torch_rng" in payload:
        torch.set_rng_state(payload["torch_rng"])
    if "cuda_rng" in payload and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(payload["cuda_rng"])
    phase_index, step = int(payload["phase_index"]), int(payload["step"])
    print(
        f"Checkpoint loaded successfully: phase_index={phase_index}, next_step={step}, "
        f"saved_at={payload.get('saved_at', 'unknown')}", flush=True,
    )
    return phase_index, step, payload.get("metrics", {}), payload


def compile_model(model: CubicHierLM):
    if not COMPILE:
        return model
    valid_modes = {
        "default", "reduce-overhead", "max-autotune",
        "max-autotune-no-cudagraphs",
    }
    if COMPILE_MODE not in valid_modes:
        raise ValueError(
            f"Unsupported CUBIC_COMPILE_MODE={COMPILE_MODE!r}; "
            f"choose one of {sorted(valid_modes)}"
        )
    print(
        f"Compiling faithful CubicV7-Hier + QKNorm/MTP graph "
        f"(mode={COMPILE_MODE}; first step will be slow) ..."
    )
    return torch.compile(model, mode=COMPILE_MODE, dynamic=False)


@torch.no_grad()
def evaluate_lm(
    run_model, model: CubicHierLM, sampler, loss_fn: FusedTokenLoss,
    batches: int = EVAL_BATCHES, batch_size: int = EVAL_BATCH_SIZE,
) -> float:
    model.eval()
    previous_checkpointing = model.gradient_checkpointing
    model.gradient_checkpointing = False
    losses = []
    for index in range(batches):
        tokens, targets, mask = sampler.batch(9_000_000 + index, batch_size)
        with amp_context():
            hidden = run_model(tokens)
            loss = loss_fn(model, hidden, targets, mask)
        losses.append(float(loss))
    model.gradient_checkpointing = previous_checkpointing
    model.train()
    return sum(losses) / len(losses)


def train_language_phase(
    phase_index: int,
    phase: Phase,
    model: CubicHierLM,
    run_model,
    sampler,
    loss_fn: FusedTokenLoss,
    start_step: int,
    resume_payload: dict[str, Any] | None,
    metrics: dict[str, Any],
    val_sampler,
) -> dict[str, Any]:
    optimizer = build_optimizer(model, phase.lr_multiplier)
    if resume_payload and start_step and "muon_optimizer" in resume_payload:
        optimizer.load_state_dict(resume_payload["muon_optimizer"])
        optimizer.adamw.load_state_dict(resume_payload["adamw_optimizer"])
        print(f"restored optimizer at {phase.name} step {start_step}")
    model.train()
    started = time.perf_counter()
    last_checkpoint = time.monotonic()
    running_loss = running_lm_loss = running_mtp_loss = 0.0
    tokens_seen = 0
    first_micro_index = start_step * phase.accumulation
    prefetcher = CudaBatchPrefetcher(sampler, phase.micro_batch, first_micro_index)
    lr_backoff = 1.0
    consecutive_nonfinite = 0
    for step in range(start_step, phase.steps):
        optimizer.zero_grad(set_to_none=True)
        step_loss = step_lm_loss = step_mtp_loss = 0.0
        invalid_reason = None
        mtp_weight = mtp_loss_weight(phase.name, step, phase.steps)
        for accumulation_index in range(phase.accumulation):
            micro_index = step * phase.accumulation + accumulation_index
            tokens, targets, mask = prefetcher.get(micro_index + 1)
            with amp_context():
                hidden = run_model(tokens)
                lm_loss = loss_fn(model, hidden, targets, mask)
                mtp_mask = mask[:, 1:]
                mtp_hidden = model.mtp_conditioned_hidden(hidden[:, :-1], tokens[:, 1:])
                mtp_loss = loss_fn(model, mtp_hidden, targets[:, 1:], mtp_mask)
                loss = lm_loss + mtp_weight * mtp_loss
                scaled_loss = loss / phase.accumulation
            loss_value = float(loss.detach())
            lm_value = float(lm_loss.detach())
            mtp_value = float(mtp_loss.detach())
            if not (
                math.isfinite(loss_value)
                and math.isfinite(lm_value)
                and math.isfinite(mtp_value)
            ):
                invalid_reason = (
                    f"non-finite loss: total={loss_value}, "
                    f"lm={lm_value}, mtp={mtp_value}"
                )
                break
            scaled_loss.backward()
            step_loss += loss_value / phase.accumulation
            step_lm_loss += lm_value / phase.accumulation
            step_mtp_loss += mtp_value / phase.accumulation
            tokens_seen += int(mask.numel() if phase.name == "pretrain" else mask.sum())
        if invalid_reason is None:
            total_grad_norm = torch.nn.utils.clip_grad_norm_(
                model.parameters(), GRAD_CLIP, foreach=True,
            )
            grad_norm_value = float(total_grad_norm.detach())
            if not math.isfinite(grad_norm_value):
                invalid_reason = f"non-finite gradient norm: {grad_norm_value}"

        if invalid_reason is not None:
            optimizer.zero_grad(set_to_none=True)
            consecutive_nonfinite += 1
            lr_backoff = max(0.125, lr_backoff * 0.5)
            # Re-align double buffering if an accumulation cycle was aborted.
            prefetcher = CudaBatchPrefetcher(
                sampler, phase.micro_batch, (step + 1) * phase.accumulation,
            )
            print(
                f"\nSTABILITY GUARD: skipped {phase.name} step {step + 1}: "
                f"{invalid_reason}; lr_backoff={lr_backoff:.3f}",
                flush=True,
            )
            if consecutive_nonfinite >= 8:
                raise FloatingPointError(
                    "Eight consecutive non-finite steps were blocked. "
                    "The last finite checkpoint was preserved; stopping safely."
                )
            continue

        consecutive_nonfinite = 0
        scale = (
            wsd_schedule(step, phase.steps, phase.warmup)
            if phase.name == "pretrain"
            else cosine_schedule(step, phase.steps, phase.warmup)
        )
        optimizer.set_lr_scale(
            scale * lr_backoff,
            gate_schedule(step, scale) * lr_backoff,
            depth_boost(step),
        )
        optimizer.step()
        lr_backoff = min(1.0, lr_backoff * 1.002)
        running_loss += step_loss
        running_lm_loss += step_lm_loss
        running_mtp_loss += step_mtp_loss
        completed = step + 1

        if completed == 1 or completed % LOG_EVERY == 0:
            elapsed = time.perf_counter() - started
            average = running_loss / (1 if completed == 1 else LOG_EVERY)
            print(
                f"[{phase.name}] {completed:7d}/{phase.steps} loss={average:.4f} "
                f"lm={running_lm_loss / (1 if completed == 1 else LOG_EVERY):.4f} "
                f"mtp={running_mtp_loss / (1 if completed == 1 else LOG_EVERY):.4f}@{mtp_weight:.2f} "
                f"lr={scale * lr_backoff:.4f} grad={grad_norm_value:.3f} "
                f"tok/s={tokens_seen / max(elapsed, 1e-6):,.0f} elapsed={elapsed / 60:.1f}m",
                flush=True,
            )
            running_loss = running_lm_loss = running_mtp_loss = 0.0

        if completed % EVAL_EVERY == 0 or completed == phase.steps:
            validation_loss = evaluate_lm(run_model, model, val_sampler, loss_fn)
            metrics[f"{phase.name}_step_{completed}"] = {
                "validation_loss": validation_loss,
                "perplexity": math.exp(min(20.0, validation_loss)),
            }
            print(f"  validation loss={validation_loss:.4f} ppl={math.exp(min(20, validation_loss)):.2f}")
            best_loss = float(metrics.get("best_validation_loss", float("inf")))
            if validation_loss < best_loss:
                metrics["best_validation_loss"] = validation_loss
                # Best is a fully resumable last-known-good fallback, including
                # both Muon and AdamW states—not model weights alone.
                save_checkpoint(model, optimizer, phase_index, completed, metrics, BEST)

        if time.monotonic() - last_checkpoint >= CHECKPOINT_MINUTES * 60:
            save_checkpoint(model, optimizer, phase_index, completed, metrics)
            last_checkpoint = time.monotonic()

    save_checkpoint(model, None, phase_index + 1, 0, metrics)
    del optimizer
    gc.collect()
    torch.cuda.empty_cache()
    return metrics


def train_dpo_phase(
    phase_index: int,
    phase: Phase,
    model: CubicHierLM,
    run_model,
    sampler: PreferenceSampler,
    start_step: int,
    resume_payload: dict[str, Any] | None,
    metrics: dict[str, Any],
) -> dict[str, Any]:
    print("Loading/freezing the immutable post-reasoning DPO reference ...")
    reference = copy.deepcopy(model).to(DEVICE).eval().requires_grad_(False)
    if start_step > 0:
        if not DPO_REFERENCE.exists():
            raise RuntimeError("Cannot resume DPO: immutable reference checkpoint is missing")
        reference_payload = torch.load(DPO_REFERENCE, map_location="cpu", weights_only=False)
        reference.load_state_dict(reference_payload["model"], strict=True)
    else:
        atomic_torch_save({
            "format": ARCH_VERSION + "-dpo-reference",
            "config": asdict(CFG),
            "model": cpu_model_state(reference),
        }, DPO_REFERENCE)
    reference.gradient_checkpointing = False
    optimizer = build_optimizer(model, phase.lr_multiplier)
    if resume_payload and start_step and "muon_optimizer" in resume_payload:
        optimizer.load_state_dict(resume_payload["muon_optimizer"])
        optimizer.adamw.load_state_dict(resume_payload["adamw_optimizer"])
    model.train()
    started = time.perf_counter()
    last_checkpoint = time.monotonic()
    running_loss = running_margin = 0.0
    first_micro_index = start_step * phase.accumulation
    prefetcher = CudaBatchPrefetcher(sampler, phase.micro_batch, first_micro_index)
    lr_backoff = 1.0
    consecutive_nonfinite = 0
    for step in range(start_step, phase.steps):
        optimizer.zero_grad(set_to_none=True)
        step_loss = step_margin = 0.0
        invalid_reason = None
        for accumulation_index in range(phase.accumulation):
            micro_index = step * phase.accumulation + accumulation_index
            chosen, chosen_mask, rejected, rejected_mask = prefetcher.get(micro_index + 1)
            combined = torch.cat((chosen, rejected), dim=0)
            combined_mask = torch.cat((chosen_mask, rejected_mask), dim=0)
            inputs = combined[:, :-1]
            with amp_context():
                policy_hidden = run_model(inputs)
                policy_logps = response_logps(model, policy_hidden, combined, combined_mask)
                with torch.no_grad():
                    reference_hidden = reference(inputs)
                    reference_logps = response_logps(reference, reference_hidden, combined, combined_mask)
                batch = chosen.shape[0]
                policy_margin = policy_logps[:batch] - policy_logps[batch:]
                reference_margin = reference_logps[:batch] - reference_logps[batch:]
                preference_logits = DPO_BETA * (policy_margin - reference_margin)
                dpo_loss = -F.logsigmoid(preference_logits).mean()
                # Small chosen-response NLL anchor prevents terse preference
                # optimization from erasing the SFT model's language quality.
                sft_anchor = -policy_logps[:batch].mean()
                loss = dpo_loss + DPO_SFT_WEIGHT * sft_anchor
            loss_value = float(loss.detach())
            margin_value = float(policy_margin.detach().mean())
            if not math.isfinite(loss_value) or not math.isfinite(margin_value):
                invalid_reason = (
                    f"non-finite DPO values: loss={loss_value}, "
                    f"margin={margin_value}"
                )
                break
            (loss / phase.accumulation).backward()
            step_loss += loss_value / phase.accumulation
            step_margin += margin_value / phase.accumulation
        if invalid_reason is None:
            total_grad_norm = torch.nn.utils.clip_grad_norm_(
                model.parameters(), GRAD_CLIP, foreach=True,
            )
            grad_norm_value = float(total_grad_norm.detach())
            if not math.isfinite(grad_norm_value):
                invalid_reason = f"non-finite DPO gradient norm: {grad_norm_value}"
        if invalid_reason is not None:
            optimizer.zero_grad(set_to_none=True)
            consecutive_nonfinite += 1
            lr_backoff = max(0.125, lr_backoff * 0.5)
            prefetcher = CudaBatchPrefetcher(
                sampler, phase.micro_batch, (step + 1) * phase.accumulation,
            )
            print(
                f"\nSTABILITY GUARD: skipped DPO step {step + 1}: "
                f"{invalid_reason}; lr_backoff={lr_backoff:.3f}",
                flush=True,
            )
            if consecutive_nonfinite >= 8:
                raise FloatingPointError(
                    "Eight consecutive non-finite DPO steps were blocked; "
                    "the last finite checkpoint remains intact."
                )
            continue
        consecutive_nonfinite = 0
        scale = cosine_schedule(step, phase.steps, phase.warmup)
        optimizer.set_lr_scale(
            scale * lr_backoff, gate_schedule(step, scale) * lr_backoff, 1.0,
        )
        optimizer.step()
        lr_backoff = min(1.0, lr_backoff * 1.002)
        running_loss += step_loss
        running_margin += step_margin
        completed = step + 1
        if completed == 1 or completed % LOG_EVERY == 0:
            divisor = 1 if completed == 1 else LOG_EVERY
            print(
                f"[dpo] {completed:7d}/{phase.steps} loss={running_loss / divisor:.4f} "
                f"chosen_margin={running_margin / divisor:+.4f} elapsed={(time.perf_counter()-started)/60:.1f}m",
                flush=True,
            )
            running_loss = running_margin = 0.0
        if time.monotonic() - last_checkpoint >= CHECKPOINT_MINUTES * 60:
            metrics[f"dpo_step_{completed}"] = {"loss": step_loss, "chosen_margin": step_margin}
            save_checkpoint(model, optimizer, phase_index, completed, metrics)
            last_checkpoint = time.monotonic()
    metrics["dpo_final"] = {"loss": step_loss, "chosen_margin": step_margin}
    save_checkpoint(model, None, phase_index + 1, 0, metrics)
    del optimizer, reference
    gc.collect()
    torch.cuda.empty_cache()
    return metrics


# -------------------------------------------------------------------------------------------------
# Generation, export and Hugging Face publishing
# -------------------------------------------------------------------------------------------------


def sample_token(logits: torch.Tensor, temperature: float, top_p: float, top_k: int, generated: Sequence[int]) -> int:
    logits = logits.float().clone()
    if generated:
        counts: dict[int, int] = defaultdict(int)
        for token in generated[-256:]:
            counts[int(token)] += 1
        for token, count in counts.items():
            logits[token] /= 1.08 ** min(4, count)
    logits /= max(temperature, 1e-4)
    if top_k > 0:
        threshold = logits.topk(min(top_k, logits.numel())).values[-1]
        logits[logits < threshold] = float("-inf")
    sorted_logits, sorted_indices = logits.sort(descending=True)
    probabilities = sorted_logits.softmax(-1)
    cumulative = probabilities.cumsum(-1)
    remove = cumulative - probabilities > top_p
    sorted_logits[remove] = float("-inf")
    selected = torch.multinomial(sorted_logits.softmax(-1), 1)
    return int(sorted_indices[selected])


@torch.no_grad()
def generate(
    model: CubicHierLM,
    tokenizer,
    prompt: str,
    max_new_tokens: int = 512,
    temperature: float = 0.60,
    top_p: float = 0.95,
    top_k: int = 50,
    reasoning: bool = False,
) -> str:
    model.eval()
    previous_checkpointing = model.gradient_checkpointing
    model.gradient_checkpointing = False
    prefix, _ = encode_chat(tokenizer, [
        {"role": "system", "content": REASONING_SYSTEM if reasoning else DIRECT_SYSTEM},
        {"role": "user", "content": prompt},
    ], supervise_assistant=False)
    prefix.append(tokenizer.token_to_id("<|assistant|>"))
    generated: list[int] = []
    eos = tokenizer.token_to_id("<eos>")
    for _ in range(max_new_tokens):
        context = (prefix + generated)[-CFG.max_context:]
        tokens = torch.tensor([context], dtype=torch.long, device=DEVICE)
        with amp_context():
            hidden = model(tokens)
            logits = model.logits(hidden[:, -1:])[0, -1]
        token = sample_token(logits, temperature, top_p, top_k, generated)
        if token == eos:
            break
        generated.append(token)
    model.gradient_checkpointing = previous_checkpointing
    model.train()
    return tokenizer.decode(generated, skip_special_tokens=False).strip()


def record_samples(model: CubicHierLM, tokenizer, phase: str, max_new_tokens: int = 256) -> None:
    prompts = [
        ("Explain in simple terms why the sky is blue, and mention one common misconception.", False),
        ("Напиши короткую, но полезную инструкцию: как проверить Python-функцию на граничных случаях?", False),
        ("A farmer has 17 sheep. All but 9 run away. Think carefully and give the answer.", True),
    ]
    with REPORT.open("a", encoding="utf-8") as handle:
        for prompt, reasoning in prompts:
            answer = generate(
                model, tokenizer, prompt, max_new_tokens=max_new_tokens, reasoning=reasoning,
            )
            item = {
                "phase": phase, "prompt": prompt, "reasoning": reasoning,
                "answer": answer, "time": time.time(),
            }
            handle.write(json.dumps(item, ensure_ascii=False) + "\n")
            print(f"\n--- {phase} sample ---\nUSER: {prompt}\nASSISTANT: {answer}\n")


MODEL_CARD = """---
license: apache-2.0
language:
- en
- ru
pipeline_tag: text-generation
tags:
- custom-code
- cubic-attention
- muon
- dpo
---

# Cubic Hier 150M

An experimental bilingual dialogue/reasoning language model with approximately
157M parameters. It keeps the faithful CubicV7 differentiable depth-memory path,
uses hierarchical cosine retrieval in the two penultimate attention layers, and
finishes with full global causal attention. Every sequence head uses QK RMSNorm.

## Training recipe

1. Base pretraining on a 72/13/10/5 stream of FineWeb-Edu, Russian FineWeb-2,
   deduplicated CodeParrot Python and FineMath-4+.
2. Assistant-only SFT on Smol-SmolTalk and English/Russian Aya examples.
3. Assistant-only reasoning SFT on the shortest verified-correct traces from
   OpenR1-Math-220k.
4. DPO with an SFT anchor on Anthropic HH-RLHF and corrected UltraFeedback.

The optimizer is the project's hybrid orthogonalized Muon (large matrices and
a boosted depth group) plus AdamW for embeddings, norms, scalars and gates.
Base training uses warmup-stable-decay and a one-token-ahead MTP auxiliary loss;
embeddings, RMSNorm scales and controls are excluded from AdamW decay.
SFT teaches explicit direct and `<think>...</think>` system-prompt modes; set
`CUBIC_REASONING=1` in chat mode to request the reasoning format.

## Important

This is a custom PyTorch architecture, not a drop-in Transformers model. Run
`train_and_chat.py` with `CUBIC_MODE=chat`; it downloads/loads all required
files and starts an interactive console. The model is small and experimental:
verify factual, safety-critical and mathematical answers independently.

Each source dataset retains its own license/terms. FineWeb corpora inherit the
Common Crawl terms described on their dataset cards; SmolTalk, Aya and OpenR1
are Apache-2.0; HH-RLHF and UltraFeedback use their published dataset licenses.
"""


def export_model(model: CubicHierLM, tokenizer, metrics: dict[str, Any], phases: Sequence[Phase]) -> Path:
    from safetensors.torch import save_file

    if EXPORT.exists():
        shutil.rmtree(EXPORT)
    EXPORT.mkdir(parents=True)
    state = {name: value.detach().to(torch.bfloat16).cpu().contiguous() for name, value in model.state_dict().items()}
    save_file(state, str(EXPORT / "model.safetensors"), metadata={
        "format": "pt", "architecture": "CubicHierLM", "version": ARCH_VERSION,
    })
    tokenizer.save(str(EXPORT / "tokenizer.json"))
    config = asdict(CFG) | {
        "architectures": ["CubicHierLM"],
        "model_type": "cubic_hier_v7_gemma4",
        "torch_dtype": "bfloat16",
        "tie_word_embeddings": True,
        "parameter_count": sum(parameter.numel() for parameter in model.parameters()),
    }
    atomic_json(config, EXPORT / "config.json")
    atomic_json({
        "temperature": 0.60, "top_p": 0.95, "top_k": 50,
        "repetition_penalty": 1.08, "max_new_tokens": 512,
        "bos_token": "<bos>", "eos_token": "<eos>", "pad_token": "<pad>",
    }, EXPORT / "generation_config.json")
    manifest = {
        "architecture_version": ARCH_VERSION,
        "architecture": "faithful CubicV7 + QKNorm + penultimate hierarchical retrieval + final global attention",
        "config": config,
        "pretraining_mix": {"web_en": 0.72, "web_ru": 0.13, "code": 0.10, "math": 0.05},
        "optimization": {
            "base_schedule": "warmup-stable-cosine-decay-last-10-percent",
            "mtp_pretrain_weight": MTP_PRETRAIN_WEIGHT,
            "mtp_decay_weight": MTP_DECAY_WEIGHT,
            "mtp_sft_weight": MTP_SFT_WEIGHT,
            "adamw_no_decay": ["embeddings", "RMSNorm", "biases", "scalar controls"],
        },
        "phases": [asdict(phase) for phase in phases],
        "metrics": metrics,
        "datasets": [
            "HuggingFaceFW/fineweb-edu:sample-10BT", "HuggingFaceFW/fineweb-2:rus_Cyrl",
            "codeparrot/codeparrot-clean", "HuggingFaceTB/finemath:finemath-4plus",
            "HuggingFaceTB/smol-smoltalk", "CohereLabs/aya_dataset",
            "open-r1/OpenR1-Math-220k:default", "Anthropic/hh-rlhf",
            "HuggingFaceH4/ultrafeedback_binarized:train_prefs",
        ],
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }
    atomic_json(manifest, EXPORT / "training_manifest.json")
    (EXPORT / "README.md").write_text(MODEL_CARD, encoding="utf-8")
    shutil.copy2(Path(__file__).resolve(), EXPORT / "train_and_chat.py")
    if REPORT.exists():
        shutil.copy2(REPORT, EXPORT / "training_samples.jsonl")
    if os.environ.get("CUBIC_UPLOAD_RESUME", "1") == "1" and FINAL.exists():
        shutil.copy2(FINAL, EXPORT / "final_training_checkpoint.pt")
    print(f"Export ready: {EXPORT}")
    return EXPORT


def upload_export(folder: Path) -> None:
    from huggingface_hub import HfApi

    api = HfApi(token=os.environ.get("HF_TOKEN") or None)
    api.create_repo(HF_REPO, repo_type="model", private=False, exist_ok=True)
    print(f"Uploading complete export to https://huggingface.co/{HF_REPO} ...")
    api.upload_folder(
        repo_id=HF_REPO,
        repo_type="model",
        folder_path=str(folder),
        commit_message="Publish faithful CubicV7-Hier 157M with QKNorm, MTP, SFT, reasoning and DPO",
    )
    print(f"Published: https://huggingface.co/{HF_REPO}")


def load_export_for_chat() -> tuple[CubicHierLM, Any]:
    from huggingface_hub import snapshot_download
    from safetensors.torch import load_file
    from tokenizers import Tokenizer

    local = Path(os.environ.get("CUBIC_MODEL_DIR", str(EXPORT)))
    if not (local / "model.safetensors").exists():
        local = Path(snapshot_download(
            repo_id=HF_REPO,
            token=os.environ.get("HF_TOKEN") or None,
            allow_patterns=["model.safetensors", "tokenizer.json", "config.json"],
        ))
    config_data = json.loads((local / "config.json").read_text(encoding="utf-8"))
    structural = {field: config_data[field] for field in asdict(CFG) if field in config_data}
    config = Config(**structural)
    model = CubicHierLM(config).to(DEVICE, dtype=AMP_DTYPE)
    model.load_state_dict(load_file(str(local / "model.safetensors"), device="cpu"), strict=True)
    model.gradient_checkpointing = False
    model.eval()
    tokenizer = Tokenizer.from_file(str(local / "tokenizer.json"))
    print(f"Loaded {sum(p.numel() for p in model.parameters()):,} parameters from {local}")
    return model, tokenizer


def chat_main() -> None:
    setup_runtime()
    model, tokenizer = load_export_for_chat()
    print("Interactive Cubic chat. Empty line exits.")
    while True:
        try:
            prompt = input("\nYou> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not prompt:
            break
        print("Cubic> ", end="", flush=True)
        print(generate(
            model, tokenizer, prompt,
            max_new_tokens=env_int("CUBIC_MAX_NEW_TOKENS", 512),
            temperature=env_float("CUBIC_TEMPERATURE", 0.60),
            top_p=env_float("CUBIC_TOP_P", 0.95),
            top_k=env_int("CUBIC_TOP_K", 50),
            reasoning=os.environ.get("CUBIC_REASONING", "0") == "1",
        ))


def train_main() -> None:
    setup_runtime()
    if not TOKENIZER_PATH.exists():
        fetch_hub_resume()
        if find_resume_checkpoint() is not None and not TOKENIZER_PATH.exists():
            raise RuntimeError(
                "A resumable checkpoint exists but tokenizer.json is missing. Training a "
                "fresh tokenizer would silently invalidate every learned embedding. "
                f"Restore {TOKENIZER_PATH} from work/tokenizer.json in {HF_CHECKPOINT_REPO}."
            )
    tokenizer = train_tokenizer()
    meta = prepare_everything(tokenizer)
    phases = phase_plan(meta)
    print("=" * 100)
    print("FAITHFUL CUBIC V7-HIER ~157M — GEMMA4/DEEPSEEK/SMOLLM3 IMPROVED PIPELINE")
    print(f"device={torch.cuda.get_device_name(0)} dtype={AMP_DTYPE} seq={CFG.seq_len:,} max_context={CFG.max_context:,}")
    print(
        f"compile={COMPILE_MODE if COMPILE else 'off'} "
        f"checkpoint_every={CHECKPOINT_EVERY} "
        f"effective_pretrain_batch={MICRO_BATCH * GRAD_ACCUM}"
    )
    print(f"data={meta}")
    for index, phase in enumerate(phases):
        print(f"  [{index}] {phase.name}: steps={phase.steps:,} micro={phase.micro_batch} accum={phase.accumulation}")
    print("=" * 100)

    model = CubicHierLM(CFG).to(DEVICE)
    parameter_count = sum(parameter.numel() for parameter in model.parameters())
    print(f"parameters={parameter_count:,} ({parameter_count / 1e6:.2f}M)")
    if parameter_count != EXPECTED_PARAMETER_COUNT:
        raise RuntimeError(
            f"architecture drift: got {parameter_count:,} parameters, "
            f"expected {EXPECTED_PARAMETER_COUNT:,}"
        )
    resume_phase, resume_step, metrics, resume_payload = load_checkpoint(model)
    run_model = compile_model(model)
    loss_fn = FusedTokenLoss()

    pretrain_train = PretrainSampler(PRETRAIN_TOKENS_FILE, CFG.seq_len, SEED + 100)
    pretrain_val = PretrainSampler(PRETRAIN_VAL_FILE, CFG.seq_len, SEED + 101)
    dialogue = RecordSampler(SFT_TOKENS_FILE, SFT_MASK_FILE, SEED + 200)
    reasoning = RecordSampler(REASON_TOKENS_FILE, REASON_MASK_FILE, SEED + 300)
    preference = PreferenceSampler(SEED + 400)
    samplers = {
        "pretrain": pretrain_train,
        # 10% base-LM replay during dialogue tuning and 20% dialogue replay
        # during reasoning tuning prevent curriculum-stage forgetting.
        "dialogue_sft": InterleavedSampler(dialogue, pretrain_train, replay_every=10),
        "reasoning_sft": InterleavedSampler(reasoning, dialogue, replay_every=5),
    }
    validation = {"pretrain": pretrain_val, "dialogue_sft": dialogue, "reasoning_sft": reasoning}

    for phase_index, phase in enumerate(phases):
        if phase_index < resume_phase:
            print(f"Skipping completed phase: {phase.name}")
            continue
        start = resume_step if phase_index == resume_phase else 0
        payload = resume_payload if phase_index == resume_phase else None
        print(f"\n{'=' * 36} PHASE {phase_index + 1}: {phase.name.upper()} {'=' * 36}")
        if phase.name == "dpo":
            metrics = train_dpo_phase(
                phase_index, phase, model, run_model, preference, start, payload, metrics,
            )
        else:
            metrics = train_language_phase(
                phase_index, phase, model, run_model, samplers[phase.name], loss_fn,
                start, payload, metrics, validation[phase.name],
            )
        record_samples(model, tokenizer, phase.name)
        resume_payload = None
        resume_step = 0

    metrics["parameter_count"] = parameter_count
    metrics["completed_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    save_checkpoint(model, None, len(phases), 0, metrics, FINAL)
    # Ensure the final state reaches the resume repository even if the last
    # regular hourly boundary was less than an hour ago.
    HOURLY_CHECKPOINT_PUBLISHER.finish(FINAL, len(phases), 0)
    export = export_model(model, tokenizer, metrics, phases)
    if UPLOAD:
        upload_export(export)
    print("\nDONE. Model, tokenizer, script, report and checkpoint are complete.")


if __name__ == "__main__":
    if MODE == "chat":
        chat_main()
    elif MODE == "train":
        train_main()
    else:
        raise SystemExit("CUBIC_MODE must be 'train' or 'chat'")