Apply Tier 2 research findings to train_gpt.py

- Warmdown 1200 → 3500 (proven by both our research and openai#2 leaderboard entry)
- Muon weight decay WD=0.04 (proven at both Tier 1 and Tier 2 scales)
- Adam embedding weight decay WD=0.01 (proven to stack with Muon WD)
- LeakyReLU(0.5) activation (used by openai#1 leaderboard entry)

Made-with: Cursor

Author: adityamhn

train_gpt.py

@@ -52,7 +52,7 @@ class Hyperparameters:
 
     # Training length.
     iterations = int(os.environ.get("ITERATIONS", 20000))
-    warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200))
+    warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 3500))
     warmup_steps = int(os.environ.get("WARMUP_STEPS", 20))
     train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288))
     train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024))
@@ -84,6 +84,8 @@ class Hyperparameters:
     beta1 = float(os.environ.get("BETA1", 0.9))
     beta2 = float(os.environ.get("BETA2", 0.95))
     adam_eps = float(os.environ.get("ADAM_EPS", 1e-8))
+    muon_wd = float(os.environ.get("MUON_WD", 0.04))
+    adam_embed_wd = float(os.environ.get("ADAM_EMBED_WD", 0.01))
     grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0))
 
 # -----------------------------
@@ -110,10 +112,10 @@ def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -
 
 
 class Muon(torch.optim.Optimizer):
-    def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True):
+    def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True, wd: float = 0.0):
         super().__init__(
             params,
-            dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov),
+            dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov, wd=wd),
         )
 
     @torch.no_grad()
@@ -135,6 +137,7 @@ def step(self, closure=None):
             momentum = group["momentum"]
             backend_steps = group["backend_steps"]
             nesterov = group["nesterov"]
+            wd = group.get("wd", 0.0)
 
             total_params = sum(int(p.numel()) for p in params)
             updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16)
@@ -162,6 +165,8 @@ def step(self, closure=None):
             curr = 0
             for p in params:
                 g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype)
+                if wd > 0:
+                    p.mul_(1.0 - lr * wd)
                 p.add_(g, alpha=-lr)
                 curr += p.numel()
 
@@ -613,8 +618,8 @@ def __init__(self, dim: int, mlp_mult: int):
         self.proj._zero_init = True
 
     def forward(self, x: Tensor) -> Tensor:
-        x = torch.relu(self.fc(x))
-        return self.proj(x.square())
+        x = F.leaky_relu(self.fc(x), negative_slope=0.5)
+        return self.proj(x * x)
 
 
 class Block(nn.Module):
@@ -873,6 +878,7 @@ def log0(msg: str, console: bool = True) -> None:
         lr=args.matrix_lr,
         momentum=args.muon_momentum,
         backend_steps=args.muon_backend_steps,
+        wd=args.muon_wd,
     )
     for group in optimizer_muon.param_groups:
         group["base_lr"] = args.matrix_lr
@@ -1031,6 +1037,9 @@ def lr_mul(step: int, elapsed_ms: float) -> float:
             torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm)
         for opt in optimizers:
             opt.step()
+        if args.adam_embed_wd > 0:
+            with torch.no_grad():
+                base_model.tok_emb.weight.mul_(1.0 - token_lr * scale * args.adam_embed_wd)
         zero_grad_all()
 
         step += 1