[submit] 0bfc0e455cbf — val by hotkey 5Fbh5xeEbXa7…

configs/run270.json

@@ -0,0 +1,43 @@
+{
+ "vocab_size": 50257,
+ "dim": 1024,
+ "n_layers": 16,
+ "n_heads": 16,
+ "head_dim": 64,
+ "ffn_mult": 2.6667,
+ "max_seq_len": 1024,
+ "seq_len": 512,
+ "batch_size": 1024,
+ "micro_batch_size": 128,
+ "total_steps": 5050,
+ "warmup_steps": 700,
+ "max_lr": 0.003,
+ "min_lr": 1e-05,
+ "weight_decay": 0.1,
+ "beta1": 0.9,
+ "beta2": 0.95,
+ "grad_clip": 1.0,
+ "schedule": "wsd",
+ "stable_frac": 0.55,
+ "decay_frac": 0.45,
+ "decay_curve": "1-sqrt",
+ "embed_lr": 0.009,
+ "embed_optimizer": "adamw",
+ "optimizer": "muon",
+ "muon_lr": 0.04,
+ "muon_momentum": 0.95,
+ "muon_ns_steps": 5,
+ "muon_weight_decay": 0.05,
+ "muon_momentum_start": null,
+ "data_seed": 270270,
+ "init_seed": 270270,
+ "use_bf16": true,
+ "compile": true,
+ "log_every": 100,
+ "ema_decay": 0.0,
+ "deterministic": false,
+ "compile_mode": "max-autotune",
+ "logit_z_coef": 0.0,
+ "data_sampling": "replacement",
+ "_comment": "run270: EXACT replication of current king max_mfu (uid124, val_bpb 1.023). muon_lr0.04 + embed_lr0.009 + warmup700 + wsd0.55/0.45 + NO ema/zloss/perm/momentum. batch1024/5050st. Target ~1.023 = king level (huge jump from our ~1.071)."
+}

data/dataset.py

@@ -41,6 +41,7 @@ def __init__(
         base_dir: Path | str,
         seq_len: int,
         seed: int,
+        sampling: str = "replacement",
     ):
         from .manifest import DataManifest, verify_manifest
 
@@ -54,6 +55,10 @@ def __init__(
         self._total = int(self._cum[-1])
         self.seq_len = seq_len
         self.seed = seed
+        self.sampling = sampling
+        self._n_windows = self._total // (self.seq_len + 1)
+        self._perm_epoch = -1
+        self._perm = None
         if self._total < seq_len + 1:
             raise ValueError(f"not enough tokens ({self._total}) for seq_len {seq_len}")
 
@@ -103,14 +108,38 @@ def get(self, step: int) -> tuple[torch.Tensor, torch.Tensor]:
         ids = torch.from_numpy(chunk.astype(np.int64))
         return ids[:-1], ids[1:]
 
+    def _perm_for_epoch(self, epoch: int) -> np.ndarray:
+        """Deterministic seed-keyed permutation of the non-overlapping window
+        indices for a given epoch. Re-derivable by the validator on audit from
+        (seed, epoch)."""
+        if self._perm_epoch != epoch:
+            self._perm = np.random.default_rng(
+                np.array([self.seed, 0xE9EC, epoch], dtype=np.uint64)
+            ).permutation(self._n_windows)
+            self._perm_epoch = epoch
+        return self._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)
+        if self.sampling == "permutation":
+            # Without-replacement full-coverage epoching: deal non-overlapping
+            # windows in a seed-keyed shuffled order. Deterministic from
+            # (seed, seq_len, batch_size, step) -> validator re-derivable.
+            stride = self.seq_len + 1
+            starts = np.empty(batch_size, dtype=np.int64)
+            for b in range(batch_size):
+                p = step * batch_size + b
+                epoch = p // self._n_windows
+                within = p % self._n_windows
+                win = int(self._perm_for_epoch(epoch)[within])
+                starts[b] = win * stride
+        else:
+            rng = np.random.default_rng(np.array([self.seed, step], dtype=np.uint64))
+            starts = rng.integers(0, self._total, size=batch_size)
         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):

model/_v4skip.py

@@ -38,6 +38,7 @@ 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)
+    logit_z_coef: float = 0.0     # z-loss: coef * logsumexp(logits)^2 (0 = off)
 
 
 def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
@@ -227,6 +228,9 @@ def forward(self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None) ->
                 targets.view(-1),
                 ignore_index=-100,
             )
+            z_coef = getattr(self.cfg, "logit_z_coef", 0.0)
+            if z_coef:
+                loss = loss + z_coef * (torch.logsumexp(logits, dim=-1).float() ** 2).mean()
         return logits, loss