# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "datasets==5.0.0",
#     "huggingface-hub==1.24.0",
#     "tokenizers==0.23.1",
#     "triton==3.7.1",
# ]
# ///

import marimo

__generated_with = "0.23.15"
app = marimo.App(
    width="medium",
    layout_file="layouts/notebook.grid.json",
    auto_download=["html"],
)


@app.cell
def _():
    #!/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'")
    return


if __name__ == "__main__":
    app.run()