[submit] 3f607e133b32 — val by hotkey 5DCytiALFXhm…

configs/h100_proxy.json

@@ -1,22 +1,37 @@
 {
-  "_comment": "b300_16k_lr6: 254M H100 proxy model, 16k steps, seq_len 512.",
+  "_comment": "crown8k_mwd FULL 8000st: clean_v4+MuonWD0.05+1sqrt bf16-NS default-compile",
   "vocab_size": 50257,
-  "dim": 1024,
+  "dim": 768,
   "n_layers": 16,
-  "n_heads": 16,
+  "n_heads": 12,
   "head_dim": 64,
   "ffn_mult": 2.6875,
   "max_seq_len": 1024,
   "seq_len": 512,
   "batch_size": 512,
-  "micro_batch_size": 128,
-  "total_steps": 2400,
-  "warmup_steps": 240,
+  "micro_batch_size": 64,
+  "total_steps": 11500,
+  "warmup_steps": 600,
   "max_lr": 0.0008,
-  "min_lr": 3e-05,
+  "min_lr": 1e-05,
   "weight_decay": 0.1,
   "beta1": 0.9,
   "beta2": 0.98,
   "grad_clip": 1.0,
-  "log_every": 50
-}
+  "log_every": 500,
+  "schedule": "wsd",
+  "stable_frac": 0.55,
+  "decay_frac": 0.45,
+  "decay_curve": "1-sqrt",
+  "optimizer": "muon",
+  "muon_lr": 0.04,
+  "muon_momentum": 0.95,
+  "muon_weight_decay": 0.05,
+  "embed_optimizer": "adamw",
+  "embed_lr": 0.006,
+  "compile": true,
+  "init_seed": 6565,
+  "data_seed": 6565,
+  "data_curriculum": "random",
+  "_resubmit": "retry-scoring-1244"
+}

data/dataset.py

@@ -56,6 +56,17 @@ def __init__(
         self.seed = seed
         if self._total < seq_len + 1:
             raise ValueError(f"not enough tokens ({self._total}) for seq_len {seq_len}")
+        # No-replacement permuted-sweep bookkeeping. We tile the token stream into
+        # `n_blocks` non-overlapping contiguous blocks of (seq_len+1) tokens each
+        # and sweep them in a deterministic seed-keyed permutation per epoch. This
+        # yields full-coverage, no-repeat-within-epoch sampling (a strictly better
+        # gradient signal than with-replacement draws — no wasted duplicate windows
+        # and no unseen holes) while remaining a pure deterministic function of
+        # (seed, step): the validator re-derives the exact same order on audit.
+        # Per-epoch permutations are MEMOIZED — never rebuilt per get_batch call —
+        # so throughput is unaffected.
+        self._n_blocks = max(1, self._total // (self.seq_len + 1))
+        self._perm_cache: dict[int, np.ndarray] = {}
 
     @property
     def total_tokens(self) -> int:
@@ -103,18 +114,43 @@ def get(self, step: int) -> tuple[torch.Tensor, torch.Tensor]:
         ids = torch.from_numpy(chunk.astype(np.int64))
         return ids[:-1], ids[1:]
 
+    def _epoch_perm(self, epoch: int) -> np.ndarray:
+        """MEMOIZED deterministic permutation of the `n_blocks` contiguous blocks
+        for the given epoch. Keyed by (seed, 0xE9C, epoch) so it is bit-identical
+        across runs (validator re-derivation) and never rebuilt once cached."""
+        perm = self._perm_cache.get(epoch)
+        if perm is None:
+            rng = np.random.default_rng(
+                np.array([self.seed, 0xE9C, epoch], dtype=np.uint64)
+            )
+            perm = rng.permutation(self._n_blocks)
+            self._perm_cache[epoch] = perm
+        return perm
+
     def get_batch(
         self,
         step: int,
         batch_size: int,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Return a batch of (B, T) input + target tensors at the given step."""
-        rng = np.random.default_rng(np.array([self.seed, step], dtype=np.uint64))
-        starts = rng.integers(0, self._total, size=batch_size)
+        """Return a batch of (B, T) input + target tensors at the given step.
+
+        No-replacement permuted sweep: for row `b` the global block index is
+        gi = step*batch_size + b, decomposed into epoch = gi // n_blocks and
+        pos = gi % n_blocks; the sampled block is memoized_perm(epoch)[pos] and
+        the window starts at block*(seq_len+1). Deterministic in (seed, step, b)
+        and bit-reproducible; per-epoch perms are memoized so tok/s is unaffected.
+        """
+        n_blocks = self._n_blocks
+        stride = self.seq_len + 1
         inputs = np.empty((batch_size, self.seq_len), dtype=np.int64)
         targets = np.empty((batch_size, self.seq_len), dtype=np.int64)
-        for b, s in enumerate(starts):
-            chunk = self._read_range(int(s), self.seq_len + 1).astype(np.int64)
+        for b in range(batch_size):
+            gi = step * batch_size + b
+            epoch = gi // n_blocks
+            pos = gi % n_blocks
+            block = int(self._epoch_perm(epoch)[pos])
+            start = block * stride
+            chunk = self._read_range(start, stride).astype(np.int64)
             inputs[b] = chunk[:-1]
             targets[b] = chunk[1:]
         return torch.from_numpy(inputs), torch.from_numpy(targets)

model/_v4skip.py

@@ -38,6 +38,13 @@ class RalphConfig:
     tie_embeddings: bool = True
     unet_skip: bool = True        # recipe-v4: U-Net learnable skip connections
     logit_softcap: float = 30.0   # recipe-v4: tanh soft-cap on logits (0 = off)
+    # recipe-v5: value embeddings (modded-nanogpt VE) — a separate token-id table
+    # whose lookup is blended into each block's attention V by learnable per-layer
+    # mixing scalars. Param-heavy but compute-cheap (an embedding gather, no matmul);
+    # a scale-robust val_bpb win that stacks on top of the v4 stack.
+    value_embeddings: bool = True   # off => no value_embed / ve_lambda params
+    ve_lambda_v: float = 1.0        # per-layer init: v_out = ve_lambda_v * v + ve_lambda_ve * ve
+    ve_lambda_ve: float = 1.0
 
 
 def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
@@ -104,7 +111,8 @@ def __init__(self, cfg: RalphConfig):
         self.q_norm = RMSNorm(cfg.head_dim, cfg.rms_norm_eps)
         self.k_norm = RMSNorm(cfg.head_dim, cfg.rms_norm_eps)
 
-    def forward(self, x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, rope_cache: torch.Tensor,
+                ve: Optional[torch.Tensor] = None, ve_lam: Optional[torch.Tensor] = None) -> torch.Tensor:
         B, T, C = x.shape
         qkv = self.qkv(x)  # (B, T, 3C)
         q, k, v = qkv.split(self.dim, dim=-1)
@@ -115,6 +123,11 @@ def forward(self, x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
         k = self.k_norm(k)
         q = apply_rope(q, rope_cache)
         k = apply_rope(k, rope_cache)
+        # recipe-v5: value embeddings — blend the per-token value-table lookup into V
+        # with the layer's learnable mixing scalars (v is not rotated, so blend here).
+        if ve is not None:
+            vh = ve.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+            v = ve_lam[0] * v + ve_lam[1] * vh
         # Causal self-attention via SDPA (uses flash on supported hardware).
         y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
         y = y.transpose(1, 2).contiguous().view(B, T, C)
@@ -145,9 +158,10 @@ def __init__(self, cfg: RalphConfig):
         self.ffn_norm = RMSNorm(cfg.dim, cfg.rms_norm_eps)
         self.ffn = SwiGLU(cfg)
 
-    def forward(self, x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
-        x = x + self.attn(self.attn_norm(x), rope_cache)
-        x = x + self.ffn(self.ffn_norm(x))
+    def forward(self, x: torch.Tensor, rope_cache: torch.Tensor,
+                ve: Optional[torch.Tensor] = None, ve_lam: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x = x + self.attn_norm(self.attn(x, rope_cache, ve, ve_lam))
+        x = x + self.ffn_norm(self.ffn(x))
         return x
 
 
@@ -167,11 +181,34 @@ def __init__(self, cfg: RalphConfig):
         self.unet_skip = getattr(cfg, "unet_skip", False)
         if self.unet_skip:
             self.skip_gate = nn.Parameter(torch.zeros(cfg.n_layers - cfg.n_layers // 2))
+        # recipe-v5: value-embedding table + per-layer learnable mixing scalars,
+        # gathered once per forward and blended into every block's attention V.
+        self.value_embeddings = getattr(cfg, "value_embeddings", False)
+        if self.value_embeddings:
+            self.value_embed = nn.Embedding(cfg.vocab_size, cfg.dim)
+            self.ve_lambda = nn.Parameter(
+                torch.tensor([[cfg.ve_lambda_v, cfg.ve_lambda_ve]] * cfg.n_layers, dtype=torch.float32)
+            )
         self.final_norm = RMSNorm(cfg.dim, cfg.rms_norm_eps)
         if cfg.tie_embeddings:
             self.lm_head = None
         else:
             self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
+        # Learned scalar readout temperature for the (tied) head. With weight
+        # tying, the embedding matrix's norm sets BOTH the input-embedding scale
+        # (init_std=0.02) and the readout/logit scale; the model cannot move one
+        # without moving the other. This single scalar decouples the readout gain
+        # from the embedding norm so the softmax temperature is fit directly.
+        # exp() parameterisation keeps it strictly positive; 0-init => gain 1.0,
+        # so the run starts bit-identical to the reordered base and learns away.
+        # Shape () => 1D => routed to AdamW(no-decay) by build_optimizer.
+        self.logit_scale = nn.Parameter(torch.zeros(()))
+        # Per-vocab readout calibration (readcal): per-token multiplicative gain +
+        # additive bias on the tied head. 0-init => exp(0)=1 and +0 => identity at
+        # step 0. Both shape (vocab,) => 1D => AdamW(no-decay). Adds state-dict keys
+        # so op4 routes to the patched-eval path (scored as the real arch).
+        self.readout_gain = nn.Parameter(torch.zeros(cfg.vocab_size))
+        self.readout_bias = nn.Parameter(torch.zeros(cfg.vocab_size))
         self.register_buffer(
             "rope_cache",
             precompute_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_base, torch.device("cpu")),
@@ -201,22 +238,32 @@ def num_parameters(self, exclude_embeddings: bool = False) -> int:
     def forward(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)
+        ve = self.value_embed(idx) if self.value_embeddings else None
         if self.unet_skip:
             n = len(self.blocks); half = n // 2; enc = []
             for i, block in enumerate(self.blocks):
+                lam = self.ve_lambda[i] if self.value_embeddings else None
                 if i < half:
-                    x = block(x, self.rope_cache); enc.append(x)
+                    x = block(x, self.rope_cache, ve, lam); enc.append(x)
                 else:
                     x = x + self.skip_gate[i - half] * enc[n - 1 - i]
-                    x = block(x, self.rope_cache)
+                    x = block(x, self.rope_cache, ve, lam)
         else:
-            for block in self.blocks:
-                x = block(x, self.rope_cache)
+            for i, block in enumerate(self.blocks):
+                lam = self.ve_lambda[i] if self.value_embeddings else None
+                x = block(x, self.rope_cache, ve, lam)
         x = self.final_norm(x)
         if self.lm_head is None:
             logits = F.linear(x, self.tok_embed.weight)
         else:
             logits = self.lm_head(x)
+        # Apply the learned readout temperature before the soft-cap. exp(0)=1 at
+        # init so this is an identity at step 0; the optimizer then sets the
+        # readout gain independently of the tied embedding norm.
+        # Per-vocab readout gain (per-token temperature) + bias (per-token prior),
+        # then the global temperature, all before the soft-cap. Identity at init.
+        logits = logits * torch.exp(self.readout_gain) + self.readout_bias
+        logits = logits * torch.exp(self.logit_scale)
         cap = getattr(self.cfg, "logit_softcap", 0.0)  # recipe-v4: logit soft-cap
         if cap and cap > 0:
             logits = cap * torch.tanh(logits / cap)