Axon-352M research baseline — DuoNeural 2026-07-05

README.md

@@ -0,0 +1,77 @@
+---
+language:
+- en
+license: apache-2.0
+tags:
+- causal-lm
+- pretraining
+- duoneural
+- research
+- custom-architecture
+- gqa
+- rope
+base_model: []
+pipeline_tag: text-generation
+---
+
+# Axon-352M
+
+**DuoNeural Research | 2026-07-05 | Archon**
+
+Research baseline model trained on smollm-corpus as a counterpart to CDM-based architectures. Used to establish a transformer baseline for comparative CDM studies.
+
+## Architecture
+
+Custom GPT-style transformer (not a standard HuggingFace architecture):
+
+| Parameter | Value |
+|-----------|-------|
+| Layers | 30 |
+| Hidden dim | 1024 |
+| FFN dim | 2560 |
+| Attention heads | 8Q / 4KV (GQA) |
+| Head dim | 128 |
+| Vocab size | 49,152 (SmolLM tokenizer) |
+| Max seq length | 2,048 |
+| Activation | ReLU² |
+| Normalization | RMSNorm + QK-norm |
+| Position encoding | RoPE (θ=10000) |
+| Logit cap | 30.0 |
+| Total params | ~352M |
+
+## Training
+
+- **Data**: smollm-corpus (FineWeb-edu-dedup 50%, Cosmopedia-v2 30%, OpenWebMath 10%, Python-edu 10%)
+- **Tokens**: ~8.5B
+- **Optimizer**: MuonH (matrix params) + AdamW (embeddings)
+- **Peak LR**: 3e-4 (trapezoidal: 5% warmup, 85% stable, 10% decay)
+- **Hardware**: RTX 3090 (1×) on vast.ai
+
+## Loading
+
+This model uses a custom architecture not directly loadable via `AutoModel`. To load:
+
+```python
+import torch
+from safetensors.torch import load_file
+
+# Load state dict
+state_dict = load_file("model.safetensors")
+
+# Architecture must be defined from training script
+# See train_axon_300m.py for the full model class
+```
+
+Full training script and architecture code available at [DuoNeural GitHub](https://github.com/DuoNeural).
+
+## Research Context
+
+Trained as a transformer baseline for the CDM (Competitive Docking Memory) research program. See:
+- [DuoNeural/CDM-Paper-1](https://huggingface.co/papers) — CDM architecture
+- [Zenodo DOI 10.5281/zenodo.21158430](https://zenodo.org/record/21158430) — CDM Paper 1
+
+## Authors
+
+Archon (DuoNeural Lab Director), Jesse Caldwell
+
+{DuoNeural standard footer}

config.json

@@ -0,0 +1,23 @@
+{
+  "model_type": "axon",
+  "architectures": [
+    "AxonForCausalLM"
+  ],
+  "hidden_size": 1024,
+  "intermediate_size": 2560,
+  "num_hidden_layers": 30,
+  "num_attention_heads": 8,
+  "num_key_value_heads": 4,
+  "vocab_size": 49152,
+  "max_position_embeddings": 2048,
+  "rope_theta": 10000.0,
+  "logit_softcap": 30.0,
+  "hidden_act": "relu2",
+  "rms_norm_eps": 1e-05,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "custom",
+  "training_tokens": 8499888128,
+  "best_train_loss": 4.409422159194946,
+  "training_step": 64849
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:873ad70c4bbe75a533c3ad07b2f1ce25b8bc8ccef177df702cc0f4e30e4ab0e2
+size 704813784

tokenizer_config.json

@@ -0,0 +1,6 @@
+{
+  "tokenizer_class": "GPT2Tokenizer",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048
+}

tokenizer_info.json

@@ -0,0 +1,4 @@
+{
+  "tokenizer_name": "HuggingFaceTB/smollm-corpus",
+  "vocab_size": 49152
+}

train_axon_300m.py

@@ -0,0 +1,627 @@
+"""
+train_300m.py — DuoNeural 300M Pretraining on RTX 5060Ti
+Archon, 2026-06-26. Aura's blueprint, my code.
+
+Hyperparams at a glance:
+  model: 302M params, 30L, H=1024, F=2560, GQA 8Q/4KV, ReLU², RMSNorm+QK-norm, logit_cap=30
+  seq: 2048  micro_batch: 4  grad_accum: 16  effective_batch: 64  tokens_per_step: 131072
+  optimizer: MuonH (2D weights) + AdamW (embeddings/1D)
+  LR: 3e-4 trapezoidal (5% warmup, 85% stable, 10% decay)
+  target: ~48B tokens over 15 days at ~35% MFU BF16
+  data: smollm-corpus (fineweb-edu-dedup 50% + cosmopedia-v2 30% + open-web-math 10% + python-edu 10%)
+"""
+
+import os, sys, math, time, json, logging, argparse, random
+from pathlib import Path
+from datetime import datetime
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+# ── liger (fused kernels) ──────────────────────────────────────────────────────
+# Triton JIT doesn't compile for sm_120 (Blackwell) on CUDA < 12.9.
+# PyTorch built-in ops (SDPA, matmul, RMSNorm-via-torch) work fine pre-built.
+# We implement chunked CE instead of liger fused CE to avoid the 800MB logits tensor.
+HAS_LIGER = False
+
+# ── datasets ───────────────────────────────────────────────────────────────────
+from datasets import load_dataset, interleave_datasets
+from transformers import AutoTokenizer
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# ARCHITECTURE
+# ══════════════════════════════════════════════════════════════════════════════
+
+HIDDEN    = 1024
+FFN_DIM   = 2560
+N_LAYERS  = 30
+N_HEADS   = 8
+N_KV      = 4
+HEAD_DIM  = HIDDEN // N_HEADS  # 128
+VOCAB     = 49152
+SEQ_LEN   = 2048
+LOGIT_CAP = 30.0
+
+
+# use PyTorch 2.4+ native RMSNorm (fused CUDA kernel, no Triton needed)
+RMSNorm = torch.nn.RMSNorm
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, max_seq: int = SEQ_LEN, base: float = 10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        t = torch.arange(max_seq).float()
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None], persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[None, None], persistent=False)
+
+    @staticmethod
+    def _rotate_half(x):
+        h = x.shape[-1] // 2
+        return torch.cat([-x[..., h:], x[..., :h]], dim=-1)
+
+    def forward(self, q, k, T):
+        cos = self.cos_cached[:, :, :T]
+        sin = self.sin_cached[:, :, :T]
+        q = (q * cos) + (self._rotate_half(q) * sin)
+        k = (k * cos) + (self._rotate_half(k) * sin)
+        return q, k
+
+
+class Attention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(HIDDEN, N_HEADS * HEAD_DIM, bias=False)
+        self.k_proj = nn.Linear(HIDDEN, N_KV   * HEAD_DIM, bias=False)
+        self.v_proj = nn.Linear(HIDDEN, N_KV   * HEAD_DIM, bias=False)
+        self.o_proj = nn.Linear(N_HEADS * HEAD_DIM, HIDDEN, bias=False)
+        self.q_norm = RMSNorm(HEAD_DIM)  # QK-norm
+        self.k_norm = RMSNorm(HEAD_DIM)
+        self.rope   = RotaryEmbedding(HEAD_DIM)
+        self._groups = N_HEADS // N_KV
+
+    def forward(self, x):
+        B, T, _ = x.shape
+        q = self.q_proj(x).view(B, T, N_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, N_KV,   HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, N_KV,   HEAD_DIM).transpose(1, 2)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q, k = self.rope(q, k, T)
+        k = k.repeat_interleave(self._groups, dim=1)
+        v = v.repeat_interleave(self._groups, dim=1)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.o_proj(out.transpose(1, 2).contiguous().view(B, T, -1))
+
+
+class FFN(nn.Module):
+    """ReLU² — 90%+ sparsity, max bandwidth efficiency on GDDR7."""
+    def __init__(self):
+        super().__init__()
+        self.up   = nn.Linear(HIDDEN, FFN_DIM, bias=False)
+        self.down = nn.Linear(FFN_DIM, HIDDEN, bias=False)
+
+    def forward(self, x):
+        return self.down(F.relu(self.up(x)).pow(2))
+
+
+class Block(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn_norm = RMSNorm(HIDDEN)
+        self.attn      = Attention()
+        self.ffn_norm  = RMSNorm(HIDDEN)
+        self.ffn       = FFN()
+
+    def forward(self, x):
+        x = x + self.attn(self.attn_norm(x))
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class DuoNeural300M(nn.Module):
+    """302M param LM. Aura's architecture. Archon's implementation. DuoNeural's model."""
+    def __init__(self, grad_ckpt=True):
+        super().__init__()
+        self.grad_ckpt = grad_ckpt
+        self.embed   = nn.Embedding(VOCAB, HIDDEN)
+        self.blocks  = nn.ModuleList([Block() for _ in range(N_LAYERS)])
+        self.norm    = RMSNorm(HIDDEN)
+        self.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+        self._init_weights()
+
+    def _init_weights(self):
+        std = 0.02
+        nn.init.normal_(self.embed.weight, std=std)
+        nn.init.normal_(self.lm_head.weight, std=std)
+        for block in self.blocks:
+            for name in ["q_proj", "k_proj", "v_proj"]:
+                nn.init.normal_(getattr(block.attn, name).weight, std=std)
+            nn.init.normal_(block.attn.o_proj.weight, std=std / math.sqrt(2 * N_LAYERS))
+            nn.init.normal_(block.ffn.up.weight,      std=std)
+            nn.init.normal_(block.ffn.down.weight,    std=std / math.sqrt(2 * N_LAYERS))
+
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def forward_hidden(self, input_ids):
+        """Returns final hidden states — caller applies lm_head + fused CE."""
+        x = self.embed(input_ids)
+        for block in self.blocks:
+            if self.grad_ckpt and self.training:
+                x = checkpoint(block, x, use_reentrant=False)
+            else:
+                x = block(x)
+        return self.norm(x)
+
+    def forward(self, input_ids):
+        """Inference-only forward. logit soft-cap applied."""
+        x = self.forward_hidden(input_ids)
+        logits = self.lm_head(x)
+        return LOGIT_CAP * torch.tanh(logits / LOGIT_CAP)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# MUON OPTIMIZER (Newton-Schulz orthogonalization, BF16)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def _ns_orthogonalize(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
+    """
+    Newton-Schulz iteration to orthogonalize gradient G.
+    Runs in BF16 for bandwidth efficiency on Blackwell.
+    a, b, c coefficients from the Muon paper.
+    """
+    assert G.ndim >= 2
+    orig_dtype = G.dtype
+    X = G.bfloat16()
+    # normalize so NS iterations are numerically stable
+    X = X / (X.norm() + 1e-7)
+    transposed = X.shape[0] > X.shape[1]
+    if transposed:
+        X = X.T
+    a, b, c = 3.4445, -4.7750, 2.0315
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+    if transposed:
+        X = X.T
+    return X.to(orig_dtype)
+
+
+class MuonH(torch.optim.Optimizer):
+    """
+    Muon Hyperball optimizer — NS-orthogonalized SGD with momentum.
+    Apply to ALL 2D weight matrices EXCEPT embeddings and lm_head.
+    Frobenius norm bounds enforced post-step (the 'Hyperball' variant).
+    """
+    def __init__(self, params, lr=3e-4, momentum=0.95, ns_steps=5, norm_bound=None):
+        defaults = dict(lr=lr, momentum=momentum, ns_steps=ns_steps, norm_bound=norm_bound)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr         = group["lr"]
+            momentum   = group["momentum"]
+            ns_steps   = group["ns_steps"]
+            norm_bound = group["norm_bound"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                g = p.grad
+
+                state = self.state[p]
+                if "momentum_buf" not in state:
+                    state["momentum_buf"] = torch.zeros_like(p)
+                    if norm_bound is not None:
+                        state["init_norm"] = p.data.norm(p=2).item()
+
+                buf = state["momentum_buf"]
+                buf.mul_(momentum).add_(g)
+
+                # NS orthogonalize the accumulated gradient
+                update = _ns_orthogonalize(buf, steps=ns_steps)
+                # scale to match gradient magnitude for lr calibration
+                update = update * (g.norm() / (update.norm() + 1e-8))
+
+                p.data.add_(update, alpha=-lr)
+
+                # Hyperball: re-project weight onto fixed-norm sphere after update
+                if norm_bound is not None:
+                    curr_norm = p.data.norm(p=2)
+                    target    = state["init_norm"] if norm_bound == "init" else norm_bound
+                    if curr_norm > target:
+                        p.data.mul_(target / (curr_norm + 1e-8))
+
+        return loss
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# LR SCHEDULE — trapezoidal (5% warmup | 85% stable | 10% cosine decay)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def get_lr(step: int, total_steps: int, peak_lr: float, min_lr: float = 0.0) -> float:
+    warmup_end  = int(0.05 * total_steps)
+    stable_end  = int(0.90 * total_steps)
+
+    if step <= warmup_end:
+        return peak_lr * (step / max(warmup_end, 1))
+    elif step <= stable_end:
+        return peak_lr
+    else:
+        # cosine decay from stable_end to total_steps
+        decay_steps = total_steps - stable_end
+        progress    = (step - stable_end) / max(decay_steps, 1)
+        return min_lr + 0.5 * (peak_lr - min_lr) * (1 + math.cos(math.pi * progress))
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# DATA PIPELINE
+# ══════════════════════════════════════════════════════════════════════════════
+
+DATASET_MIX = [
+    # (hf_name, config_or_None, text_field, sampling_weight)
+    # open-web-math is its own repo (not a smollm-corpus config)
+    ("HuggingFaceTB/smollm-corpus", "fineweb-edu-dedup", "text", 0.50),
+    ("HuggingFaceTB/smollm-corpus", "cosmopedia-v2",     "text", 0.30),
+    ("open-web-math/open-web-math",  None,                "text", 0.10),
+    ("HuggingFaceTB/smollm-corpus", "python-edu",        "text", 0.10),
+]
+
+
+def build_streaming_dataset(tokenizer, seq_len: int, seed: int = 42):
+    """
+    Builds a packed streaming dataset from the 4-source mixture.
+    Manual weighted interleaving — bypasses interleave_datasets schema issues.
+    Documents are concatenated with EOS tokens and chunked to seq_len+1.
+    No padding ever — GPU sees 100% real tokens.
+    """
+    rng = random.Random(seed)
+    eos_id = tokenizer.eos_token_id
+
+    streams      = []
+    text_fields  = []
+    weights      = []
+    for name, config, text_field, w in DATASET_MIX:
+        if config is not None:
+            ds = load_dataset(name, config, split="train", streaming=True)
+        else:
+            ds = load_dataset(name, split="train", streaming=True)
+        streams.append(ds)
+        text_fields.append(text_field)
+        weights.append(w)
+
+    def pack_documents_multi(streams, text_fields, weights):
+        """Generator: weighted random sampling across streams, pack into seq chunks."""
+        iters = [iter(s) for s in streams]
+        buf   = []
+        while True:
+            # sample which stream to pull from (weighted)
+            idx = rng.choices(range(len(streams)), weights=weights, k=1)[0]
+            try:
+                example = next(iters[idx])
+            except StopIteration:
+                # reset exhausted stream
+                iters[idx] = iter(streams[idx])
+                example = next(iters[idx])
+
+            text = example.get(text_fields[idx]) or ""
+            if not isinstance(text, str):
+                text = str(text)
+            if not text.strip():
+                continue
+
+            ids = tokenizer.encode(text, add_special_tokens=False)
+            buf.extend(ids)
+            buf.append(eos_id)
+            while len(buf) >= seq_len + 1:
+                yield buf[:seq_len + 1]
+                buf = buf[seq_len + 1:]
+
+    return streams, lambda _: pack_documents_multi(streams, text_fields, weights)
+
+
+class StreamingTokenDataLoader:
+    """
+    Wraps the streaming HF dataset into a simple micro-batch iterator.
+    pin_memory for PCIe 5.0 x8 transfer speed — keeps GPU fed.
+    """
+    def __init__(self, dataset, pack_fn, micro_batch: int, seq_len: int,
+                 num_workers: int = 4, prefetch: int = 64):
+        self.dataset     = dataset
+        self.pack_fn     = pack_fn
+        self.micro_batch = micro_batch
+        self.seq_len     = seq_len
+        self.prefetch    = prefetch
+
+    def __iter__(self):
+        buffer = []
+        gen    = self.pack_fn(self.dataset)
+        for chunk in gen:
+            buffer.append(chunk)
+            if len(buffer) >= self.micro_batch:
+                t = torch.tensor(buffer[:self.micro_batch], dtype=torch.long)
+                input_ids = t[:, :-1]
+                labels    = t[:, 1:].contiguous()
+                yield input_ids.pin_memory(), labels.pin_memory()
+                buffer = buffer[self.micro_batch:]
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# LOSS
+# ══════════════════════════════════════════════════════════════════════════════
+
+def compute_loss(model: DuoNeural300M, input_ids: torch.Tensor, labels: torch.Tensor):
+    """
+    Training loss. With grad checkpointing we use only ~6GB peak VRAM, so we can
+    afford to materialize the full (B*T, 49152) logits tensor (~800MB at mb=4).
+    Logit soft-cap (tanh ×30) applied before CE to prevent late-training NaN.
+    """
+    raw_model = getattr(model, "_orig_mod", model)
+    hidden  = model.forward_hidden(input_ids)             # (B, T, H)
+    logits  = raw_model.lm_head(hidden)                   # (B, T, V)
+    logits  = LOGIT_CAP * torch.tanh(logits / LOGIT_CAP)
+    return F.cross_entropy(logits.reshape(-1, VOCAB), labels.reshape(-1))
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# CHECKPOINT
+# ══════════════════════════════════════════════════════════════════════════════
+
+def save_ckpt(model, muon_opt, adamw_opt, step, loss, ckpt_dir: Path):
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    path = ckpt_dir / f"step_{step:07d}.pt"
+    torch.save({
+        "step":        step,
+        "loss":        loss,
+        "model":       model.state_dict(),
+        "muon":        muon_opt.state_dict(),
+        "adamw":       adamw_opt.state_dict(),
+        "timestamp":   datetime.utcnow().isoformat(),
+    }, path)
+    # keep only last 3 checkpoints to save disk
+    ckpts = sorted(ckpt_dir.glob("step_*.pt"))
+    for old in ckpts[:-3]:
+        old.unlink()
+    print(f"✓ checkpoint saved: {path}")
+    return path
+
+
+def load_ckpt(model, muon_opt, adamw_opt, ckpt_dir: Path):
+    ckpts = sorted(ckpt_dir.glob("step_*.pt"))
+    if not ckpts:
+        return 0
+    path = ckpts[-1]
+    ckpt = torch.load(path, map_location="cpu", weights_only=True)
+    model.load_state_dict(ckpt["model"])
+    muon_opt.load_state_dict(ckpt["muon"])
+    adamw_opt.load_state_dict(ckpt["adamw"])
+    step = ckpt["step"]
+    print(f"✓ resumed from {path} at step {step}, loss={ckpt['loss']:.4f}")
+    return step
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# TRAINING
+# ══════════════════════════════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt_dir",     default="/workspace/ckpts_300m")
+    parser.add_argument("--log_file",     default="/workspace/train_300m.log")
+    parser.add_argument("--micro_batch",  type=int, default=8)
+    parser.add_argument("--grad_accum",   type=int, default=8)
+    parser.add_argument("--total_tokens", type=float, default=8.5e9,
+                        help="total training tokens (default: 8.5B = Chinchilla-optimal for 352M)")
+    parser.add_argument("--peak_lr_muon",  type=float, default=3e-4)
+    parser.add_argument("--peak_lr_adam",  type=float, default=3e-4)
+    parser.add_argument("--save_every",   type=int, default=1000)
+    parser.add_argument("--log_every",    type=int, default=50)
+    parser.add_argument("--seed",         type=int, default=42)
+    parser.add_argument("--compile",      action="store_true",
+                        help="torch.compile the model — faster but slower startup")
+    args = parser.parse_args()
+
+    # ── setup ────────────────────────────────────────────────────────────────
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(message)s",
+        handlers=[
+            logging.StreamHandler(sys.stdout),
+            logging.FileHandler(args.log_file),
+        ]
+    )
+    log = logging.getLogger("train_300m")
+
+    tokens_per_step = args.micro_batch * args.grad_accum * SEQ_LEN
+    total_steps     = int(args.total_tokens / tokens_per_step)
+
+    log.info(f"{'='*60}")
+    log.info(f"DuoNeural 300M Pretraining")
+    log.info(f"device:         {device} ({torch.cuda.get_device_name(0)})")
+    log.info(f"micro_batch:    {args.micro_batch}")
+    log.info(f"grad_accum:     {args.grad_accum}")
+    log.info(f"tokens_per_step:{tokens_per_step:,}")
+    log.info(f"total_steps:    {total_steps:,}")
+    log.info(f"total_tokens:   {args.total_tokens/1e9:.1f}B")
+    log.info(f"liger_kernel:   {HAS_LIGER}")
+    log.info(f"torch.compile:  {args.compile}")
+    log.info(f"{'='*60}")
+
+    # ── tokenizer ────────────────────────────────────────────────────────────
+    log.info("loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        "HuggingFaceTB/SmolLM2-135M",
+        cache_dir="/workspace/hf_cache",
+    )
+    assert tokenizer.vocab_size == VOCAB, \
+        f"tokenizer vocab {tokenizer.vocab_size} != model vocab {VOCAB}"
+    log.info(f"tokenizer: vocab_size={tokenizer.vocab_size}")
+
+    # ── model ─────────────────────────────────────────────────────────────────
+    log.info("building model...")
+    model = DuoNeural300M(grad_ckpt=True).to(device)
+    model = model.bfloat16()
+    n_params = model.n_params()
+    log.info(f"model params: {n_params/1e6:.1f}M")
+
+    if args.compile:
+        log.info("compiling model (this takes 1-2 min on first run)...")
+        model = torch.compile(model, mode="reduce-overhead")
+
+    # ── optimizers ────────────────────────────────────────────────────────────
+    # Muon: all 2D weight matrices EXCEPT embed and lm_head (vocab dim too large for NS)
+    muon_params  = []
+    adamw_params = []
+    for name, p in model.named_parameters():
+        if p.ndim == 2 and "embed" not in name and "lm_head" not in name:
+            muon_params.append(p)
+        else:
+            adamw_params.append(p)
+
+    log.info(f"muon params:  {sum(p.numel() for p in muon_params)/1e6:.1f}M")
+    log.info(f"adamw params: {sum(p.numel() for p in adamw_params)/1e6:.1f}M")
+
+    muon_opt  = MuonH(muon_params,  lr=args.peak_lr_muon,  momentum=0.95, ns_steps=5)
+    adamw_opt = torch.optim.AdamW(adamw_params, lr=args.peak_lr_adam,
+                                   betas=(0.9, 0.95), weight_decay=0.1, eps=1e-8)
+
+    # ── resume ────────────────────────────────────────────────────────────────
+    ckpt_dir  = Path(args.ckpt_dir)
+    start_step = load_ckpt(model, muon_opt, adamw_opt, ckpt_dir)
+
+    # ── data ──────────────────────────────────────────────────────────────────
+    log.info("building data pipeline...")
+    raw_ds, pack_fn = build_streaming_dataset(tokenizer, SEQ_LEN, seed=args.seed)
+    loader = StreamingTokenDataLoader(raw_ds, pack_fn,
+                                      micro_batch=args.micro_batch,
+                                      seq_len=SEQ_LEN)
+
+    # ── training loop ─────────────────────────────────────────────────────────
+    model.train()
+    scaler = torch.amp.GradScaler("cuda", enabled=False)  # BF16 doesn't need GradScaler
+
+    total_tokens_seen = start_step * tokens_per_step
+    step = start_step
+    accum_loss  = 0.0
+    accum_count = 0
+    t0 = time.time()
+
+    log.info(f"starting training from step {step}/{total_steps}")
+    log.info(f"{'step':>7} {'loss':>8} {'lr':>10} {'tok/s':>8} {'tokens':>12}")
+
+    data_iter = iter(loader)
+
+    while step < total_steps:
+        # ── LR update ────────────────────────────────────────────────────────
+        lr = get_lr(step, total_steps, args.peak_lr_muon)
+        for pg in muon_opt.param_groups:
+            pg["lr"] = lr
+        lr_adam = get_lr(step, total_steps, args.peak_lr_adam)
+        for pg in adamw_opt.param_groups:
+            pg["lr"] = lr_adam
+
+        # ── gradient accumulation ─────────────────────────────────────────────
+        muon_opt.zero_grad()
+        adamw_opt.zero_grad()
+
+        step_loss = 0.0
+        for micro_step in range(args.grad_accum):
+            try:
+                input_ids, labels = next(data_iter)
+            except StopIteration:
+                # dataset exhausted — restart (48B tokens should outlast 15 days)
+                log.info("dataset exhausted, restarting...")
+                data_iter = iter(loader)
+                input_ids, labels = next(data_iter)
+
+            input_ids = input_ids.to(device, non_blocking=True)
+            labels    = labels.to(device, non_blocking=True)
+
+            with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+                loss = compute_loss(model, input_ids, labels)
+                loss = loss / args.grad_accum
+
+            loss.backward()
+            step_loss += loss.item()
+
+        # gradient clipping — essential with Muon's aggressive LR
+        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+        muon_opt.step()
+        adamw_opt.step()
+
+        step += 1
+        total_tokens_seen += tokens_per_step
+        accum_loss  += step_loss
+        accum_count += 1
+
+        # ── logging ───────────────────────────────────────────────────────────
+        if step % args.log_every == 0:
+            elapsed    = time.time() - t0
+            tok_s      = (args.log_every * tokens_per_step) / elapsed
+            avg_loss   = accum_loss / accum_count
+            log.info(f"{step:>7} {avg_loss:>8.4f} {lr:>10.2e} {tok_s:>8.0f} {total_tokens_seen/1e9:>10.3f}B")
+            accum_loss  = 0.0
+            accum_count = 0
+            t0 = time.time()
+
+        # ── checkpoint ────────────────────────────────────────────────────────
+        if step % args.save_every == 0:
+            save_ckpt(model, muon_opt, adamw_opt, step, step_loss, ckpt_dir)
+
+    log.info(f"training complete! {total_tokens_seen/1e9:.2f}B tokens, {step} steps")
+    save_ckpt(model, muon_opt, adamw_opt, step, step_loss, ckpt_dir)
+
+    # final save as HF-compatible (for SFT afterwards)
+    _save_hf_model(model, tokenizer, "/workspace/duoneural_300m_final")
+    log.info("model saved to /workspace/duoneural_300m_final")
+
+
+def _save_hf_model(model, tokenizer, output_dir: str):
+    """Save model in transformers-compatible format for later SFT."""
+    from transformers import PreTrainedModel, PretrainedConfig
+
+    class DN300MConfig(PretrainedConfig):
+        model_type = "duoneural_300m"
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.hidden_size   = HIDDEN
+            self.ffn_dim       = FFN_DIM
+            self.num_layers    = N_LAYERS
+            self.num_heads     = N_HEADS
+            self.num_kv_heads  = N_KV
+            self.vocab_size    = VOCAB
+            self.seq_len       = SEQ_LEN
+            self.logit_cap     = LOGIT_CAP
+
+    config = DN300MConfig()
+    path   = Path(output_dir)
+    path.mkdir(parents=True, exist_ok=True)
+    config.save_pretrained(str(path))
+
+    # unwrap compile if needed
+    raw_model = getattr(model, "_orig_mod", model)
+    torch.save(raw_model.state_dict(), path / "model_weights.pt")
+    tokenizer.save_pretrained(str(path))
+    print(f"HF-style checkpoint saved to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()