Upload rung8_grouped_g4.py with huggingface_hub

rung8_grouped_g4.py

@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+"""rung8_grouped_g4.py — coarse-grained grouped expert routing on Gemma-4.
+
+Each MLP's D_FFN neurons are partitioned into K groups (cluster assignments
+loaded from analyze_activation_groups.py output). Per token, top-K_active
+groups are selected; all neurons within a selected group are activated.
+
+Vs rung7's per-neuron mask:
+  - Coarser-grained → potentially compute/memory-bandwidth efficient
+    (skip whole groups, not individual neurons)
+  - Same per-token density target but expressed via group selection
+
+Usage mirrors rung7 + adds:
+    --K_groups 64 --group_assignments_dir logs/groups --group_tag s25
+"""
+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 GroupedMaskedMLP(nn.Module):
+    """Top-K_active group routing. Group score = max |gate_act| within group.
+
+    Forward:
+        gate_act = gelu(gate_proj(x))                     # [B, T, D_FFN]
+        for each token, group_g_score = max over j in group g of |gate_act[j]|
+        select top K_active groups → mask all neurons in selected groups
+        h = gate_act * up_proj(x) * mask
+        out = down_proj(h)
+    """
+    def __init__(self, base_mlp, K_groups, K_active, group_assignments, 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_groups = int(K_groups)
+        self.K_active = int(K_active)
+        # group_assignments: [D_FFN] long, in [0, K_groups)
+        self.register_buffer("group_assignments", group_assignments.long())
+        # Build group → neuron map (one-hot) for vectorized scatter
+        # neuron_in_group[d, g] = 1 if neuron d is in group g, else 0  shape [D_FFN, K_groups]
+        D = group_assignments.shape[0]
+        nig = torch.zeros(D, K_groups)
+        nig.scatter_(1, group_assignments.long().unsqueeze(1), 1.0)
+        self.register_buffer("neuron_in_group", nig)
+        self.tau = 0.01  # used only for sigmoid relaxation; defaults hard
+
+    def forward(self, x):
+        gate_act = F.gelu(self.gate_proj(x), approximate="tanh")        # [B, T, D_FFN]
+        up_act   = self.up_proj(x)
+        gate_abs = gate_act.abs().to(torch.float32)
+        B, T, D = gate_abs.shape
+        BT = B * T
+        flat = gate_abs.view(BT, D)
+        # Group score = max within group  (vectorized via scatter_reduce)
+        group_score = torch.full((BT, self.K_groups), -float("inf"),
+                                  device=gate_act.device, dtype=torch.float32)
+        group_score.scatter_reduce_(1, self.group_assignments.unsqueeze(0).expand(BT, -1),
+                                    flat, reduce="amax", include_self=False)
+        # Top-K_active groups per token
+        top_vals, top_idx = group_score.topk(self.K_active, dim=-1)  # [BT, K_active]
+        # Sigmoid relaxation around the K_active-th largest group-score:
+        # neuron_score[d] = group_score[group_of_d]
+        # mask = sigmoid((neuron_score - kth_thr) / tau)
+        kth_thr = top_vals[..., -1:]                                 # [BT, 1]
+        neuron_score = group_score.gather(1, self.group_assignments.unsqueeze(0).expand(BT, -1).long())
+        mask_flat = torch.sigmoid((neuron_score - kth_thr) / max(self.tau, 1e-3))
+        mask = mask_flat.view(B, T, D)
+        h = gate_act * up_act * mask.to(gate_act.dtype)
+        return self.down_proj(h)
+
+
+def install_grouped(model, K_groups, K_active, group_assignments_dir, group_tag, freeze_base=False):
+    mlp_modules = []
+    for i in range(N_LAYERS):
+        d_ffn = _d_ffn_at(i)
+        path = f"{group_assignments_dir}/{group_tag}_K{K_groups}_layer{i}.pt"
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Missing group assignments: {path}")
+        assignments = torch.load(path, map_location="cpu", weights_only=False)
+        if assignments.numel() != d_ffn:
+            raise ValueError(f"Layer {i}: assignments size {assignments.numel()} != D_FFN {d_ffn}")
+        new_mlp = GroupedMaskedMLP(model.layers[i].mlp,
+                                    K_groups=K_groups, K_active=K_active,
+                                    group_assignments=assignments,
+                                    freeze_base=freeze_base)
+        new_mlp = new_mlp.to(DEVICE)
+        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="G1")
+    parser.add_argument("--K_groups", type=int, required=True)
+    parser.add_argument("--density", type=float, default=0.25,
+                        help="Target per-token density: K_active = density * K_groups (rounded)")
+    parser.add_argument("--K_active", type=int, default=0,
+                        help="Override K_active explicitly (else computed from density)")
+    parser.add_argument("--group_assignments_dir", default="logs/groups")
+    parser.add_argument("--group_tag", required=True)
+    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")
+    parser.add_argument("--freeze_embeddings", action="store_true")
+    parser.add_argument("--use_lora", action="store_true")
+    parser.add_argument("--lora_targets", type=str, default="")
+    parser.add_argument("--lora_rank", type=int, default=16)
+    parser.add_argument("--lora_alpha", type=float, default=16.0)
+    parser.add_argument("--tau", type=float, default=0.01)
+    parser.add_argument("--max_steps", type=int, default=500)
+    parser.add_argument("--lr", type=float, default=1e-5)
+    parser.add_argument("--shuffle_seed", type=int, default=42)
+    parser.add_argument("--save_every", type=int, default=200)
+    parser.add_argument("--eval_every", type=int, default=100)
+    parser.add_argument("--eval_max_seqs", type=int, default=200)
+    parser.add_argument("--calib_path", required=True)
+    parser.add_argument("--eval_calib_path", required=True)
+    parser.add_argument("--load_checkpoint", type=str, default="")
+    parser.add_argument("--save_checkpoint", type=str, default="")
+    args = parser.parse_args()
+
+    K_active = args.K_active if args.K_active > 0 else max(1, round(args.K_groups * args.density))
+    print(f"=== Rung 8 Grouped — phase={args.phase} ===")
+    print(f"  K_groups={args.K_groups}  K_active={K_active}  effective_density={K_active/args.K_groups:.3f}")
+    print(f"  loss={args.loss}  unfreeze_base={args.unfreeze_base}  use_lora={args.use_lora}")
+    print(f"  group_assignments_dir={args.group_assignments_dir}  group_tag={args.group_tag}")
+    if args.load_checkpoint:
+        print(f"  load_checkpoint={args.load_checkpoint}")
+
+    # Teacher only needed for KL; for CE we skip
+    teacher_needed = (args.loss == "kl")
+    if teacher_needed:
+        print("Loading teacher & student...")
+        teacher, tokenizer = load_gemma4()
+        teacher.eval()
+        for p in teacher.parameters(): p.requires_grad_(False)
+        student, _ = load_gemma4()
+    else:
+        print("Loading student only (CE loss; teacher skipped)...")
+        teacher = None
+        student, tokenizer = load_gemma4()
+
+    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)
+
+    freeze_base_in_mlp = not args.unfreeze_base
+    mlp_modules = install_grouped(student,
+                                  K_groups=args.K_groups, K_active=K_active,
+                                  group_assignments_dir=args.group_assignments_dir,
+                                  group_tag=args.group_tag,
+                                  freeze_base=freeze_base_in_mlp)
+    print(f"  Installed GroupedMaskedMLP on {N_LAYERS} layers")
+
+    if args.load_checkpoint:
+        print(f"  Loading checkpoint from {args.load_checkpoint}...")
+        ckpt = torch.load(args.load_checkpoint, map_location=DEVICE, weights_only=False)
+        # The loaded ckpt has GateMaskedMLP state (no group_assignments, no neuron_in_group).
+        # Load with strict=False — only base proj weights match, group buffers stay as we set them.
+        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} Linear modules")
+
+    if args.use_lora:
+        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} ({n_lora} modules, {n_lora_p/1e6:.2f}M)")
+
+    if args.load_checkpoint:
+        missing2, unexp2 = student.load_state_dict(ckpt["student_state"], strict=False)
+        print(f"  re-loaded after wrappers: missing={len(missing2)} unexpected={len(unexp2)}")
+
+    for m in mlp_modules: m.tau = args.tau
+
+    n_train = sum(p.numel() for p in student.parameters() if p.requires_grad)
+    print(f"  Trainable params: {n_train/1e6:.3f}M")
+
+    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: {args.calib_path}\n  Eval: {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)} sequences")
+    g = torch.Generator(); g.manual_seed(args.shuffle_seed)
+    loader = torch.utils.data.DataLoader(seqs, BATCH, shuffle=True, generator=g)
+    loader_iter = iter(loader)
+
+    step = 0
+    t0 = time.time()
+    curve = []
+    optimizer.zero_grad()
+
+    while step < args.max_steps:
+        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)
+
+        if teacher is not None 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=2.0, mask=mask)
+        else:
+            loss = ce_loss(s_logits, labels)
+
+        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 if args.eval_max_seqs > 0 else None)
+            elapsed = time.time() - t0
+            print(f"  step={step:5d}  loss={loss.item():.4f}  ppl={ppl:.4f}  t={elapsed:.0f}s", flush=True)
+            curve.append({"step": step, "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)
+
+    final_ppl = eval_ppl(student, tokenizer, calib_path=args.eval_calib_path,
+                         max_seqs=args.eval_max_seqs if args.eval_max_seqs > 0 else None)
+    print(f"\n=== Final PPL: {final_ppl:.4f} ===")
+
+    out = {"phase": args.phase, "config": vars(args), "final_ppl": final_ppl, "ppl_curve": curve}
+    os.makedirs("logs", exist_ok=True)
+    with open(f"logs/rung8_{args.phase}_results.json", "w") as f: json.dump(out, f, indent=2)
+
+    if args.save_checkpoint:
+        torch.save({"student_state": student.state_dict(),
+                    "config": vars(args), "final_ppl": final_ppl}, args.save_checkpoint)
+        print(f"Saved {args.save_checkpoint}")
+        interim = args.save_checkpoint.replace(".pt", "_intermediate.pt")
+        if os.path.exists(interim): os.remove(interim)
+
+
+if __name__ == "__main__":
+    main()