Upload rung7_swiglu_g4.py with huggingface_hub

rung7_swiglu_g4.py

@@ -0,0 +1,364 @@
+#!/usr/bin/env python3
+"""rung7_swiglu_g4.py — per-token top-K gate mask for Gemma-4 (no expert structure).
+
+Each neuron is its own "expert". Per-token mask = top-K by |gate_act| magnitude,
+relaxed via sigmoid((|gate| - kth_threshold) / τ) for differentiability.
+At τ→0 it converges to hard top-K. No router, no MECE partition, no A matrix.
+
+Mirrors rung6_moe_g4.py CLI/training loop but installs GateMaskedMLP instead of
+MoEMLP. Reuses load_seqs / eval_ppl / wrap_int4 / get_tau / kl_loss / ce_loss.
+
+Strong prior: Gemma-3 Design 6 (this exact mechanism) hit PPL 7.26 vs base 7.89
+= 0.92× base. Best result on Gemma-3. Never tried on Gemma-4.
+"""
+import argparse, json, math, os, time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+from gemma4_hf import load_gemma4, DEVICE, N_LAYERS
+from rung6_moe_g4 import (
+    Int4QuantLinear, wrap_int4, apply_int4_inplace,
+    LoRALinear, wrap_lora,
+    load_seqs, eval_ppl, kl_loss, ce_loss, get_tau,
+    _d_ffn_at,
+    MAX_SEQ_LEN, BATCH, LR, BASELINE_PPL, CLEAN_PPL,
+)
+
+
+class GateMaskedMLP(nn.Module):
+    """Per-token top-K mask on |gate_act|. τ-annealed sigmoid relaxation.
+
+    Forward:
+        gate_act = gelu(gate_proj(x))
+        threshold[t] = kth-largest |gate_act[t]| (k = k_keep)
+        mask[t,j] = sigmoid((|gate_act[t,j]| - threshold[t]) / τ)
+        h = gate_act * up_proj(x) * mask
+        out = down_proj(h)
+    """
+    def __init__(self, base_mlp, k_keep, freeze_base=False):
+        super().__init__()
+        self.gate_proj = base_mlp.gate_proj
+        self.up_proj   = base_mlp.up_proj
+        self.down_proj = base_mlp.down_proj
+        if freeze_base:
+            for p in self.gate_proj.parameters(): p.requires_grad_(False)
+            for p in self.up_proj.parameters():   p.requires_grad_(False)
+            for p in self.down_proj.parameters(): p.requires_grad_(False)
+        self.k_keep = int(k_keep)
+        self.tau = 1.0  # set externally each step
+
+    def forward(self, x):
+        gate_raw = self.gate_proj(x)
+        gate_act = F.gelu(gate_raw, approximate="tanh")    # [B, T, D_FFN]
+        up_act   = self.up_proj(x)
+        gate_abs = gate_act.abs().to(torch.float32)
+        # Per-token kth-largest threshold (non-differentiable wrt selection,
+        # but mask values around the threshold ARE differentiable via sigmoid).
+        threshold = gate_abs.topk(self.k_keep, dim=-1).values[..., -1:]   # [B, T, 1]
+        mask = torch.sigmoid((gate_abs - threshold) / max(self.tau, 1e-3))  # [B, T, D_FFN]
+        h = gate_act * up_act * mask.to(gate_act.dtype)
+        return self.down_proj(h)
+
+
+def install_gate_mask(model, density, freeze_base=False):
+    mlp_modules = []
+    for i in range(N_LAYERS):
+        d_ffn = _d_ffn_at(i)
+        k_keep = max(1, int(round(d_ffn * density)))
+        new_mlp = GateMaskedMLP(model.layers[i].mlp, k_keep=k_keep, freeze_base=freeze_base)
+        model.layers[i].mlp = new_mlp
+        mlp_modules.append(new_mlp)
+    return mlp_modules
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--phase", type=str, default="S1")
+    parser.add_argument("--density", type=float, default=0.75,
+                        help="Fraction of MLP neurons to keep per token (e.g. 0.75 ≈ Aconst4 density)")
+    parser.add_argument("--loss", choices=["kl", "ce"], default="ce")
+    parser.add_argument("--int4_qat", action="store_true")
+    parser.add_argument("--int4_group_size", type=int, default=32)
+    parser.add_argument("--unfreeze_base", action="store_true",
+                        help="Train base weights (gate/up/down + attn). Default freezes them.")
+    parser.add_argument("--freeze_embeddings", action="store_true")
+    parser.add_argument("--gate_only_train", action="store_true",
+                        help="Override: freeze entire model, only gate_proj across all layers trains. "
+                             "Tests whether the gate alone can route + adapt.")
+    parser.add_argument("--gate_lora_train", action="store_true",
+                        help="Override: freeze entire model, train gate_proj + LoRA adapters on "
+                             "up_proj/down_proj. Tests whether LoRA on the masked weights "
+                             "compensates for aggressive masking at low density.")
+    parser.add_argument("--lora_targets", type=str, default="",
+                        help="Comma-separated substrings of Linear names to wrap with LoRA. "
+                             "Default (empty) uses the rung6 wrap_lora default. For gate_lora_train, "
+                             "set to 'up_proj,down_proj' to skip gate_proj.")
+    parser.add_argument("--use_lora", action="store_true")
+    parser.add_argument("--lora_rank", type=int, default=16)
+    parser.add_argument("--lora_alpha", type=float, default=16.0)
+    parser.add_argument("--tau_start", type=float, default=1.0)
+    parser.add_argument("--tau_end", type=float, default=0.01)
+    parser.add_argument("--tau_hold_frac", type=float, default=0.2)
+    parser.add_argument("--max_steps", type=int, default=10000)
+    parser.add_argument("--lr", type=float, default=LR)
+    parser.add_argument("--main_kl_temp", type=float, default=2.0)
+    parser.add_argument("--shuffle_seed", type=int, default=42)
+    parser.add_argument("--data_skip", type=int, default=0)
+    parser.add_argument("--save_every", type=int, default=2500)
+    parser.add_argument("--eval_every", type=int, default=2500)
+    parser.add_argument("--eval_max_seqs", type=int, default=0,
+                        help="Cap eval to first N sequences (0 = no cap, current behavior). "
+                             "Set e.g. 200 to keep mid-training evals fast; the final "
+                             "post-training eval line still runs full unless capped here.")
+    parser.add_argument("--calib_path", type=str, required=True)
+    parser.add_argument("--eval_calib_path", type=str, required=True)
+    parser.add_argument("--load_checkpoint", type=str, default="")
+    parser.add_argument("--save_checkpoint", type=str, default="")
+    parser.add_argument("--diverse_calib_path", type=str, default="")
+    parser.add_argument("--diverse_every_n", type=int, default=4)
+    parser.add_argument("--kl_base_lambda", type=float, default=0.5)
+    parser.add_argument("--kl_base_temp", type=float, default=2.0)
+    parser.add_argument("--w_drift_lambda", type=float, default=0.0)
+    args = parser.parse_args()
+
+    print(f"=== Rung 7 SWIGLU gate-mask — phase={args.phase} ===")
+    print(f"  density={args.density:.2f}  loss={args.loss}")
+    print(f"  tau: {args.tau_start} → {args.tau_end} over {args.max_steps} steps "
+          f"(hold last {args.tau_hold_frac*100:.0f}%)")
+    print(f"  unfreeze_base={args.unfreeze_base}  freeze_embeddings={args.freeze_embeddings}")
+    print(f"  int4_qat={args.int4_qat}  use_lora={args.use_lora}")
+    if args.load_checkpoint:
+        print(f"  load_checkpoint={args.load_checkpoint}")
+    if args.save_checkpoint:
+        print(f"  save_checkpoint={args.save_checkpoint}")
+
+    # Teacher is only needed if main loss is KL, a diverse-corpus KL-to-base
+    # regularizer is configured, or the W-drift penalty needs the teacher's
+    # snapshot. With --loss ce and no diverse / drift, the teacher forward is
+    # dead compute and ~9GB of dead weight; skip loading it.
+    teacher_ever_needed = (
+        args.loss == "kl"
+        or bool(args.diverse_calib_path)
+        or args.w_drift_lambda > 0
+    )
+
+    if teacher_ever_needed:
+        print("Loading teacher & student on cuda...")
+        teacher, tokenizer = load_gemma4()
+        teacher.eval()
+        for p in teacher.parameters(): p.requires_grad_(False)
+    else:
+        print("Loading student only on cuda (teacher not needed: --loss ce, no diverse calib)...")
+        # Tokenizer comes from a lightweight load; reuse the student load below.
+        teacher = None
+
+    student, tokenizer_s = load_gemma4()
+    if teacher is None:
+        tokenizer = tokenizer_s
+    if args.freeze_embeddings:
+        for n, p in student.named_parameters():
+            if "embed_tokens" in n or "lm_head" in n:
+                p.requires_grad_(False)
+        n_frozen = sum(p.numel() for n, p in student.named_parameters()
+                       if ("embed_tokens" in n or "lm_head" in n))
+        print(f"  Froze embeddings: {n_frozen/1e9:.2f}B params")
+
+    freeze_base_in_mlp = not args.unfreeze_base
+    mlp_modules = install_gate_mask(student, density=args.density,
+                                    freeze_base=freeze_base_in_mlp)
+    print(f"  Installed GateMaskedMLP on {N_LAYERS} layers; "
+          f"k_keep range = [{min(m.k_keep for m in mlp_modules)}, {max(m.k_keep for m in mlp_modules)}]")
+
+    if args.load_checkpoint:
+        print(f"  Loading checkpoint from {args.load_checkpoint}...")
+        ckpt = torch.load(args.load_checkpoint, map_location=DEVICE, weights_only=False)
+        missing, unexpected = student.load_state_dict(ckpt["student_state"], strict=False)
+        print(f"    missing={len(missing)} unexpected={len(unexpected)}")
+
+    if args.int4_qat:
+        Int4QuantLinear._group_size = args.int4_group_size
+        n_wrap = wrap_int4(student)
+        print(f"  Int4 QAT: wrapped {n_wrap} nn.Linear modules (group_size={args.int4_group_size})")
+
+    if args.use_lora or args.gate_lora_train:
+        if args.lora_targets:
+            targets = tuple(t.strip() for t in args.lora_targets.split(",") if t.strip())
+            n_lora, n_lora_p = wrap_lora(student, rank=args.lora_rank,
+                                         alpha=args.lora_alpha, target_substrings=targets)
+        else:
+            n_lora, n_lora_p = wrap_lora(student, rank=args.lora_rank, alpha=args.lora_alpha)
+        print(f"  LoRA: rank={args.lora_rank} alpha={args.lora_alpha} "
+              f"({n_lora} modules, {n_lora_p/1e6:.2f}M params)")
+
+    if args.load_checkpoint:
+        # Re-load after wrappers (LoRA / int4 add new keys)
+        missing2, unexp2 = student.load_state_dict(ckpt["student_state"], strict=False)
+        print(f"  re-loaded after wrappers: missing={len(missing2)} unexpected={len(unexp2)}")
+
+    if args.gate_only_train:
+        for p in student.parameters():
+            p.requires_grad_(False)
+        for n, p in student.named_parameters():
+            if "gate_proj" in n:
+                p.requires_grad_(True)
+        n_gate = sum(p.numel() for n, p in student.named_parameters() if p.requires_grad)
+        print(f"  --gate_only_train override: only gate_proj trains ({n_gate/1e6:.2f}M params)")
+
+    if args.gate_lora_train:
+        for p in student.parameters():
+            p.requires_grad_(False)
+        for n, p in student.named_parameters():
+            # Gate projection trains directly (router specialization).
+            # LoRA adapters on up/down_proj train (compensate for aggressive masking).
+            # NOTE: a Linear named "..mlp.gate_proj" wrapped by LoRA becomes "..mlp.gate_proj.base"
+            # — to avoid ambiguity we use --lora_targets up_proj,down_proj so gate isn't wrapped.
+            if "gate_proj" in n or "lora_a" in n or "lora_b" in n:
+                p.requires_grad_(True)
+        n_train = sum(p.numel() for n, p in student.named_parameters() if p.requires_grad)
+        n_gate_p = sum(p.numel() for n, p in student.named_parameters()
+                       if p.requires_grad and "gate_proj" in n)
+        n_lora_p = sum(p.numel() for n, p in student.named_parameters()
+                       if p.requires_grad and ("lora_a" in n or "lora_b" in n))
+        print(f"  --gate_lora_train override: gate_proj + LoRA adapters train "
+              f"({n_train/1e6:.2f}M total — gate {n_gate_p/1e6:.2f}M + LoRA {n_lora_p/1e6:.2f}M)")
+
+    n_train = sum(p.numel() for p in student.parameters() if p.requires_grad)
+    print(f"  Trainable params: {n_train/1e6:.3f}M (no router; mask is non-parametric)")
+
+    optimizer = AdamW([p for p in student.parameters() if p.requires_grad],
+                      lr=args.lr, weight_decay=0.01)
+    scheduler = CosineAnnealingLR(optimizer, T_max=args.max_steps, eta_min=args.lr * 0.1)
+
+    print(f"  Train data: {args.calib_path}")
+    print(f"  Eval  data: {args.eval_calib_path}")
+    train_split = "all" if args.calib_path != args.eval_calib_path else "train"
+    seqs = load_seqs(tokenizer, train_split, calib_path=args.calib_path)
+    print(f"  Loaded {len(seqs)} train sequences of {MAX_SEQ_LEN} tokens "
+          f"= {len(seqs)*MAX_SEQ_LEN/1e6:.2f}M tokens (split={train_split})")
+    g = torch.Generator(); g.manual_seed(args.shuffle_seed)
+    loader = torch.utils.data.DataLoader(seqs, BATCH, shuffle=True, generator=g)
+    loader_iter = iter(loader)
+    if args.data_skip > 0:
+        for _ in range(args.data_skip):
+            try: next(loader_iter)
+            except StopIteration:
+                loader_iter = iter(loader); next(loader_iter)
+        print(f"  Skipped first {args.data_skip} samples")
+
+    diverse_loader_iter = None
+    if args.diverse_calib_path:
+        print(f"  Diverse corpus (KL-to-base): {args.diverse_calib_path}")
+        diverse_seqs = load_seqs(tokenizer, "all", calib_path=args.diverse_calib_path, raw_text=True)
+        print(f"    {len(diverse_seqs)} sequences, every {args.diverse_every_n} steps, "
+              f"λ={args.kl_base_lambda}, T={args.kl_base_temp}")
+        diverse_loader = torch.utils.data.DataLoader(diverse_seqs, BATCH, shuffle=True)
+        diverse_loader_iter = iter(diverse_loader)
+
+    teacher_param_map = None
+    if args.w_drift_lambda > 0:
+        print(f"  W-drift penalty active: λ={args.w_drift_lambda}")
+        teacher_param_map = {n: p.detach() for n, p in teacher.named_parameters()}
+
+    step = 0
+    t0 = time.time()
+    curve = []
+    optimizer.zero_grad()
+
+    while step < args.max_steps:
+        tau = get_tau(step, args.max_steps, args.tau_start, args.tau_end,
+                      hold_frac=args.tau_hold_frac)
+        for m in mlp_modules: m.tau = tau
+
+        try: batch = next(loader_iter)
+        except StopIteration:
+            loader_iter = iter(loader); batch = next(loader_iter)
+        input_ids = batch["input_ids"].to(DEVICE)
+        labels    = batch["labels"].to(DEVICE)
+
+        # Teacher forward is needed only if the main loss is KL or if a diverse
+        # KL-to-base regularizer is firing this step. With --loss ce and no
+        # --diverse_calib_path, the teacher logits are computed-then-discarded —
+        # ~half the per-step compute (4.65B params) for nothing. Short-circuit
+        # in that case. Numerically equivalent to dropping a dead branch; no
+        # change to the gradient that reaches the student.
+        diverse_active_this_step = (
+            diverse_loader_iter is not None and step % args.diverse_every_n == 0
+        )
+        teacher_needed = (args.loss == "kl") or diverse_active_this_step
+
+        if teacher_needed and args.loss == "kl":
+            with torch.no_grad():
+                t_logits = teacher(input_ids)
+        s_logits = student(input_ids)
+
+        if args.loss == "kl":
+            mask = (labels != -100)
+            loss = kl_loss(s_logits, t_logits, temp=args.main_kl_temp, mask=mask)
+        else:
+            loss = ce_loss(s_logits, labels)
+
+        if diverse_loader_iter is not None and step % args.diverse_every_n == 0:
+            try: dbatch = next(diverse_loader_iter)
+            except StopIteration:
+                diverse_loader_iter = iter(diverse_loader); dbatch = next(diverse_loader_iter)
+            d_ids = dbatch["input_ids"].to(DEVICE)
+            with torch.no_grad():
+                t_d_logits = teacher(d_ids)
+            s_d_logits = student(d_ids)
+            d_kl = kl_loss(s_d_logits, t_d_logits, temp=args.kl_base_temp)
+            loss = loss + args.kl_base_lambda * d_kl
+
+        if teacher_param_map is not None:
+            drift = 0.0
+            for n, p in student.named_parameters():
+                if not p.requires_grad: continue
+                if n in teacher_param_map:
+                    drift = drift + (p - teacher_param_map[n]).pow(2).sum()
+            loss = loss + args.w_drift_lambda * drift
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_([p for p in student.parameters() if p.requires_grad], 1.0)
+        optimizer.step()
+        optimizer.zero_grad()
+        scheduler.step()
+
+        step += 1
+        if step % args.eval_every == 0 or step == args.max_steps:
+            ppl = eval_ppl(student, tokenizer, calib_path=args.eval_calib_path,
+                           max_seqs=(args.eval_max_seqs or None))
+            elapsed = time.time() - t0
+            print(f"  step={step:5d}  tau={tau:.4f}  loss={loss.item():.4f}  "
+                  f"ppl={ppl:.4f}  t={elapsed:.0f}s")
+            curve.append({"step": step, "tau": tau, "loss": float(loss.item()), "ppl": float(ppl)})
+            if args.save_checkpoint and step % args.save_every == 0 and step < args.max_steps:
+                interim = args.save_checkpoint.replace(".pt", "_intermediate.pt")
+                torch.save({"student_state": student.state_dict(),
+                            "config": vars(args), "step": step, "ppl": ppl}, interim)
+                print(f"  [intermediate] overwrote {interim} (step {step})")
+
+    final_ppl = eval_ppl(student, tokenizer, calib_path=args.eval_calib_path,
+                         max_seqs=(args.eval_max_seqs or None))
+    print(f"\n=== Final PPL (tau={args.tau_end}): {final_ppl:.4f} ===")
+
+    out = {"phase": args.phase, "config": vars(args), "final_ppl": final_ppl,
+           "ppl_curve": curve}
+    os.makedirs("logs", exist_ok=True)
+    out_path = f"logs/rung7_swiglu_{args.phase}_results.json"
+    with open(out_path, "w") as f: json.dump(out, f, indent=2)
+    print(f"Saved to {out_path}")
+
+    if args.save_checkpoint:
+        torch.save({"student_state": student.state_dict(),
+                    "config": vars(args), "final_ppl": final_ppl}, args.save_checkpoint)
+        print(f"Saved checkpoint to {args.save_checkpoint}")
+        interim = args.save_checkpoint.replace(".pt", "_intermediate.pt")
+        if os.path.exists(interim):
+            os.remove(interim)
+            print(f"Removed {interim}")
+
+
+if __name__ == "__main__":
+    main()