model-def (GPT-ref Qwen3-arch) for same-harness loader

train_gpt_ref.py

@@ -60,6 +60,25 @@ def get_args():
                    help="adamw (default, backward-compat) | muon (Newton-Schulz ortho dla 2D-weights + AdamW dla reszty)")
     p.add_argument("--muon-lr", type=float, default=0.02,
                    help="peak LR dla Muon (macierzowe params); AdamW-aux uzywa --lr. Muon skalowany ta sama cosine-schedule co AdamW przez lr_mult=muon_lr/lr")
+    p.add_argument("--norm", choices=["layernorm", "rmsnorm"], default="layernorm",
+                   help="layernorm (default, backward-compat) | rmsnorm (Qwen3-style, fp32-compute)")
+    p.add_argument("--norm-eps", type=float, default=1e-6)
+    p.add_argument("--pos", choices=["learned", "rope"], default="learned",
+                   help="learned (default) | rope (parameter-free RoPE na q,k; usuwa learned pos-embedding)")
+    p.add_argument("--rope-theta", type=float, default=100000.0,
+                   help="RoPE theta; top-3-board (JugnuLM/GPT-X2) uzywaja 100000 (nie-default-10K)")
+    p.add_argument("--ffn", choices=["gelu", "swiglu"], default="gelu",
+                   help="gelu (default 4x MLP) | swiglu (Qwen3 gated-MLP, hidden=--ffn-mult*d)")
+    p.add_argument("--ffn-mult", type=float, default=2.667,
+                   help="mnoznik hidden dla swiglu (~param-parity z 4x-gelu przy 8/3)")
+    p.add_argument("--value-residual", action="store_true",
+                   help="ResFormer value-residuals: v_l += lambda_l*v0 (lambda init 0), ARC-targeted")
+    p.add_argument("--qk-norm", action="store_true",
+                   help="Qwen3 QK-Norm: RMSNorm per-head na Q,K przed-attention (stabilnosc z Muon/high-LR)")
+    p.add_argument("--events-jsonl", default=None,
+                   help="jesli podane: emituj events.jsonl (fabryka-track sidecar-format: update/evaluation/checkpoint/end)")
+    p.add_argument("--compile", action="store_true",
+                   help="torch.compile model (2-3x throughput; state_dict zapisywany bez _orig_mod prefix via raw_model)")
     return p.parse_args()
 
 
@@ -179,38 +198,103 @@ def build_optimizer(a, model):
     return MuonWithAuxAdam(muon, adamw)
 
 
+class RMSNorm(nn.Module):
+    """Qwen3-style RMSNorm (fp32-compute dla stabilnosci). 1D weight -> AdamW w split-Muon."""
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+def make_norm(d, cfg):
+    return RMSNorm(d, cfg.norm_eps) if cfg.norm == "rmsnorm" else nn.LayerNorm(d)
+
+
+def apply_rope(x, base=100000.0):
+    """Parameter-free RoPE na [B,H,T,D] (interleaved-conv, port z qwen_model.py). Train==eval
+    MUSZA uzywac tej samej konwencji (self-contained eval -> spojne)."""
+    _, _, T, dim = x.shape
+    pos = torch.arange(T, device=x.device, dtype=torch.float32)
+    freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
+    ang = torch.outer(pos, freq)
+    cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None]
+    even, odd = x[..., ::2], x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class SwiGLU(nn.Module):
+    """Qwen3 gated-MLP: down(silu(gate(x))*up(x)). 3x 2D bez-bias -> wszystkie do Muon."""
+    def __init__(self, d, hidden):
+        super().__init__()
+        self.gate = nn.Linear(d, hidden, bias=False)
+        self.up = nn.Linear(d, hidden, bias=False)
+        self.down = nn.Linear(hidden, d, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
 class Block(nn.Module):
-    def __init__(self, d, nh, block):
+    def __init__(self, d, nh, block, cfg, is_first=False):
         super().__init__()
-        self.ln1 = nn.LayerNorm(d)
-        self.ln2 = nn.LayerNorm(d)
+        self.ln1 = make_norm(d, cfg)
+        self.ln2 = make_norm(d, cfg)
         self.qkv = nn.Linear(d, 3 * d)
         self.proj = nn.Linear(d, d)
-        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
+        if cfg.ffn == "swiglu":
+            self.mlp = SwiGLU(d, int(round(cfg.ffn_mult * d)))
+        else:
+            self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
         self.nh = nh
         self.d = d
-
-    def forward(self, x):
+        self.cfg = cfg
+        self.is_first = is_first
+        if cfg.value_residual and not is_first:
+            self.vr_lambda = nn.Parameter(torch.zeros(1))
+        if cfg.qk_norm:
+            hd = d // nh
+            self.q_norm = RMSNorm(hd, cfg.norm_eps)
+            self.k_norm = RMSNorm(hd, cfg.norm_eps)
+
+    def forward(self, x, v0=None):
         B, T, D = x.size()
         h = self.ln1(x)
         q, k, v = self.qkv(h).split(self.d, dim=2)
-        q = q.view(B, T, self.nh, D // self.nh).transpose(1, 2)
-        k = k.view(B, T, self.nh, D // self.nh).transpose(1, 2)
-        v = v.view(B, T, self.nh, D // self.nh).transpose(1, 2)
+        hd = D // self.nh
+        q = q.view(B, T, self.nh, hd).transpose(1, 2)
+        k = k.view(B, T, self.nh, hd).transpose(1, 2)
+        v = v.view(B, T, self.nh, hd).transpose(1, 2)
+        if self.cfg.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if self.cfg.pos == "rope":
+            q = apply_rope(q, self.cfg.rope_theta)
+            k = apply_rope(k, self.cfg.rope_theta)
+        if self.cfg.value_residual:
+            if self.is_first:
+                v0 = v
+            else:
+                v = v + self.vr_lambda * v0
         y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
         y = y.transpose(1, 2).contiguous().view(B, T, D)
         x = x + self.proj(y)
         x = x + self.mlp(self.ln2(x))
-        return x
+        return x, v0
 
 
 class GPT(nn.Module):
-    def __init__(self, vocab, n_layer, n_embd, n_head, block):
+    def __init__(self, vocab, n_layer, n_embd, n_head, block, cfg):
         super().__init__()
+        self.cfg = cfg
         self.tok = nn.Embedding(vocab, n_embd)
-        self.pos = nn.Embedding(block, n_embd)
-        self.blocks = nn.ModuleList([Block(n_embd, n_head, block) for _ in range(n_layer)])
-        self.lnf = nn.LayerNorm(n_embd)
+        self.use_rope = cfg.pos == "rope"
+        if not self.use_rope:
+            self.pos = nn.Embedding(block, n_embd)
+        self.blocks = nn.ModuleList([Block(n_embd, n_head, block, cfg, is_first=(i == 0)) for i in range(n_layer)])
+        self.lnf = make_norm(n_embd, cfg)
         self.head = nn.Linear(n_embd, vocab, bias=False)
         self.head.weight = self.tok.weight  # tie
         self.block = block
@@ -226,10 +310,13 @@ class GPT(nn.Module):
 
     def forward(self, idx, targets=None):
         B, T = idx.size()
-        pos = torch.arange(T, device=idx.device)
-        x = self.tok(idx) + self.pos(pos)[None]
+        x = self.tok(idx)
+        if not self.use_rope:
+            pos = torch.arange(T, device=idx.device)
+            x = x + self.pos(pos)[None]
+        v0 = None
         for b in self.blocks:
-            x = b(x)
+            x, v0 = b(x, v0)
         logits = self.head(self.lnf(x))
         loss = None
         if targets is not None:
@@ -251,6 +338,18 @@ def main():
         with open(log_path, "a", encoding="utf-8") as f:
             f.write(line + "\n")
 
+    events_path = a.events_jsonl
+
+    def emit(kind, step, metrics=None, **extra):
+        if not events_path:
+            return
+        rec = {"kind": kind, "updates": int(step), "tokens": int(step) * a.block * a.batch}
+        if metrics:
+            rec["metrics"] = metrics
+        rec.update(extra)
+        with open(events_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(rec) + "\n")
+
     torch.manual_seed(a.seed)
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
@@ -337,16 +436,19 @@ def main():
             except queue.Empty:
                 pass
 
-    model = GPT(a.vocab, a.n_layer, a.n_embd, a.n_head, a.block).to(device)
-    nparam = sum(p.numel() for p in model.parameters())
+    raw_model = GPT(a.vocab, a.n_layer, a.n_embd, a.n_head, a.block, a).to(device)
+    nparam = sum(p.numel() for p in raw_model.parameters())
     log(f"model GPT-ref: {nparam/1e6:.1f}M param (L{a.n_layer} d{a.n_embd} h{a.n_head})")
-    opt = build_optimizer(a, model)
+    opt = build_optimizer(a, raw_model)
     log(f"optimizer={a.optimizer}" + (f" muon_lr={a.muon_lr} (mult={a.muon_lr/a.lr:.1f}x)" if a.optimizer == "muon" else ""))
+    model = torch.compile(raw_model) if a.compile else raw_model
+    if a.compile:
+        log("torch.compile enabled")
 
     start_step = 0
     if a.resume and os.path.exists(ckpt_path):
         ck = torch.load(ckpt_path, map_location=device)
-        model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); start_step = ck["step"]
+        raw_model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); start_step = ck["step"]
         log(f"RESUME @ {start_step}")
 
     def lr_at(s):
@@ -395,14 +497,31 @@ def main():
             log(f"step {step+1}/{a.steps} loss {running/a.log_every:.4f} lr {lr:.2e} gnorm {float(gn):.2f} {tok_s:,.0f} tok/s peakVRAM {mem:.1f}GB")
             with open(metrics_path, "a", encoding="utf-8") as f:
                 f.write(json.dumps({"step": step+1, "loss": running/a.log_every, "lr": lr, "tok_s": tok_s}) + "\n")
+            emit("update", step + 1, metrics={"loss": running / a.log_every, "tokens_per_second": tok_s,
+                                              "gradient_norm": float(gn), "learning_rate": lr})
             running = 0.0; t0 = time.time()
         if (step + 1) % a.eval_every == 0:
-            log(f"  >> VAL loss {eval_val():.4f} @ {step+1}")
+            vloss = eval_val()
+            log(f"  >> VAL loss {vloss:.4f} @ {step+1}")
+            emit("evaluation", step + 1, metrics={"loss": vloss})
         if (step + 1) % a.ckpt_every == 0 and not a.synthetic:
-            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step+1}, ckpt_path)
+            torch.save({"model": raw_model.state_dict(), "opt": opt.state_dict(), "step": step + 1,
+                        "config": {"vocab": a.vocab, "n_layer": a.n_layer, "n_embd": a.n_embd,
+                                   "n_head": a.n_head, "block": a.block, "norm": a.norm,
+                                   "norm_eps": a.norm_eps, "pos": a.pos, "rope_theta": a.rope_theta,
+                                   "ffn": a.ffn, "ffn_mult": a.ffn_mult, "value_residual": a.value_residual,
+                                   "qk_norm": a.qk_norm}}, ckpt_path)
             log(f"ckpt @ {step+1}")
+            if a.events_jsonl:
+                import hashlib as _hl
+                _h = _hl.sha256()
+                with open(ckpt_path, "rb") as _cf:
+                    for _chunk in iter(lambda: _cf.read(1 << 20), b""):
+                        _h.update(_chunk)
+                emit("checkpoint", step + 1, sha256=_h.hexdigest())
     prefetcher.close()
     log("DONE")
+    emit("end", a.steps, status="completed")
 
 
 if __name__ == "__main__":