Upload scripts/fast4gpu_boost_v4.py with huggingface_hub

scripts/fast4gpu_boost_v4.py

@@ -0,0 +1,333 @@
+#!/usr/bin/env python3
+"""Fractus-1B boost trainer v4 — AR gap.
+
+v3 kernels/data/ckpt semantics, PLUS the losses that actually target free-run:
+
+  - P0 routing on (phase carry + per-token MoE + switch LB)
+  - SS_RATE default 1.0 (every step), SS_PROB ramps 0.2 → 0.5
+  - anti-repeat: λ · mean log p(class = last input token)
+  - unique@40 greedy probe (no ban) on a timer — THIS is the go/no-go, not ema_tf
+
+Env (v3 plus):
+  SS_RATE=1.0  SS_PROB_START=0.2  SS_PROB_END=0.5  SS_RAMP_TOKENS=50000000
+  REPEAT_COEF=0.1  PROBE_EVERY=200
+  P0=1  (set 0 to ablate routing surgery)
+
+One GPU:
+  CUDA_VISIBLE_DEVICES=$i GPU_ID=$i START_TOKEN=<manifest> \\
+    BATCH=8 CE_CHUNK=2048 FRACTUS_ATTN_IMPL=chunked BLOCK_CKPT=1 \\
+    python -u scripts/fast4gpu_boost_v4.py
+
+Smoke ONE gpu 1–2h and read unique@40 before touching the other seven.
+"""
+from __future__ import annotations
+
+import os
+import sys
+import time
+import json
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+os.chdir(ROOT)
+
+os.environ.setdefault("FRACTUS_ATTN_IMPL", os.environ.get("FRACTUS_ATTN_IMPL", "cumsum"))
+from fractus.continuous_engine import ContinuousThoughtEngine
+from fractus.generate_aligned import unique40_probe
+from fractus.train.ar_loss import ss_prob_at
+from fractus.train.v4_step import v4_forward_losses, v4_ss_pass, should_ss, snapshot_carry, restore_carry
+
+GPU = int(os.environ.get("GPU_ID", "0"))
+LB_COEF = float(os.environ.get("LB_COEF", "0.02"))
+GATE_TEMP = float(os.environ.get("GATE_TEMP", "2.5"))
+LR = float(os.environ.get("LR", "7e-4"))
+EMA_BETA = float(os.environ.get("EMA_BETA", "0.98"))
+SS_RATE = float(os.environ.get("SS_RATE", "1.0"))
+SS_PROB_START = float(os.environ.get("SS_PROB_START", "0.2"))
+SS_PROB_END = float(os.environ.get("SS_PROB_END", "0.5"))
+SS_RAMP_TOKENS = int(os.environ.get("SS_RAMP_TOKENS", "50000000"))
+REPEAT_COEF = float(os.environ.get("REPEAT_COEF", "0.1"))
+PROBE_EVERY = int(os.environ.get("PROBE_EVERY", "200"))
+P0 = os.environ.get("P0", "1") == "1"
+B = int(os.environ.get("BATCH", "4"))
+SEQ = int(os.environ.get("SEQ", "128"))
+CE_CHUNK = int(os.environ.get("CE_CHUNK", "2048"))
+ACCUM = max(1, int(os.environ.get("ACCUM", "1")))
+BLOCK_CKPT = os.environ.get("BLOCK_CKPT", "0") == "1"
+USE_COMPILE = os.environ.get("COMPILE", "0") == "1"
+ATTN_IMPL = os.environ.get("FRACTUS_ATTN_IMPL", "cumsum")
+
+TARGET = dict(
+    d_model=1280, n_heads=20, d_head=64, n_levels=2,
+    n_oscillators=16, coupling_rank=8, n_experts=128, top_k=2,
+    expert_d_ff=2048, siren_rank=64, n_layers=16,
+)
+
+torch.manual_seed(42 + GPU)
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+    device = torch.device("cuda:0")
+    autocast = lambda: torch.autocast("cuda", dtype=torch.bfloat16)
+else:
+    device = torch.device("cpu")
+    from contextlib import nullcontext
+    autocast = nullcontext
+
+default_merged = ROOT / "checkpoints" / "FRACTUS_1B_STAGE2_MERGED.pt"
+default_gpu = ROOT / "checkpoints" / f"fractus_1b_gpu{GPU}.pt"
+CKPT_IN = Path(os.environ.get("CKPT_IN", str(default_gpu if default_gpu.exists() else default_merged)))
+CKPT_OUT = Path(os.environ.get("CKPT_OUT", str(default_gpu)))
+SHARD = Path(os.environ.get("SHARD", str(ROOT / "data" / f"shard_gpu{GPU}.npy")))
+MANIFEST_OUT = Path(os.environ.get(
+    "MANIFEST_OUT", str(CKPT_OUT.parent / f"RESUME_MANIFEST_gpu{GPU}.json")))
+
+print(
+    f"GPU {GPU}: BOOSTv4 B={B} SEQ={SEQ} LR={LR} SS_RATE={SS_RATE} "
+    f"ss_prob={SS_PROB_START}->{SS_PROB_END} repeat={REPEAT_COEF} P0={P0} "
+    f"attn={ATTN_IMPL} ce_chunk={CE_CHUNK} block_ckpt={BLOCK_CKPT}",
+    flush=True,
+)
+print(f"GPU {GPU}: load {CKPT_IN}", flush=True)
+
+ck = torch.load(CKPT_IN, map_location="cpu", weights_only=False)
+sd = ck.get("model_state", ck)
+clean = {(k[10:] if k.startswith("_orig_mod.") else k): v for k, v in sd.items()}
+
+eng = ContinuousThoughtEngine(vocab_size=50257, **TARGET)
+own = eng.state_dict()
+loaded = 0
+for k, v in clean.items():
+    if k in own and own[k].shape == v.shape:
+        own[k] = v
+        loaded += 1
+    elif (
+        k in own and v.dim() >= 1 and own[k].dim() >= 1
+        and v.shape[0] > own[k].shape[0] and v.shape[1:] == own[k].shape[1:]
+    ):
+        own[k] = v[: own[k].shape[0]].contiguous()
+        loaded += 1
+eng.load_state_dict(own, strict=False)
+print(f"GPU {GPU}: loaded_tensors={loaded}", flush=True)
+
+eng.set_p0_routing(
+    P0, lb_mode=("switch_topk" if P0 else "soft_var"),
+)
+with torch.no_grad():
+    for blk in eng.blocks:
+        if hasattr(blk, "moe") and hasattr(blk.moe, "temperature"):
+            blk.moe.temperature = GATE_TEMP
+
+eng = eng.to(device)
+eng.reset_thought(B)
+
+if USE_COMPILE:
+    try:
+        eng = torch.compile(eng)
+        print(f"GPU {GPU}: torch.compile ON", flush=True)
+    except Exception as e:
+        print(f"GPU {GPU}: compile skip: {e}", flush=True)
+
+payload = lambda: (eng._orig_mod if hasattr(eng, "_orig_mod") else eng)
+
+opt = torch.optim.SGD(eng.parameters(), lr=LR, momentum=0.9)
+
+if not SHARD.exists():
+    raise FileNotFoundError(f"Shard not found: {SHARD}")
+
+import numpy as np
+if not str(SHARD).endswith(".npy"):
+    raise FileNotFoundError(f"v4 expects .npy int32 shards, got {SHARD}")
+shard_mm = np.load(str(SHARD), mmap_mode="r")
+shard_len = int(shard_mm.shape[0])
+print(f"GPU {GPU}: memmap shard {SHARD} len={shard_len:,} dtype={shard_mm.dtype}", flush=True)
+
+step_tokens = B * SEQ
+
+
+def fetch(start: int, count: int) -> torch.Tensor:
+    view = np.asarray(shard_mm[start : start + count])
+    return torch.from_numpy(view).to(torch.int64, non_blocking=True).to(device)
+
+
+start_token = int(os.environ.get("START_TOKEN", "0"))
+start_token = (start_token // step_tokens) * step_tokens
+print(f"GPU {GPU}: RESUME start_token={start_token} step={step_tokens} shard_len={shard_len:,}",
+      flush=True)
+
+t0 = time.time()
+ema_tf = ema_ss = ema_rep = None
+n = 0
+tok_sess = 0
+pending_backward = False
+CKPT_OUT.parent.mkdir(parents=True, exist_ok=True)
+
+
+def save_ckpt(tokens_done: int):
+    tmp = CKPT_OUT.with_suffix(CKPT_OUT.suffix + ".tmp")
+    torch.save(
+        {
+            "model_state": payload().state_dict(),
+            "config": {
+                **TARGET, "gpu": GPU, "boost_v4": True, "batch": B, "lr": LR,
+                "ss_rate": SS_RATE, "repeat_coef": REPEAT_COEF, "p0": P0,
+                "tokens_processed": tokens_done,
+            },
+        },
+        tmp,
+    )
+    os.replace(tmp, CKPT_OUT)
+    mtmp = MANIFEST_OUT.with_suffix(MANIFEST_OUT.suffix + ".tmp")
+    mtmp.write_text(json.dumps({
+        "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "gpu": GPU, "trainer": "fast4gpu_boost_v4",
+        "attn_impl": ATTN_IMPL, "p0": P0, "batch": B,
+        "tokens_processed": tokens_done, "start_token_next": tokens_done,
+        "shard": str(SHARD), "shard_len": shard_len, "ckpt": CKPT_OUT.name,
+    }, indent=1))
+    os.replace(mtmp, MANIFEST_OUT)
+    print(f"GPU {GPU}: saved [boostv4] -> {CKPT_OUT} @ {tokens_done:,} tok", flush=True)
+
+
+def run_probe(tokens_done: int):
+    eng_p = payload()
+    thought = eng_p.thought_state
+    carries = [(b.attn_S.clone(), b.attn_z.clone(), b.kuramoto_phases.clone())
+               for b in eng_p.blocks]
+    B_live = thought.shape[0]
+    probe = unique40_probe(eng_p, max_new=40, mode="prefix")
+    probe_c = unique40_probe(eng_p, max_new=40, mode="carry")
+    # restore live train state
+    eng_p.thought_state = thought
+    for b, (S, z, ph) in zip(eng_p.blocks, carries):
+        b.attn_S, b.attn_z, b.kuramoto_phases = S, z, ph
+    eng_p.reset_thought(B_live)  # batch may have been set to 1
+    # reset_thought zeros — put live carries back
+    eng_p.thought_state = thought
+    for b, (S, z, ph) in zip(eng_p.blocks, carries):
+        b.attn_S, b.attn_z, b.kuramoto_phases = S, z, ph
+    gate = "GO" if probe["gate_go"] else "NO-GO"
+    print(
+        f"GPU {GPU}: UNIQUE@40 PREFIX {gate} mean_u={probe['mean_unique']:.1f} "
+        f"echo={probe['mean_echo_frac']:.2f} | "
+        f"CARRY u={probe_c['mean_unique']:.1f} echo={probe_c['mean_echo_frac']:.2f} "
+        f"@ {tokens_done:,} tok",
+        flush=True,
+    )
+    for r in probe["rows"]:
+        print(
+            f"  {r['prompt']!r}: unique={r['unique']} echo={r['echo_frac']:.2f} head={r['head']}",
+            flush=True,
+        )
+    rs = eng_p.routing_stats()
+    print(
+        f"  routing alive={rs.get('alive_experts')} "
+        f"H={rs.get('dispatch_entropy', float('nan')):.3f} "
+        f"max_frac={rs.get('max_frac', float('nan')):.3f}",
+        flush=True,
+    )
+    return probe
+
+
+for start in range(start_token, shard_len - step_tokens - SEQ - 1, step_tokens):
+    block = fetch(start, step_tokens + 1)
+    chunk = block[:step_tokens].view(B, SEQ).long()
+    target = block[1:].view(B, SEQ)
+    tokens_now = start + step_tokens
+    ss_prob = ss_prob_at(tokens_now, SS_PROB_START, SS_PROB_END, SS_RAMP_TOKENS)
+    carry_snap = snapshot_carry(eng)
+
+    with autocast():
+        loss, extras = v4_forward_losses(
+            payload(), chunk, target,
+            lb_coef=LB_COEF, repeat_coef=REPEAT_COEF,
+            ce_chunk=CE_CHUNK, block_ckpt=BLOCK_CKPT,
+        )
+    ce_tf, lb, rep, h = extras["ce"], extras["lb"], extras["repeat"], extras["h"]
+
+    ss_fired = False
+    ce_ss_v = None
+    if should_ss(SS_RATE):
+        ss_fired = True
+
+    if ACCUM == 1:
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0)
+        opt.step()
+        opt.zero_grad(set_to_none=True)
+        if ss_fired:
+            restore_carry(eng, carry_snap)
+            with autocast():
+                loss2, ce_ss = v4_ss_pass(
+                    payload(), chunk, target, h.detach(),
+                    ss_prob=ss_prob, lb_coef=LB_COEF,
+                    ce_chunk=CE_CHUNK, block_ckpt=False,
+                )
+            loss2.backward()
+            torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0)
+            opt.step()
+            opt.zero_grad(set_to_none=True)
+            ce_ss_v = float(ce_ss.item())
+            ema_ss = ce_ss_v if ema_ss is None else EMA_BETA * ema_ss + (1 - EMA_BETA) * ce_ss_v
+    else:
+        (loss / ACCUM).backward()
+        if ss_fired:
+            restore_carry(eng, carry_snap)
+            with autocast():
+                loss2, ce_ss = v4_ss_pass(
+                    payload(), chunk, target, h.detach(),
+                    ss_prob=ss_prob, lb_coef=LB_COEF,
+                    ce_chunk=CE_CHUNK, block_ckpt=False,
+                )
+            (loss2 / ACCUM).backward()
+            ce_ss_v = float(ce_ss.item())
+            ema_ss = ce_ss_v if ema_ss is None else EMA_BETA * ema_ss + (1 - EMA_BETA) * ce_ss_v
+        pending_backward = True
+
+    tf_v = float(ce_tf.detach().item())
+    lb_v = float(lb.detach().item()) if torch.is_tensor(lb) else float(lb)
+    rp_v = float(rep.detach().item())
+    ema_tf = tf_v if ema_tf is None else EMA_BETA * ema_tf + (1 - EMA_BETA) * tf_v
+    ema_rep = rp_v if ema_rep is None else EMA_BETA * ema_rep + (1 - EMA_BETA) * rp_v
+
+    n += 1
+    tok_sess += step_tokens
+
+    if ACCUM > 1 and n % ACCUM == 0:
+        torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0)
+        opt.step()
+        opt.zero_grad(set_to_none=True)
+        pending_backward = False
+
+    if n % 40 == 0:
+        tps = tok_sess / max(time.time() - t0, 1e-6)
+        extra = f" ss={ce_ss_v:.3f} ema_ss={ema_ss:.3f}" if ce_ss_v is not None else ""
+        try:
+            mem_s = f" mem={torch.cuda.max_memory_allocated() / 1e9:.1f}GB"
+        except Exception:
+            mem_s = ""
+        print(
+            f"GPU {GPU}: {tokens_now:>12,} tf={tf_v:.3f} ema_tf={ema_tf:.3f}{extra} "
+            f"rep={rp_v:.3f} ema_rep={ema_rep:.3f} lb={lb_v:.3f} "
+            f"ssp={ss_prob:.2f} {tps:.0f} tok/s{mem_s} [boostv4]",
+            flush=True,
+        )
+
+    if PROBE_EVERY > 0 and n % PROBE_EVERY == 0:
+        run_probe(tokens_now)
+
+    # All 8 GPUs save, staggered so they never write 8x4.3G at once.
+    if n % 4000 == (GPU * 500) % 4000:
+        save_ckpt(tokens_now)
+
+if pending_backward:
+    torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0)
+    opt.step()
+    opt.zero_grad(set_to_none=True)
+
+save_ckpt(start_token + n * step_tokens)
+print(f"GPU {GPU}: DONE", flush=True)