aide-moe-aux-loss-stability git:20260917.4308e77C

Immutable. This exact content is served forever at /api/v1/blob/afbaf6757b32e785.

# MoE Aux-Loss Stability — Flagship Twin-Mind Pretrain (local140a50)

## When to use
- Any flagship twin-mind MoE run showing instability: eval CE regression while train loss falls, pre-clip grad-norm spikes (1e5 -> 1e7), LBL climbing toward double digits, router z loss pinned ~2-4 (or reported-4 with router lse ~5.5 per layer, NOT ~20), xpop -> ~0 (router collapse).
- BEFORE changing ANY aux-loss weight, LR/clip/warmup, or adding/removing a loss on this model.
- Before every flagship relaunch (patch B: resume-LR override) and before any D4 SFT marker retrain.

## Verified facts (2026-09-10, this machine, this repo)
- BUG: `cipher_arch.py:760` = `layer_z = (lse ** 2).mean() * 0.01` — internal `* 0.01`. Trainer then multiplies z by 0.001 (`train_v6.py:505`). Stacked effective weight = **1e-5 = 100x below ST-MoE canonical 0.001**.
- Consequence measured on chunk1c: router logits grew under-penalized (eval z 4.2 was already ×0.01 → true Σ_layer mean(lse^2) ≈ 420 summed over 14 MoE layers → per-layer ~30 → logsumexp ≈ 5.5, NOT ~20; the "z ~2-4" printed at training IS the ×0.01-scaled number), LBL climbed 2.1 -> 10.8, pre-clip grad-norm reached 1.1M -> 11.2M, eval CE regressed 7.7749@1000 -> 8.3567@1750 -> 10.1639@2000 (best eval 7.7031). Clip 1.0 prevented NaN (zero NaNs/aborts across chunk1c).
- SCALE CHANGE ON THE FIX (do not misread): removing the internal `*0.01` makes the LOGGED z metric read ~100x LARGER (~214 train / ~420 eval at the current logit magnitude), and the z-loss TERM in `loss = ce + 0.01*lbl + 0.001*z + 0.01*xpop` jumps from ~0.004 to ~0.42 (100x stronger regularizer -> the actual fix mechanism). The rein-in target: as router logits shrink toward a healthy compact range, Σ mean(lse^2) should DECLINE from ~420 toward well under ~100-150 (per-layer mean(lse^2) toward ~4-9, lse ~2-3) — the z column TREND downward is the signal, never compare raw post-fix z to the pre-fix ×0.01-scaled numbers.
- Canonical reference implementations already in-repo and correct: `verify_v6.py:58` (weight 0.001*z) and `src_colonymind/colony_moe.py:78` (`torch.logsumexp(logits, dim=-1).square().mean()`, no multiplier). ONLY `cipher_arch.py:760` carries the extra 0.01.
- Population split: `local140a50` = `n_routed=48, n_pop=2` => pop_size=24 (Sheldon-pop = w1p/routed indices 0..23, Spock-pop = 24..47). (The 600a150/1200a200 presets use 62 => 31/31; the code reads `pop_size = n_routed // cfg["n_pop"]` so it never hardcodes.) xpop cosine = one global per-layer cosine over `w1p[:pop_size].flatten() @ w1p[pop_size:].flatten()` (`cipher_arch.py:514-516`) — 14 MoE layers, single cosine each. SAFE. NOT a per-pair 576/2304-term sum.
- Locked training shape (probe-verified 2026-09-10): **B=2, S=512 = 3.12GB / 1322 tok/s**. B=4 S=512 = 5.27GB OVER; B=2 S=1024 = 5.50GB OVER. Do not change without a probe.
- Host-RAM: NOT a dataset/mmap leak (probe-verified 2026-09-10). Cause = Windows CUDA host-shadow commit (~5-7GB persistent) + full-state ckpt-save transient (~1.7GB). Collapse signature: FreePhysicalMemory -> ~1GB, trainer PrivateMemorySize64 -> 13.7GB, WS 6.3GB but RSS 0.92GB (paged out), tok/s 1322 -> 170-227 while GPU 92-100%. Doctrine: keep FreePhysicalMemory >= 5GB, kill stale llama-server/python before launch.
- Auto-resume (plain relaunch, no --base-ckpt) preserves the cosine schedule; log says `resuming from step_..._NNNNNN.pt` + `resumed at step N best=...`. `--base-ckpt` = model-only + FRESH optimizer/cosine = full restart with the model weights — NEVER use for continuation.
- VERIFIED FIX IS LIVE (chunk2a 2026-09-10, resume 2000 -> stop-at 2500, 21.3 min, 0 NaNs): after both patches the canonical z-loss reined the router in — step z 427.6@2010 -> 295.2@2500, EVAL z 289.1, LBL 10.5-11.3 quiet (not climbing), popS 0.497/popK 0.503 balanced, dead 0.000 experts, eval ce best 7.7095@2500 (+0.006 flat vs 7.7031 pre-fix best — the 5e-4->3e-4 LR cut regressed nothing). xpop measured 0.0001-0.018 (never a dominator). The ckpt round-trip: `step_local140a50_002500.pt` stores keys `[step, model, optimizer, scheduler, loss, cfg]` — NO top-level lr/base_lrs keys; optimizer lr/initial_lr and scheduler base_lrs/_last_lr/_step_count carry the state (verified reload: optimizer lr 2.61e-4, initial_lr 3e-4, base_lrs [3e-4], step_count 2501). Generation probe on the step ckpt: finite logits, top5 p 0.0315->0.0128, entropy 8.0, vocab cv 8.88 = peaked, not collapsed.
- **OVERSHOOT FAILURE (chunk2b 2026-09-10, resume 2500 -> 4000): the effective-canonical 1e-3 OVER-reined the router and killed it.** Measured: z collapsed 295.2@2500 -> 26.2@2750 -> 0.082@3000 -> 0.028@3250 -> 0.030@4000 (logits crushed to lse ~0.05/layer, ~50x past the lse 2-3 target). Best eval ce was 7.6066@3250 (hit at the moment logits were just-compacted but still discriminative), then regressed 7.6406@3500 / **8.2422@3750 / 8.2717@4000** as the router went uniform -> quorum picks random -> twin-mind debate hollow -> 48-expert average acts like a dense net -> overfits train shards (train ce ~5.5 vs eval 8.2). grad_norm collapsed 1e5-1e6 -> ~0.5-780. WHY: at THIS scale (B=2 S=512 = only 1024 token/step task signal) the ST-MoE canonical 1e-3 behaves like the ~1e-2 regime of ST-MoE Fig.7 (c_z >= 1e-2 shrinks z to near zero) — the canonical coefficient is tuned for giant-scale where the router sees billions of diverse tokens/step, not 1024. **z-weight is a PER-SCALE tune, not a universal 0.001.** Bisect decision: 1e-4, verified with a real chunk.
- **Crash bug on that run (fixed):** a `print("\u26a0 WARNING")` hit `UnicodeEncodeError: 'charmap' codec can't encode '\u26a0'` under Windows cp1252-redirected stdout AFTER the step-4000 eval but BEFORE the step-4000 ckpt write -> `step_local140a50_004000.pt` is MISSING. Rule: **NO non-ASCII characters in trainer prints** (cp1252 redirect), fixed to `[WARN]`. Metrics CSV survived (rows through 4000).
- Recovery constraints learned: `best_local140a50.pt` stores only `{step, model, cfg}` — NO optimizer/scheduler, NOT resumable. **The last full-state ckpt with a LIVE router is `step_local140a50_002500.pt`** (`_003000`/`_003500` are post-collapse). train_v6.py now has `--resume-from <path>` for full-state resume from a specific ckpt (auto-resume would glob the LATEST step_* = `_003500`, wrong) and `--z-weight <float>` (loss line `ce + 0.01*lbl + args.z_weight*z + 0.01*xpop`, default 1e-4). A copy `recovery_step002500.pt` guards the recovery point against prune.

## Canonical recipe (research-grounded, primary sources)
- z-loss: ST-MoE `arXiv:2202.08906` Sec 3.4 Eq.5 — mean per token of (logsumexp router logits)^2, coefficient c_z = **0.001** (ST-MoE Fig. 7 best; OLMoE `arXiv:2409.02060` uses 0.001). No internal extra multipliers. **2026-09-10 CORRECTION: 0.001 is the GIANT-scale coefficient; at this 146.5M scale (B=2 S=512, 1024 tok/step) it behaves like ST-MoE's ~1e-2 (dead-uniform router). Tune per scale — for local140a50, bisect at 1e-4 (`--z-weight 1e-4`, the trainer default). Healthy band to verify empirically: z col in post-fix scale should DECLINE toward ~5-150 (per-layer mean(lse^2) ~1-10, lse ~1-3) and NOT blow past it; dead = <1; hot = >150. Verify with a real chunk before locking.
- LBL: switch-style f_i*g_i summed, multiplied by n_routed, coefficient **0.01** (Switch Transformer; balanced healthy value ~1.0). Scope = the EFFECTIVE full batch. In `train_v6.py` batch=2 IS the full batch (no accumulation), bincount over 1024 tokens — correct, keep.
- Grad clip: **1.0** canonical (OLMoE). NEVER raise to 5.0 to silence spikes — that hides the real problem. MiniMax anomalous-grad-norm precedent: the fix is LOWER LR (2e-4 -> 1.3e-4), not a bigger clip. queen-bee 150M dense used clip 5.0 — different regime, do not transplant.
- LR on instability: cut 5e-4 -> 3e-4 (logits still ~5.5e-3 at 3e-4 — a mild cut, not a stall). Fresh runs: warmup 500-1000 recommended by research; the existing 10k schedule keeps 200 (already past on resume, no action needed this run).
- Telemetry (professional path): clip fraction (% steps where pre-clip grad-norm > clip), per-population utilization (Sheldon 0..23 vs Spock 24..47 share of routed tokens for local140a50), dead-expert fraction (experts with zero routed tokens during eval), and a 3-consecutive-eval-regression warning (do NOT auto-kill on warning).

## What to do (exact patch list, then relaunch)
1. `cipher_arch.py:760` — remove the `* 0.01`: `layer_z = (lse ** 2).mean()`.
2. `train_v6.py` resume block — after `optimizer.load_state_dict` + `scheduler.load_state_dict`: if `--lr` differs from the saved base, override `param_groups[i]['lr']` and `['initial_lr']`, and `scheduler.base_lrs` + `scheduler._last_lr`, to args.lr. Log the override. (Without this, `--lr 3e-4` on an auto-resume is SILENTLY IGNORED — PyTorch restores lr/initial_lr/base_lrs/_last_lr from the ckpt.)
3. `train_v6.py` — add telemetry: (a) clip-hit counter at the clip_grad_norm_ call; (b) per-eval per-pop utilization + dead-expert fraction from the eval metas; (c) 3-consecutive-eval-regression warning line.
4. Relaunch as a PLAIN auto-resume (no --fresh, no --base-ckpt) from the latest `step_local140a50_*.pt` with `--lr 3e-4`, chunked `--stop-at`. To jump back to an earlier full-state ckpt (e.g. after a collapse) use **`--resume-from <path>`** instead (auto-resume globs the LATEST step_* — it will NOT find a recovery copy named `recovery_*.pt`). z is set via **`--z-weight`** (default 1e-4). Shape locked B=2 S=512. Sun reset: eval/ckpt cadence unchanged (250/500).
5. VERIFY WITH REAL TESTS, NOT SMOKE (verification-complete doctrine — the operator's hard rule: NO smoke tests, only ACTUAL task verification): the evidence is the REAL chunk run from the REAL checkpoint — resume log line, LR-override line, real metrics rows (lr col ~2.6-2.75e-04 on the cosine phase, NOT the restored 4.60e-04), real z rows in the canonical scale (~300-430 post-fix train-scale, NOT ~2.1 and NOT xpop ~0.02), real EVAL lines with popS/popK/dead/clip_frac (ce must not regress far past the pre-fix best 7.7031), the real step ckpt file written, a reload round-trip (step + optimizer lr/initial_lr + scheduler base_lrs/_step_count preserved), and a real generation probe proving non-degenerate output. The 500-step chunk from step 2000 -> 2500 with evals at 2250/2500 and a step ckpt at 2500 IS the verified proof (chunk2a, 2026-09-10: all 7 gates passed, 0 NaNs, 21.3 min). Only then continue the schedule.

## What NOT to do
- Raise grad clip to 5.0. Delete LBL. Add entropy-max. Stack aux multipliers (fi all used weights at one place — `train_v6.py:505`: `ce + 0.01*lbl + 0.001*z + 0.01*xpop`).
- Restart from scratch — resume the trajectory (2000 healthy steps preserved in chunk1c / `step_local140a50_002000.pt`).
- Run `--fresh` (wipes resume + truncates the metrics CSV history).
- Change B/S from locked 2/512. Run `num_workers>0` on Windows. Train FP16/BF16 (Pascal: FP32 only).
- Auto-kill or auto-react to a warning — warnings are signals, not stop conditions. Never kill/pause/restart a live run without human approval (hard rule R2).
- Compare post-fix z rows to pre-fix z rows numerically WITHOUT the scale correction (post-fix z reads ~100x larger by design after removing the internal ×0.01). Track the z TREND within the new scale — it must DECLINE as the router logits are reined in.
- Use z-weight 1e-3 at THIS scale (chunk2b proved it kills the router: uniform -> dense-like -> eval ce 8.27 while train ce 5.5). Keep 1e-3 only for giant-scale runs. After a bad weight, resume from the last FULL-STATE pre-collapse ckpt (`--resume-from`) — best_local_*.pt is model-only and cannot resume.
- Put non-ASCII characters (e.g. `⚠` in a print) in trainer source — Windows cp1252 redirect raises UnicodeEncodeError mid-run and can kill a run between an eval and its ckpt write. ASCII only (`[WARN]`).

## Dependencies
- torch + cuPy/CUDA runtime on the venv (`py -3.10 -E`), llama.cpp NOT needed for training.
- Disk for ckpts (~1.65GB per step ckpt; 5 kept via prune_checkpoints).
- AC power, clean machine (no other GPU consumers; kill stale llama-server/python first; FreePhysicalMemory >= 5GB).

## Threats / pitfalls / bugs (all hit in this repo)
- **Auto-resume glob picks by FILENAME order, not chronos or health (2026-09-10):** `sorted(glob("step_{preset}_*.pt"))[-1]` resumed the chunk2b POST-COLLAPSE `_003500` over the healthy new `_003000`. Fix already in train_v6.py: BOTH the resume and prune globs are now STRICT digits-only `step_{preset}_[0-9]{6}.pt` (match only `{step:06d}` writes; renamed/backup files like `..._003500_postcollapse.pt` or `recovery_*.pt` can never shadow the trajectory). Discipline: for trajectory jumps ALWAYS `--resume-from <path>`; use plain auto-resume only when all step_* files are one contiguous healthy run. Never blind-launch "auto" after a collapse.
- `optimizer.load_state_dict` silently restores lr + initial_lr; `scheduler.load_state_dict` restores base_lrs + _last_lr (both under sklearn-style naming in PyTorch): any `--lr` passed on resume is IGNORED without patch 2.
- metrics CSV overwrite: `--fresh` rewrites only the header, wiping history; the probe run with `--fresh` left probe rows in `metrics_local140a50.csv` (rescued copy = `metrics_local140a50_probe_20260910.csv`).
- Corrupted step ckpts from mid-save kills: already fixed via atomic `.tmp` + `os.replace` (2026-09-07 corruption of _000300/_000600). Keep using it; verify every ckpt load.
- Stale skill registry: `aide-flagship-twin-mind-train` exists on disk but the skill tool reports "not found" — always read it from disk.
- xpop is NOT the dominator on this 146.5M (measured ~0.0023 while LBL ~10.8): don't "fix" xpop.
- schedule shortfall (cosmetic, no action): 10,000 steps x 1024 tok = ~10.24M tokens is far under Chinchilla-optimal (~2.9B) for 146.5M; operator decision to train fast, documented, keep.

## Verification gates
- VERIFIED (2026-09-10 chunk2a, stop-at 2500): each of the following gates on REAL artifacts — resume line + `LR re-base: 5e-04 -> 3e-04 on resume (cosine phase preserved)` (patch B proven: metrics lr col 2.75e-4@2010 -> 2.61e-4@2500, NOT 4.60e-4); z in canonical scale and DECLINING within it (step z 427.6@2010 -> 295.2@2500, EVAL z 289.1; patch A proven); 0 NaNs, vram 4.26-4.44GB < 5GB budget, clip_frac logged (1.000 = same pre-clip regime as chunk1c; the LR cut IS the fix, not a bigger clip); real EVALs with telemetry @2250 ce 7.7913 popS 0.478/popK 0.522 dead 0.000 and @2500 ce 7.7095 (best, +0.006 vs pre-fix 7.7031 — flat, not regressed); `step_local140a50_002500.pt` written + reload round-trip (step 2500, loss 7.9127 exact log match, optimizer lr 2.61e-4 initial_lr 3e-4, scheduler base_lrs [3e-4] step_count 2501); generation probe finite (1,8,24025), top5 spread p 0.0315->0.0128, entropy 8.0, vocab cv 8.88 = peaked not collapsed, all 14 MoE layers emitted meta. NO SMOKE RUNS — the real chunk IS the gate.
- **VERIFIED (2026-09-10 chunk2d2, `--resume-from step_local140a50_003000.pt` `--z-weight 1e-4 --stop-at 4000`, 32.9 min, 0 NaNs, err.log empty): z-weight 1e-4 CONCLUSIVELY LOCKED — healthy trajectory confirmed across two consecutive full real chunks.** z Entered and stayed in the healthy band: 196@3010 -> 125@3170 (crossed <150 target) -> 100-110@3500-3700 (stable center band) -> 134@EVAL4000 (mild natural noise, well within band). LBL 7.98->6.38. eval ce improving with clean plateaus: 7.7371@2750 -> 7.6486@3000 -> 7.7053@3250 -> 7.6502@3500 -> 7.6116@3750 -> 7.6323@4000 (regress_streak 1 = single-eval noise, not a trend); **best 7.6116**. popS/popK 0.52/0.48 both evals, dead 0.000, clip_frac 0.976-0.993 (mild-clip regime now that z is in band — fewer pre-clip spikes). 1500 pre-clip grad-norm spikes seen at 3680/3970/4000 (caught by clip, no NaN). tok/s 500-540, vram 4.26GB. `step_local140a50_004000.pt` written + round-trip VERIFIED: step 4000, loss 6.6345 exact log match, optimizer lr 2.018e-4 / initial_lr 3e-4, scheduler base_lrs [3e-4] / _last_lr 2.018e-4 / _step_count 4001; generation probe: finite (1,64,24025), top5 p 0.0367->0.0113, entropy 7.622, vocab cv 10.199 peaked not collapsed, 14 MoE layers meta. err.log EMPTY = cp1252 crash bug fix confirmed across real run. `resuming at step 4000 best=6.6345` line correct. Clean DONE/exit (RAM recovered 7.18GB). The 1e-4 bisect now proven across chunk2c + chunk2d2: no collapse, no dead router, no regress streak >1, balanced pops, all 3 ckpts written.
- **VERIFIED (2026-09-10 chunk2e, `--resume-from _004000` `--z-weight 1e-4 --stop-at 5000`, 25.6 min, 0 NaNs, err.log empty): 1e-4 held across a THIRD real chunk.** z stayed planted in band: 99-105 around step 5000 (EVAL z 101.76, popS 0.530/popK 0.470, dead 0.000, clip_frac 0.996); LBL 5.7-6.0; eval ce **best 7.6033** (new overall low, beating 7.6116), EVAL@5000 7.6280 (regress_streak 1 = noise; no streak >1 all session). tok/s 684-687 (fastest yet — FreeRAM 6.1GB at launch, no host-shadow pressure). `step_local140a50_005000.pt` written + round-trip VERIFIED: step 5000, loss 6.6189 exact, optimizer lr 1.548e-4 / initial_lr 3e-4 (cosine midpoint ✓), base_lrs [3e-4], _step_count 5001; probe: finite, top5 p 0.0394->0.0141, entropy 7.496, cv 10.815 peaked, 14 MoE meta. Clean exit, no strays.
- **VERIFIED (2026-09-10 chunk2f, `--resume-from _005000` `--z-weight 1e-4 --stop-at 6000`, 25.6 min, 0 NaNs, err.log empty): 1e-4 held across a FOURTH real chunk.** z steady 92-103 all chunk (near-band, stable); EVALs z 96.78@5750 / 95.29@6000; popS 0.531/popK 0.469, dead 0.000, clip_frac 0.996; LBL 5.3-5.8; eval ce best 7.6040@5750 (global second to 7.6033), 7.6267@6000. tok/s 640-682, FreeRAM 4.6-5.2GB mid-run. `step_local140a50_006000.pt` written + round-trip VERIFIED: step 6000, loss 6.6946 exact, optimizer lr 1.073e-4 / initial_lr 3e-4 (cosine at 60% ✓), base_lrs [3e-4], _step_count 6001; probe: finite, top5 p 0.0318->0.0143, entropy 7.565, cv 9.585 peaked, 14 MoE meta. Clean exit. **Eval-best curve so far: 7.7095@2500 → 7.6486@3000 → 7.6116@3750 → 7.6033 → 7.6040@5750; cosine LR now ~1.07e-4 at 60% of run.**
- **VERIFIED (2026-09-10 chunk2j FINAL, `--resume-from _009000` `--z-weight 1e-4 --stop-at 10000`, 20.6 min, 0 NaNs, err.log empty): PRETRAIN COMPLETE.** Step 10000 hit with cosine fully ended — on-disk lr = 0.0, base_lrs [3e-4], _step_count 10001. EVAL@10000: ce 7.6289, z 45.04, popS 0.521/popK 0.479, dead 0.000. **best 7.6017 (final global low)**; `best_local140a50.pt` updated. `_010000.pt` round-trip exact (loss 6.6882); probe peaked cv 10.027. **TWIN-MIND DIAGNOSTIC: per-layer z 15.2-58.5, all 14 MoE layers engaged, zero dead; Sheldon(0-23) 52.2% / Spock(24-47) 47.8% over 458,752 real selections with layer specialization 59/41..45/54 — both populations genuinely used, not collapsed, not uniform.** Full run: chunks 2c..2j, every chunk 0 NaNs/empty err/verified ckpt/clean exit; eval best 7.7095@2500..7.6017@10k; z 279→44; LBL 7.98→4.5; tok/s 470→856; ~10.24M tokens. Follow-up: 6-check twin-mind battery, then D4 SFT marker retrain (sft_local30a12).
- Actually-PROVEN BAD (2026-09-10 chunk2b, stop-at 4000): z-weight 1e-3 crushed the router uniform (z 0.03, lse ~0.05) — eval ce best 7.6066@3250 then 8.2422@3750 / 8.2717@4000 (regress_streak 3-fired before the crash). Do NOT keep 1e-3 at this scale.
- PENDING (chunk2g onward, in flight 2026-09-10, 6000→7000→...→10000, same flags `--resume-from <prev>` `--z-weight 1e-4`): gates — z stays in band (92-150 healthy window now empirically pinched), eval ce improves toward/below 7.6033 with NO regress_streak >= 3, popS/popK ~0.5, dead 0.000, 0 NaNs, vram <5GB; each stop-at ckpt written + round-trip + generation probe; clean exit. After 10000: 6-check twin-mind battery (real generation probes for personality/format/coherence), then D4 SFT marker retrain (sft_local30a12, marker 5/3/1.0, 1200 steps, lr 2.5e-5, warmup 50, clip 0.5, Gate 1 byte-exact ≥80%).
- During the remaining schedule: eval CE must NOT keep regressing past 7.7095; z must keep declining toward the <100-150 rein-in target; clip fraction logged; per-pop util / dead-expert tracked; on 3-consecutive-eval-regression a warning fires and we investigate before the next chunk.
- After pretrain completes: 6-check twin-mind battery, then D4 SFT marker retrain (`sft_local30a12`, marker weights 5/3/1.0, 1200 steps, lr 2.5e-5, warmup 50, clip 0.5, Gate 1 byte-exact >= 80%).