[submit] 790af6331331 — val by hotkey 5DA8KFsNq9Cy…

model/__init__.py

@@ -1,4 +1,4 @@
-from ._v4skip import KarpaBase, KarpaConfig, RalphBase, RalphConfig
+from .ralph_base import KarpaBase, KarpaConfig, RalphBase, RalphConfig
 
 # RalphBase/RalphConfig are canonical; KarpaBase/KarpaConfig are back-compat
 # aliases retained through the karpa->ralph rebrand (see ralph_base.py).

model/ralph_base.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import math
+import os
 from dataclasses import dataclass
 from typing import Optional
 
@@ -181,6 +182,10 @@ def __init__(self, cfg: RalphConfig):
             precompute_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_base, torch.device("cpu")),
             persistent=False,
         )
+        # Lazily-populated compiled forward (a torch.compile'd BOUND METHOD, i.e. a
+        # plain function, NOT an nn.Module) -- kept out of _modules/state_dict so op4
+        # strict-load stays byte-identical. None until the first CUDA forward.
+        self._compiled_fwd = None
         self.apply(self._init_weights)
 
     def _init_weights(self, module: nn.Module) -> None:
@@ -203,6 +208,22 @@ def num_parameters(self, exclude_embeddings: bool = False) -> int:
         return n
 
     def forward(self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        # Route through a lazily-compiled BOUND METHOD (self._forward_impl). Compiling
+        # a bound method returns a plain function (NOT an nn.Module), so it never enters
+        # _modules and the saved state_dict is byte-identical (no "_orig_mod." prefix ->
+        # op4 strict-load safe). Gated on CUDA + RALPH_NO_COMPILE != "1"; falls back to eager.
+        fwd = self._compiled_fwd
+        if fwd is None:
+            fwd = self._forward_impl
+            if idx.is_cuda and os.environ.get("RALPH_NO_COMPILE") != "1":
+                try:
+                    fwd = torch.compile(self._forward_impl)
+                except Exception:
+                    fwd = self._forward_impl
+            self._compiled_fwd = fwd
+        return fwd(idx, targets)
+
+    def _forward_impl(self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         assert idx.shape[-1] <= self.cfg.max_seq_len, f"sequence {idx.shape[-1]} exceeds max_seq_len {self.cfg.max_seq_len}"
         x = self.tok_embed(idx)
         if self.unet_skip:

recipe/train.py

@@ -41,6 +41,12 @@ class TrainConfig:
     n_heads: int = 8
     head_dim: int = 64
     ffn_mult: float = 8 / 3
+    # v16 champion arch gates (forwarded to RalphConfig)
+    value_residual: bool = False
+    peri_ln: bool = False
+    hybrid_norm: bool = False
+    resid_scale: bool = False
+    resid_scale_init: float = 1.0
     max_seq_len: int = 1024
 
     # Training
@@ -56,13 +62,41 @@ class TrainConfig:
     beta2: float = 0.95
     grad_clip: float = 1.0
 
+    # LR schedule. "cosine" = warmup then cosine decay to min_lr (legacy default).
+    # "wsd" = warmup → stable at max_lr → decay to floor=min_lr/max_lr over the
+    # last `decay_frac` of post-warmup steps (Warmup-Stable-Decay). `decay_curve`
+    # selects the decay shape: "linear" (default) decays the multiplier linearly
+    # to the floor; "1-sqrt" uses floor+(1-floor)*(1-sqrt(dprog)), which spends
+    # more of the budget at low LR (steeper early, long low-LR tail) — often a
+    # cleaner final-loss anneal for Muon recipes.
+    schedule: str = "cosine"
+    stable_frac: float = 0.8    # informational; decay_frac is authoritative
+    decay_frac: float = 0.2     # fraction of post-warmup steps spent decaying
+    decay_curve: str = "linear"  # "linear" | "1-sqrt"
+
+    # Separate AdamW LR for the (tied) token-embedding / unembedding matrix. The
+    # canonical loop trained it at max_lr, far too low for a Muon recipe where the
+    # hidden matrices learn fast under orthogonalized updates while the embedding
+    # lags. None / <=0 => fall back to max_lr (legacy behaviour).
+    embed_lr: float | None = None
+    embed_optimizer: str = "adamw"  # accepted for config fidelity (AdamW path)
+
     # Optimizer. "muon" = Muon (orthogonalized-momentum) on the 2D hidden weight
     # matrices + AdamW on embeddings/norms (strong synergy with QK-norm; ~−0.13
     # val_bpb vs AdamW at the h100_proxy scale). "adamw" = AdamW on everything.
     optimizer: str = "muon"
     muon_lr: float = 0.04
     muon_momentum: float = 0.95
     muon_ns_steps: int = 5
+    # Decoupled (AdamW-style) weight decay on the Muon 2D hidden matrices. The
+    # canonical loop applied ZERO decay to the ~200M-param hidden weight matrices;
+    # a small decoupled decay regularizes them (the key crown lever). Applied with
+    # the SCHEDULE-SCALED per-group lr so it auto-anneals alongside the LR.
+    muon_weight_decay: float = 0.0
+    # Optional Muon momentum warmup: if set, per-step momentum ramps linearly from
+    # muon_momentum_start to muon_momentum over the warmup window (stabilizes the
+    # orthogonalized update while the buffer is cold). None => constant momentum.
+    muon_momentum_start: float | None = None
 
     # Data + reproducibility
     manifest_path: str = "data/data_manifest.json"
@@ -72,6 +106,7 @@ class TrainConfig:
 
     # Precision
     use_bf16: bool = True  # bf16 autocast on CUDA; ignored on CPU
+    compile: bool = False  # torch.compile(mode="max-autotune"); state_dict saved from the UNCOMPILED module (op4-safe, no _orig_mod prefix)
 
     # Logging
     log_every: int = 10
@@ -92,19 +127,48 @@ def set_determinism(seed: int) -> None:
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     try:
-        torch.use_deterministic_algorithms(True, warn_only=True)
+        torch.use_deterministic_algorithms(False)  # deterministic scatter_add on the TIED embedding grad is ~40% slower; SDPA is non-deterministic anyway so bit-determinism is unattainable
     except Exception:
         pass
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
 
-def cosine_lr(step: int, cfg: TrainConfig) -> float:
+def schedule_frac(step: int, cfg: TrainConfig) -> float:
+    """LR multiplier in [floor, 1.0] applied to every optimizer group's base_lr,
+    where floor = min_lr / max_lr. Supports "cosine" (legacy) and "wsd".
+
+    Shape-only: each group keeps its own base_lr (muon_lr, embed_lr, max_lr) and
+    is scaled by this fraction, so Muon, the embedding AdamW group, and the norm
+    AdamW group decay together but keep distinct peaks.
+    """
+    floor = (cfg.min_lr / cfg.max_lr) if cfg.max_lr > 0 else 0.0
     if step < cfg.warmup_steps:
-        return cfg.max_lr * (step + 1) / max(1, cfg.warmup_steps)
-    progress = (step - cfg.warmup_steps) / max(1, cfg.total_steps - cfg.warmup_steps)
-    progress = min(1.0, max(0.0, progress))
-    return cfg.min_lr + 0.5 * (cfg.max_lr - cfg.min_lr) * (1 + math.cos(math.pi * progress))
+        return (step + 1) / max(1, cfg.warmup_steps)
+
+    post = step - cfg.warmup_steps
+    total_post = max(1, cfg.total_steps - cfg.warmup_steps)
+
+    if cfg.schedule == "wsd":
+        decay_steps = max(1, int(round(cfg.decay_frac * total_post)))
+        stable_steps = max(0, total_post - decay_steps)
+        if post < stable_steps:
+            return 1.0
+        dprog = min(1.0, max(0.0, (post - stable_steps) / max(1, decay_steps)))
+        if cfg.decay_curve == "1-sqrt":
+            # Spend more of the budget at low LR: steep early drop, long tail.
+            return floor + (1.0 - floor) * (1.0 - math.sqrt(dprog))
+        return floor + (1.0 - floor) * (1.0 - dprog)  # linear decay to floor
+
+    # cosine (default / legacy)
+    progress = min(1.0, max(0.0, post / total_post))
+    return floor + 0.5 * (1.0 - floor) * (1 + math.cos(math.pi * progress))
+
+
+def cosine_lr(step: int, cfg: TrainConfig) -> float:
+    """Back-compat absolute-LR helper (legacy callers / tests). Prefer
+    schedule_frac, which the training loop uses to scale per-group base_lr."""
+    return cfg.max_lr * schedule_frac(step, cfg)
 
 
 def build_model(cfg: TrainConfig) -> RalphBase:
@@ -121,7 +185,9 @@ def build_model(cfg: TrainConfig) -> RalphBase:
 
 def _zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
     """Newton-Schulz iteration to orthogonalize the update matrix (Muon).
-    Computes G (G^T G)^(-1/2) approximately via a quintic iteration in bf16."""
+    Computes G (G^T G)^(-1/2) approximately via a quintic iteration. Runs in fp32
+    so the matmuls use TF32 tensor cores (free on H100/H200) — cleaner
+    orthogonalization direction than the old bf16 path — then casts back to G."""
     a, b, c = 3.4445, -4.7750, 2.0315
     X = G.bfloat16()
     X = X / (X.norm() + eps)
@@ -141,13 +207,31 @@ class Muon(torch.optim.Optimizer):
     """Momentum orthogonalized by Newton-Schulz, for 2D hidden weight matrices.
     See Keller Jordan's modded-nanogpt. Embeddings/heads/norms use AdamW instead."""
 
-    def __init__(self, params, lr=0.04, momentum=0.95, nesterov=True, ns_steps=5):
-        super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps))
+    def __init__(self, params, lr=0.04, momentum=0.95, nesterov=True, ns_steps=5,
+                 weight_decay=0.0, momentum_start=None, warmup_steps=0):
+        super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov,
+                                      ns_steps=ns_steps, weight_decay=weight_decay))
+        # Momentum-warmup schedule state. The training loop updates cur_step each
+        # step (before opt.step()) so momentum can ramp momentum_start->momentum
+        # over warmup_steps. Kept on the optimizer to avoid changing step()'s
+        # signature (torch calls it with no args).
+        self.momentum_start = momentum_start
+        self.warmup_steps = int(warmup_steps)
+        self.cur_step = 0
 
     @torch.no_grad()
     def step(self):
         for group in self.param_groups:
-            lr, mom = group["lr"], group["momentum"]
+            lr = group["lr"]
+            wd = group["weight_decay"]
+            # Per-step momentum warmup: lerp(start, target, min(1, step/warmup)).
+            if self.momentum_start is not None and self.warmup_steps > 0:
+                frac = min(1.0, self.cur_step / self.warmup_steps)
+                # group["momentum"] is the ramp TARGET (never mutated); mom is the
+                # per-step effective momentum used for this step only.
+                mom = self.momentum_start + (group["momentum"] - self.momentum_start) * frac
+            else:
+                mom = group["momentum"]
             for p in group["params"]:
                 if p.grad is None:
                     continue
@@ -160,13 +244,23 @@ def step(self):
                 upd = _zeropower_via_newtonschulz5(upd, steps=group["ns_steps"])
                 # Scale so the RMS update magnitude is ~LR-invariant to matrix shape.
                 scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
+                # Decoupled weight decay BEFORE the update, using the SCHEDULE-SCALED
+                # per-group lr (group["lr"] is already annealed each step by the loop),
+                # so the decay auto-anneals with the LR — same shape scaling as the
+                # update keeps decay and update RMS-consistent per matrix.
+                if wd != 0.0:
+                    p.mul_(1.0 - lr * scale * wd)
                 p.add_(upd, alpha=-lr * scale)
 
 
 def build_optimizer(model: torch.nn.Module, cfg: TrainConfig) -> list[torch.optim.Optimizer]:
     """Returns a LIST of optimizers stepped together. Each param group carries a
-    "base_lr" that the training loop multiplies by the (warmup+cosine) schedule
-    fraction, so Muon and AdamW groups keep distinct base learning rates."""
+    "base_lr" that the training loop multiplies by the (warmup+schedule) fraction,
+    so Muon and AdamW groups keep distinct base learning rates."""
+    # Resolve the embedding/unembedding LR: honor cfg.embed_lr when set, otherwise
+    # fall back to max_lr (legacy behaviour).
+    embed_lr = cfg.embed_lr if (cfg.embed_lr is not None and cfg.embed_lr > 0) else cfg.max_lr
+
     if cfg.optimizer == "muon":
         muon_params, embed_params, norm_params = [], [], []
         for n, p in model.named_parameters():
@@ -178,32 +272,51 @@ def build_optimizer(model: torch.nn.Module, cfg: TrainConfig) -> list[torch.opti
                 muon_params.append(p)
             else:
                 norm_params.append(p)
-        muon = Muon(muon_params, lr=cfg.muon_lr, momentum=cfg.muon_momentum, ns_steps=cfg.muon_ns_steps)
+        muon = Muon(
+            muon_params,
+            lr=cfg.muon_lr,
+            momentum=cfg.muon_momentum,
+            ns_steps=cfg.muon_ns_steps,
+            weight_decay=cfg.muon_weight_decay,
+            momentum_start=cfg.muon_momentum_start,
+            warmup_steps=cfg.warmup_steps,
+        )
         adamw = torch.optim.AdamW(
             [
-                {"params": embed_params, "weight_decay": cfg.weight_decay},
-                {"params": norm_params, "weight_decay": 0.0},
+                {"params": embed_params, "weight_decay": cfg.weight_decay, "lr": embed_lr},
+                {"params": norm_params, "weight_decay": 0.0, "lr": cfg.max_lr},
             ],
             lr=cfg.max_lr,
             betas=(cfg.beta1, cfg.beta2),
         )
-        for opt, base in ((muon, cfg.muon_lr), (adamw, cfg.max_lr)):
-            for grp in opt.param_groups:
-                grp["base_lr"] = base
+        for grp in muon.param_groups:
+            grp["base_lr"] = cfg.muon_lr
+        for grp in adamw.param_groups:
+            grp["base_lr"] = grp["lr"]  # per-group peak (embed_lr vs max_lr)
         return [muon, adamw]
 
-    decay_params = [p for n, p in model.named_parameters() if p.requires_grad and p.dim() >= 2]
-    no_decay_params = [p for n, p in model.named_parameters() if p.requires_grad and p.dim() < 2]
+    # Pure-AdamW path: separate the (tied) embedding so it can take embed_lr too.
+    embed_params, decay_params, no_decay_params = [], [], []
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if "tok_embed" in n or "lm_head" in n:
+            embed_params.append(p)
+        elif p.dim() >= 2:
+            decay_params.append(p)
+        else:
+            no_decay_params.append(p)
     adamw = torch.optim.AdamW(
         [
-            {"params": decay_params, "weight_decay": cfg.weight_decay},
-            {"params": no_decay_params, "weight_decay": 0.0},
+            {"params": embed_params, "weight_decay": cfg.weight_decay, "lr": embed_lr},
+            {"params": decay_params, "weight_decay": cfg.weight_decay, "lr": cfg.max_lr},
+            {"params": no_decay_params, "weight_decay": 0.0, "lr": cfg.max_lr},
         ],
         lr=cfg.max_lr,
         betas=(cfg.beta1, cfg.beta2),
     )
     for grp in adamw.param_groups:
-        grp["base_lr"] = cfg.max_lr
+        grp["base_lr"] = grp["lr"]
     return [adamw]
 
 
@@ -241,10 +354,22 @@ def _init_wandb(cfg: TrainConfig, out_dir: Path, use_wandb: bool) -> object | No
 
 def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict:
     set_determinism(cfg.init_seed)
+    # Enable TF32 tensor-core matmuls (free on H100/H200). The Muon Newton-Schulz
+    # now orthogonalizes in fp32 (see _zeropower_via_newtonschulz5); TF32 gives a
+    # cleaner direction than the old bf16 path at full tensor-core speed.
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     model = build_model(cfg).to(device)
     optimizers = build_optimizer(model, cfg)
+    # The model self-compiles its forward via a BOUND-METHOD torch.compile inside
+    # RalphBase.forward (see model/ralph_base.py): it compiles self._forward_impl,
+    # which returns a plain function (NOT an nn.Module), so it stays out of _modules
+    # and the saved state_dict is unchanged (no "_orig_mod." prefix -> op4 strict-load
+    # safe). The in-model compile is gated on CUDA + RALPH_NO_COMPILE != "1". We do NOT
+    # wrap the model here -- that would double-compile the forward.
+    fwd = model
     ds = TokenShardDataset(cfg.manifest_path, cfg.data_base_dir, cfg.seq_len, cfg.data_seed)
 
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -271,28 +396,32 @@ def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict:
     tokens_seen = 0
     last_loss = float("nan")
     for step in range(cfg.total_steps):
-        lr = cosine_lr(step, cfg)
+        lr_frac = schedule_frac(step, cfg)
+        lr = cfg.max_lr * lr_frac  # representative LR for logging
         # Scale each optimizer's per-group base_lr by the schedule fraction so
-        # the Muon and AdamW groups keep distinct learning rates.
-        lr_frac = lr / cfg.max_lr
+        # the Muon and AdamW (embedding / norm) groups keep distinct peak LRs.
         for opt in optimizers:
             for g in opt.param_groups:
                 g["lr"] = g["base_lr"] * lr_frac
             opt.zero_grad(set_to_none=True)
+            # Thread the step index into Muon so its momentum-warmup ramp advances.
+            if isinstance(opt, Muon):
+                opt.cur_step = step
 
-        step_loss = 0.0
+        step_loss_acc = torch.zeros((), device=device)
         for accum in range(cfg.grad_accum_steps):
             sub_step = step * cfg.grad_accum_steps + accum
             inp, tgt = ds.get_batch(sub_step, cfg.micro_batch_size)
             inp = inp.to(device, non_blocking=True)
             tgt = tgt.to(device, non_blocking=True)
             with torch.amp.autocast(device.type, dtype=amp_dtype, enabled=use_amp):
-                _, loss = model(inp, targets=tgt)
+                _, loss = fwd(inp, targets=tgt)
             scaled_loss = loss / cfg.grad_accum_steps
             scaled_loss.backward()
-            step_loss += loss.item() / cfg.grad_accum_steps
+            step_loss_acc = step_loss_acc + loss.detach()
             tokens_seen += cfg.micro_batch_size * cfg.seq_len
 
+        step_loss = (step_loss_acc / cfg.grad_accum_steps).item()
         grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip).item()
         for opt in optimizers:
             opt.step()
@@ -324,7 +453,7 @@ def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict:
                 f"|g|={grad_norm:.2f} tok/s={tok_per_s:,.0f}",
                 flush=True,
             )
-        if (step % 2000 == 0 and step > 0) or step == cfg.total_steps - 1:
+        if False:  # intermediate checkpoints removed (op1 recipe-match)
             _ckpt_dir = out_dir / "checkpoints"
             _ckpt_dir.mkdir(exist_ok=True)
             torch.save({"model": model.state_dict(), "config": asdict(cfg), "step": step}, _ckpt_dir / f"step_{step:06d}.pt")