ROCm-safe flip rule (kthvalue/where), --accum, --freeze-aux, --lr-horizon, launch_rig.sh/bench_rig.py for the RX 5600 XT rig

bench_rig.py

@@ -0,0 +1,51 @@
+"""bench_rig — does the v2 model train on this GPU, and how fast? (ROCm / any CUDA-like device)
+Builds the big1 architecture, runs fwd/bwd+Adam steps at a few batch sizes with random tokens,
+reports s/step, tokens/s and peak memory; then times one FlipOptimizer step (svd_lowrank on GPU).
+  python3 bench_rig.py [K] [batches e.g. 1,2,4]
+"""
+import sys, os, time, torch, torch.nn.functional as F
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from ternary_lm import Model
+from flip_lm import FlipOptimizer
+from train_live import chunked_ce
+
+K = int(sys.argv[1]) if len(sys.argv) > 1 else 2048
+batches = [int(x) for x in (sys.argv[2] if len(sys.argv) > 2 else "1,2,4").split(",")]
+dev = "cuda"
+print("device:", torch.cuda.get_device_name(0), "| torch", torch.__version__, "| hip", getattr(torch.version, "hip", None))
+m = Model(32768, d=512, blocks=8, K=K, lanes=4, heads=8, pos="rope", emb_bits=8, qkv_bits=8, res_shift=2, final_norm=True, qk_norm=True).to(dev)
+print("params", sum(p.numel() for p in m.parameters()))
+for amp in (False,):
+    for B in batches:
+        try:
+            opt = torch.optim.Adam(m.parameters(), lr=1e-4); scaler = torch.amp.GradScaler("cuda", enabled=amp)
+            torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
+            def step():
+                ctx = torch.randint(0, 32768, (B, K), device=dev); tgt = torch.randint(0, 32768, (B, K), device=dev)
+                with torch.autocast("cuda", dtype=torch.float16, enabled=amp):
+                    h = m.hidden(ctx); loss = chunked_ce(m, h, tgt, 32768)
+                opt.zero_grad(set_to_none=True); scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
+                return loss
+            for _ in range(2): step()
+            torch.cuda.synchronize(); t0 = time.time(); n = 4
+            for _ in range(n): l = step()
+            torch.cuda.synchronize(); dt = (time.time() - t0) / n
+            print(f"amp={amp} B={B} K={K}: {dt:.2f} s/step | {B*K/dt:,.0f} tok/s | peak {torch.cuda.max_memory_allocated()/2**30:.2f} GB | loss {l.item():.2f}")
+            del opt, scaler
+            for p in m.parameters(): p.grad = None
+            torch.cuda.empty_cache()
+        except Exception as e:
+            print(f"amp={amp} B={B}: FAILED {type(e).__name__}: {str(e)[:160]}"); torch.cuda.empty_cache()
+# flip step timing (codes only, no Adam state on weights)
+try:
+    for p in m.parameters(): p.grad = None
+    fo = FlipOptimizer(m, rank=64, refresh=10, rho=1e-5, aux_lr=0)
+    B = batches[0]; ctx = torch.randint(0, 32768, (B, K), device=dev); tgt = torch.randint(0, 32768, (B, K), device=dev)
+    torch.cuda.reset_peak_memory_stats(); torch.cuda.synchronize(); t0 = time.time()
+    for i in range(3):
+        h = m.hidden(ctx); loss = chunked_ce(m, h, tgt, 32768)
+        fo.zero_grad(); loss.backward(); fo.step()
+    torch.cuda.synchronize()
+    print(f"FLIP B={B}: {(time.time()-t0)/3:.2f} s/step | peak {torch.cuda.max_memory_allocated()/2**30:.2f} GB | flips {fo.total_flips} | state {fo.state_bytes()/2**20:.0f} MB")
+except Exception as e:
+    print("FLIP FAILED", type(e).__name__, str(e)[:200])

flip_lm.py

@@ -21,6 +21,22 @@ per weight = the 2-bit code. Fixed rho "boils" after a while in the source's run
 import torch
 import torch.nn as nn
 
+_SVD_ON_DEVICE = {}   # device -> bool; some ROCm builds (gfx1010) lack the LAPACK kernels svd_lowrank needs
+
+
+def _top_directions(G, rank):
+    """Top-`rank` left singular directions of G [O, I] via randomized SVD; falls back to CPU when the
+    device's linear-algebra kernels are missing (HIP 'invalid device function' on Navi 10)."""
+    dev = str(G.device)
+    if _SVD_ON_DEVICE.get(dev, True):
+        try:
+            U_, _, _ = torch.svd_lowrank(G, q=rank, niter=2)
+            return U_
+        except Exception:
+            _SVD_ON_DEVICE[dev] = False
+    U_, _, _ = torch.svd_lowrank(G.detach().cpu(), q=rank, niter=2)
+    return U_.to(G.device)
+
 
 class FlipState:
     """Low-rank Adam state + flip rule for ONE ternary weight tensor (a Lin.weight holding codes)."""
@@ -48,8 +64,7 @@ class FlipState:
         if not torch.isfinite(G).all():                                   # a bad batch must not flip anything
             return 0
         if self.P is None or step_idx % self.refresh == 0:
-            U_, _, _ = torch.svd_lowrank(G, q=self.rank, niter=2)        # top-r directions of THIS step's gradient
-            self.P = U_.contiguous()                                      # [O, r]
+            self.P = _top_directions(G, self.rank).contiguous()           # [O, r]: top-r left singular directions
         g = self.P.t() @ G                                                # [r, I]
         b1, b2 = self.betas
         self.t += 1
@@ -70,12 +85,11 @@ class FlipState:
         valid = (step != 0) & (new.abs() <= 1)
         gain = torch.where(valid, D.abs(), torch.zeros_like(D))
         k = max(1, int(rho * D.numel()))
-        top = torch.topk(gain.flatten(), k).indices
-        sel = torch.zeros(D.numel(), dtype=torch.bool, device=D.device)
-        sel[top] = True
-        sel = sel.view_as(D) & valid & (gain > 0)
+        # threshold via kthvalue (ROCm gfx1010 lacks the topk kernel for mid-sized tensors; kthvalue works)
+        thr = torch.kthvalue(gain.flatten(), max(1, D.numel() - k + 1)).values
+        sel = valid & (gain >= thr) & (gain > 0)
         n = int(sel.sum())
-        codes[sel] = new[sel]
+        codes.copy_(torch.where(sel, new, codes))          # elementwise (masked index_put is missing on ROCm gfx1010)
         self.flips += n
         return n
 
@@ -148,8 +162,9 @@ class FlipOptimizer:
 
     def assert_on_grid(self):
         for n, st in self.states.items():
-            u = set(st.w.detach().unique().tolist())
-            assert u <= {-1.0, 0.0, 1.0}, f"{n} left the ternary grid: {u}"
+            w = st.w.detach()
+            ok = bool(((w == -1) | (w == 0) | (w == 1)).all())        # (torch.unique is missing on some ROCm builds)
+            assert ok, f"{n} left the ternary grid"
 
     # ---- persistence (the codes themselves live in the model's state_dict) ----
     def state_dict(self):

launch_rig.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# launch_rig.sh — run a TernaryTest training arm on the desktop rig's RX 5600 XT (ROCm container).
+#   usage (run ON the rig, in ~/TernaryTest):  bash launch_rig.sh <name> <flip|adam> [steps] [extra train_live args...]
+#   e.g. bash launch_rig.sh flipB flip 1000            # arm B: low-rank-Adam flips from big1's step-22k codes
+#        bash launch_rig.sh adamA adam 1000            # arm A: true Adam on latents, same start, same data
+# 6 GB VRAM: fp32 (fp16 is slower on Navi 10), batch 1 x K2048 with --accum 8 = 16k tokens/step like big1.
+# Streams the same 36 sources (the rig has 46 GB RAM). Banks to HF <name>/ (token mounted from ~/.cache).
+cd "$(dirname "$0")"
+NAME=${1:?name}; MODE=${2:?flip|adam}; STEPS=${3:-1000}; shift 3 2>/dev/null || shift $#
+mkdir -p "$NAME"
+if [ "$MODE" = "flip" ]; then
+  ARM="--flip --rho 1e-5 --flip-rank 64 --flip-refresh 10 --rollback-tol 0.0 --freeze-aux"
+else
+  ARM="--freeze-aux --clip 1.0 --lr-horizon 200000"      # big1's own (nearly flat) lr over a short window
+fi
+docker rm -f "tt-$NAME" >/dev/null 2>&1
+docker run -d --name "tt-$NAME" --device=/dev/kfd --device=/dev/dri --group-add video --group-add render \
+  --security-opt seccomp=unconfined --shm-size 16g \
+  -v "$HOME/TernaryTest:/work" -v "$HOME/.cache/huggingface:/root/.cache/huggingface" -w /work \
+  rocm63-gfx1010:trainer bash -lc "python3 train_live.py --out /work/$NAME --init /work/big1/best_model.pt \
+    --tokenizer /work/tokenizer.json --pack --mix big \
+    --d 512 --blocks 8 --heads 8 --lanes 4 --emb-bits 8 --qkv-bits 8 --res-shift 2 --final-norm --qk-norm \
+    --K 2048 --batch 1 --accum 8 --steps $STEPS --warmup 0 --lr 3e-3 --eval-every 100 --eval-windows 4 \
+    --heldout-chars 100000 --buffer-chars 3000000 --amp 0 --sample-chars 200 \
+    --hf-repo CodeMasterCody3D/ternary-composer-charlm --hf-subdir $NAME --bank-every-min 30 $ARM $* \
+    > /work/$NAME/run.out 2>&1"
+echo "started container tt-$NAME (mode $MODE, $STEPS steps); log: ~/TernaryTest/$NAME/run.out"

ternary_lm.py

@@ -284,8 +284,7 @@ class Model(nn.Module):
             params[n].copy_(det_tern(params[n]))
             vals.append(params[n].flatten())
         v = torch.cat(vals)
-        u = set(v.unique().tolist())
-        assert u <= {-1.0, 0.0, 1.0}, f"non-ternary values present: {u}"
+        assert bool(((v == -1) | (v == 0) | (v == 1)).all()), "non-ternary values present"   # no torch.unique (ROCm)
         n = v.numel()
         return {"n": n, "frac_neg": (v == -1).float().mean().item(),
                 "frac_zero": (v == 0).float().mean().item(),

train_live.py

@@ -57,9 +57,11 @@ def get_args():
     ap.add_argument("--qk-norm", action="store_true", help="RMS-normalize q,k per head (bounded attention scores)")
     ap.add_argument("--float", action="store_true", help="float weights (teacher / baseline)")
     ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--accum", type=int, default=1, help="micro-batches per optimizer step (small-VRAM GPUs)")
     ap.add_argument("--steps", type=int, default=60000)
     ap.add_argument("--lr", type=float, default=3e-3)
     ap.add_argument("--warmup", type=int, default=200)
+    ap.add_argument("--lr-horizon", type=int, default=None, help="cosine horizon in steps (default: --steps); e.g. 200000 keeps big1's schedule in a short arm")
     ap.add_argument("--clip", type=float, default=1.0, help="grad-norm clip (0 = off)")
     ap.add_argument("--eval-every", type=int, default=500)
     ap.add_argument("--eval-windows", type=int, default=8, help="held-out windows per source")
@@ -79,6 +81,7 @@ def get_args():
     ap.add_argument("--rho", type=float, default=1e-5, help="fraction of each layer's weights flipped per step")
     ap.add_argument("--aux-lr", type=float, default=1e-4, help="Adam lr for the non-ternary params (emb, norms, biases)")
     ap.add_argument("--rollback-tol", type=float, default=0.0, help="on an eval worse than best*(1+tol): restore best codes, rho *= 0.5")
+    ap.add_argument("--freeze-aux", action="store_true", help="freeze non-ternary params (emb, norms, biases) in BOTH modes — equal-footing A/B")
     ap.add_argument("--fixed-corpus", default=None, help="dir from fixed_corpus.py build (no HF streams; ~3 GB RSS)")
     ap.add_argument("--hf-repo", default=None)
     ap.add_argument("--hf-subdir", default=None)
@@ -329,9 +332,17 @@ def main():
             say("NON-FINITE WEIGHTS in the resume checkpoint — refusing to continue from it"); return 5
     if ctrl.get("teacher"):
         say(attach_teacher(ctrl["teacher"]))
+    if a.freeze_aux and not a.flip:            # Adam arm on the ternary latents only (matches --flip --aux-lr 0)
+        tern = set(m.ternary_weight_names())
+        for n_, p in m.named_parameters():
+            if n_ not in tern:
+                p.requires_grad_(False)
+        opt = torch.optim.Adam([p for n_, p in m.named_parameters() if n_ in tern], lr=a.lr)
+        say(f"freeze-aux: Adam on {len(tern)} ternary latent tensors only")
     if a.flip:
         from flip_lm import FlipOptimizer
-        flip = FlipOptimizer(m, rank=a.flip_rank, refresh=a.flip_refresh, rho=float(ctrl["rho"]), aux_lr=a.aux_lr)
+        flip = FlipOptimizer(m, rank=a.flip_rank, refresh=a.flip_refresh, rho=float(ctrl["rho"]),
+                             aux_lr=0.0 if a.freeze_aux else a.aux_lr)
         if a.resume and os.path.exists(a.resume) and ck.get("flip"):
             flip.load_state_dict(ck["flip"])
         flip.assert_on_grid()
@@ -345,7 +356,7 @@ def main():
     event("start", arch=m.arch(), args=vars(a), control=ctrl, vocab=Vsz)
 
     def lr_at(s):
-        T = max(1, ctrl["steps"] - a.warmup)
+        T = max(1, (a.lr_horizon or ctrl["steps"]) - a.warmup)
         base = (s + 1) / a.warmup if s < a.warmup else 0.5 * (1 + math.cos(math.pi * min(1.0, (s - a.warmup) / T)))
         return a.lr * base * float(ctrl["lr_mult"])
 
@@ -367,28 +378,35 @@ def main():
                    os.path.join(a.out, "best_model.pt"))
 
     def train_step():
+        """One optimizer step = `a.accum` micro-batches of `batch` x K tokens (gradients accumulate)."""
         nonlocal tokens_seen
         K, B = int(ctrl["K"]), int(ctrl["batch"])
-        ctx, tgt, _, _ = mixer.batch(B, K, device=dev)
-        with torch.autocast("cuda", dtype=torch.float16, enabled=amp):
-            h = m.hidden(ctx)
-            th = teacher.hidden(ctx) if teacher is not None else None
-            loss = chunked_ce(m, h, tgt, Vsz, th, teacher, float(ctrl["kd_temp"]), float(ctrl["kd_alpha"]))
-        if flip is not None:                 # codes-only training: low-rank Adam -> flips (scale-invariant, no scaler)
+        if flip is not None:
             flip.zero_grad()
-            loss.backward()
+        else:
+            opt.zero_grad(set_to_none=True)
+        tot = 0.0
+        for _ in range(a.accum):
+            ctx, tgt, _, _ = mixer.batch(B, K, device=dev)
+            with torch.autocast("cuda", dtype=torch.float16, enabled=amp):
+                h = m.hidden(ctx)
+                th = teacher.hidden(ctx) if teacher is not None else None
+                loss = chunked_ce(m, h, tgt, Vsz, th, teacher, float(ctrl["kd_temp"]), float(ctrl["kd_alpha"])) / a.accum
+            if flip is not None:             # codes-only training: low-rank Adam -> flips (scale-invariant, no scaler)
+                loss.backward()
+            else:
+                scaler.scale(loss).backward()
+            tot += loss.item()
+            tokens_seen += B * K
+        if flip is not None:
             flip.rho = float(ctrl["rho"])
             flip.step()
-            tokens_seen += B * K
-            return loss.item()
-        opt.zero_grad(set_to_none=True)
-        scaler.scale(loss).backward()
+            return tot
         if a.clip > 0:                       # finite-but-astronomical gradients poison Adam's moments
             scaler.unscale_(opt)
             torch.nn.utils.clip_grad_norm_(m.parameters(), a.clip)
         scaler.step(opt); scaler.update()
-        tokens_seen += B * K
-        return loss.item()
+        return tot
 
     t0, run_loss, run_n, nan_run = time.time(), 0.0, 0, 0
     best_codes = flip.snapshot() if flip is not None else None