MoE: keep the host MoE forward graph for backward; hot experts beat host-only on RNGD

.research/events.jsonl

@@ -10,3 +10,4 @@
 {"id": "9f702fdb114e4b25b23508bb2bab2d14", "ts": "2026-10-21T15:37:39+09:00", "task": "#34 resident device AdamW + host-throughput step model", "result": "compile --adamw-only 150 s (was >2.5 h); 2 RNGD x 8 resident layers 4.57 s/step vs 1-card offload 5.22 s (+14% tok/s, steps 3-7, intermittent OOM #40); --update-workers 8 -8%; host-throughput model within 15% on 6 configs", "next": "fix resident OOM (#40), then 200-step run", "state": "active", "done": false, "blocker": "", "metrics": {"tok_s": "448"}, "params": {"seq_len": "512", "accum": "4", "resident_layers": "8", "cards": "2", "sdk": "furiosa-torch-2026.2.0"}, "artifacts": [], "experiments": ["experiments/2026-10-21-resident-adamw.md"], "labels": ["benchmark"], "review": true, "todos": [], "note": "", "agent": "claude", "host": "control-node", "branch": "perf/34-resident-adamw-compile", "commit": "8539644-dirty"}
 {"id": "f97382c821f344d08c88e8afc0c290ca", "ts": "2026-10-21T20:20:40+09:00", "task": "#33 Qwen2 gradient check on RNGD", "result": "npu-lora check passes (worst 3.1%, was 44%); 0 negative-loss steps in 400 (was 11); loss curve matches A100", "next": "Llama-3.1-8B check with the new head; Qwen3", "state": "active", "done": false, "blocker": "", "metrics": {"step_time_s": "2.74"}, "params": {}, "artifacts": [], "experiments": ["experiments/2026-10-21-qwen2-grad-precision.md"], "labels": ["bug", "experiment"], "review": true, "todos": [], "note": "", "agent": "claude", "host": "DESKTOP-R49IQ68", "branch": "fix/33-qwen-lora-grad", "commit": "a7e206a"}
 {"id": "b958b783d2b041f2ba43e072eb284abc", "ts": "2026-10-21T21:57:00+09:00", "task": "#35 MoE hot experts batched", "result": "batched hot experts 13.0 s/step vs host-only 13.5 (alternating steps, Qwen3-30B-A3B, RNGD); per-expert main 42.4 vs 17.2", "next": "more rows per call (longer S), hot fraction 0.5, drop the recomputed forward", "state": "active", "done": false, "blocker": "", "metrics": {"step_time_s": "13.0"}, "params": {}, "artifacts": [], "experiments": ["experiments/2026-10-21-moe-batched-hot-experts.md"], "labels": ["experiment", "benchmark"], "review": false, "todos": [], "note": "", "agent": "claude", "host": "DESKTOP-R49IQ68", "branch": "perf/35-batch-hot-experts", "commit": "9d5e498"}
+{"id": "c37c232f20a045e0aa3315cc1c786106", "ts": "2026-10-21T22:27:58+09:00", "task": "#35 MoE hot experts batched", "result": "hot+keep_graph 5.19 s/step vs host-only+keep 8.00 (paired median -2.66 s, 6/6 rounds) and vs main host-only 9.21 s; Qwen3-30B-A3B S128 RNGD", "next": "", "state": "idle", "done": true, "blocker": "", "metrics": {"step_time_s": "5.19"}, "params": {}, "artifacts": [], "experiments": ["experiments/2026-10-21-moe-batched-hot-experts.md"], "labels": ["experiment", "benchmark"], "review": false, "todos": [], "note": "", "agent": "claude", "host": "DESKTOP-R49IQ68", "branch": "perf/35-hot-experts-faster", "commit": "1fee9f2-dirty"}

STATUS.md

@@ -1,12 +1,12 @@
 # STATUS — RNGD NPU 파인튜닝 학습
 
 <!-- research-hub:header:begin (maintained by .research/report.py; keep the "- **Key**: value" format) -->
-- **State**: active
-- **Updated**: 2026-10-21 21:57+09:00
+- **State**: idle
+- **Updated**: 2026-10-21 22:27+09:00
 - **Agent**: claude@DESKTOP-R49IQ68
-- **Task**: #35 MoE hot experts batched
-- **Last result**: batched hot experts 13.0 s/step vs host-only 13.5 (alternating steps, Qwen3-30B-A3B, RNGD); per-expert main 42.4 vs 17.2
-- **Next**: more rows per call (longer S), hot fraction 0.5, drop the recomputed forward
+- **Task**: #35 MoE hot experts batched (done)
+- **Last result**: hot+keep_graph 5.19 s/step vs host-only+keep 8.00 (paired median -2.66 s, 6/6 rounds) and vs main host-only 9.21 s; Qwen3-30B-A3B S128 RNGD
+- **Next**: -
 - **Blockers**: none
 <!-- research-hub:header:end -->
 
@@ -42,6 +42,13 @@
 
 <!-- research-hub:log (report.py inserts new entries below this line, newest first) -->
 
+### 2026-10-21 22:27+09:00 — #35 MoE hot experts batched ✅
+
+- Result: hot+keep_graph 5.19 s/step vs host-only+keep 8.00 (paired median -2.66 s, 6/6 rounds) and vs main host-only 9.21 s; Qwen3-30B-A3B S128 RNGD
+- Metrics: step_time_s=5.19
+- Experiments: experiments/2026-10-21-moe-batched-hot-experts.md
+- By: claude@DESKTOP-R49IQ68 · perf/35-hot-experts-faster@1fee9f2-dirty
+
 ### 2026-10-21 21:57+09:00 — #35 MoE hot experts batched
 
 - Result: batched hot experts 13.0 s/step vs host-only 13.5 (alternating steps, Qwen3-30B-A3B, RNGD); per-expert main 42.4 vs 17.2

experiments/2026-10-21-moe-batched-hot-experts.md

@@ -78,3 +78,36 @@ Raw data in `experiments/data/moe-batched/`:
   - a larger hot fraction (E=64 graphs, about 29 GB of stacked weights on the card);
   - removing the forward that `HostMLP.backward` recomputes (one of the three calls per layer).
 - The losses of both paths match the CPU equivalence test (`test_moe_hot_experts_on_device_match_host`). Per-step losses on RNGD are not compared here: the head on `main` rounds the target logit (#33, fixed in #43), and several steps report small negative losses on both paths.
+
+## Follow-up: profile, then drop the recomputed forward (`HostMLP.keep_graph`)
+
+**Commit**: `perf/35-hot-experts-faster` `1fee9f2`. Same host, card and recipe as above, S=128, hot fraction 0.25.
+
+**Method** (`moe_prof.py`, raw data `prof-s128-roundrobin.jsonl`):
+- One process. 4 host-only warm-up steps fix the hot set.
+- Then 6 rounds of the four configurations, with the order reversed every other round.
+- Load average 53 during this run, lower than during the A/B runs above (120–150).
+
+| config (step median, 5–6 steps each) | step | attention + head on device | HostMLP fwd | HostMLP bwd | of which: recomputed fwd | cold-expert fwd | hot device fwd / bwd |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| host-only (as `main`) | 9.21 s | 1.26 | 2.48 | 5.34 | 2.48 | 2.46 | — |
+| host-only + keep_graph | 8.00 s | 1.16 | 2.46 | 4.31 | 0 | 1.31 | — |
+| hot (as #38) | 6.70 s | 1.07 | 1.73 | 3.79 | 1.66 | 1.27 | 0.55 / 0.52 |
+| **hot + keep_graph** | **5.19 s** | 1.10 | 1.81 | 2.21 | 0 | 0.64 | 0.42 / 0.59 |
+
+- The cold-expert forward column counts both the forward and the recompute, so it halves when the recompute is dropped.
+- The router takes ≤ 0.05 s per step.
+- `HostMLP.backward` recomputed the whole MoE block forward: about a third of the host MoE time, and one of the three device calls per layer.
+- `keep_graph` keeps the forward's autograd graph until that layer's backward. It is bitwise-equal to recomputing (`test_host_mlp_keep_graph_equals_recompute`) and costs a few MB of activations per layer at S=128.
+
+**Paired per round** (hot + keep_graph minus the host-only step of the same round):
+
+| baseline | per-round differences | median | rounds where hot is faster |
+|---|---|---:|---:|
+| vs host-only + keep_graph | −1.33, −3.13, −3.84, −2.50, −2.81, −1.25 s | **−2.66 s (−33%)** | 6 / 6 |
+| vs host-only as on `main` | −4.71, −5.03, −4.60, −2.58, −3.74, −3.63 s | **−4.17 s (−45%)** | 6 / 6 |
+
+**Conclusion**:
+- With the recompute dropped, hot experts on RNGD are faster than host-only by more than the round-to-round noise at this load: 6 of 6 rounds, median −2.66 s against the equally improved host-only path.
+- Even without `keep_graph`, hot (6.70 s) beat host-only (9.21 s) at load 53. In the A/B above, at load 120–150, they were at parity. The gain depends on how busy the shared host is, because the cold experts and the backward run there.
+- Not tried: E=64 and S=256/512. Both need about 40 min of compile per expert bucket, and S needs new attention graphs too. They were not needed to answer #35.

experiments/data/moe-batched/moe_prof.py

@@ -0,0 +1,119 @@
+"""#35: profile a MoE LoRA step and compare configurations round-robin in one process (one model load).
+
+Configs: host-only / hot experts on the device, each with HostMLP.keep_graph off (backward recomputes the forward)
+and on. Per step it logs the split: attention+head on the device (incl. transfers), HostMLP forward, backward
+recompute, backward autograd, and inside them the hot-expert device calls, the cold experts' forward, the router.
+usage: python3 -m npu_lora.moe_prof MODEL DATA FRAC SEQ ROUNDS OUT.jsonl
+"""
+import json
+import os
+import sys
+import time
+
+import torch
+
+from npu_lora import backends, moe as M, moe_device as MD
+from npu_lora.data import PackedDataset
+from npu_lora.trainer import LoRAConfig, LoRATrainer, TrainConfig
+
+model, data, frac, S, rounds, out = sys.argv[1], sys.argv[2], float(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5]), sys.argv[6]
+f = open(out, "a")
+T = {}
+
+
+def add(k, dt):
+    T[k] = T.get(k, 0.0) + dt
+
+
+def timed(k, fn):
+    def w(*a, **kw):
+        t0 = time.perf_counter()
+        try:
+            return fn(*a, **kw)
+        finally:
+            add(k, time.perf_counter() - t0)
+    return w
+
+
+def log(**kw):
+    kw.update(t=round(time.time(), 1), load1=os.getloadavg()[0] if hasattr(os, "getloadavg") else None)
+    f.write(json.dumps(kw) + "\n")
+    f.flush()
+    print(json.dumps(kw), flush=True)
+
+
+# timers (wrapping only; no behaviour change)
+M.HostMLP.forward = timed("mlp_fwd", M.HostMLP.forward)
+_orig_f = M.HostMLP._f
+_in_bwd = [False]
+
+
+def _f(self, i, a):
+    t0 = time.perf_counter()
+    try:
+        return _orig_f(self, i, a)
+    finally:
+        if _in_bwd[0]:
+            add("bwd_recompute", time.perf_counter() - t0)
+
+
+M.HostMLP._f = _f
+_orig_bwd = M.HostMLP._backward
+
+
+def _bwd(self, i, a, g):
+    _in_bwd[0] = True
+    try:
+        return _orig_bwd(self, i, a, g)
+    finally:
+        _in_bwd[0] = False
+
+
+M.HostMLP._backward = timed("mlp_bwd", _bwd)
+MD._HotBatch.forward = staticmethod(timed("hot_dev_fwd", MD._HotBatch.forward))
+MD._HotBatch.backward = staticmethod(timed("hot_dev_bwd", MD._HotBatch.backward))
+MD.expert_fwd = timed("cold_fwd", MD.expert_fwd)
+
+spec = os.environ.get("PROF_BACKEND", "furiosa:-1")
+be = backends.make(spec, cache_dir="/cache/furiosa") if spec.startswith("furiosa") else backends.make(spec)
+tc = TrainConfig(seq_len=S, lr=2e-4, steps=200, accum=1, warmup=2, seed=0, moe_hot_fraction=frac,
+                 moe_replan_every=10 ** 9)
+tr = LoRATrainer(model, be, LoRAConfig(r=16, alpha=32, targets=("q", "v")), tc, log=lambda s: None)
+t0 = time.time()
+tr.setup()
+for norm, mlp in tr.host_mlp.blocks:
+    mlp.gate.forward = timed("router_fwd", mlp.gate.forward)
+log(event="setup", seconds=round(time.time() - t0, 1), seq=S, **be.describe())
+from transformers import AutoTokenizer  # noqa: E402
+if os.environ.get("PROF_RANDOM_VOCAB"):                # smoke test on a tiny random model
+    _V = int(os.environ["PROF_RANDOM_VOCAB"])
+    it = iter([[(x, torch.roll(x, -1), None)] for x in (torch.randint(0, _V, (S,)) for _ in range(200))])
+else:
+    it = iter(PackedDataset(data, AutoTokenizer.from_pretrained(tr.path), S, seed=0).batches(1, 200))
+
+
+def step(name, hot, keep):
+    tr.host_mlp.keep_graph = keep
+    for w, h in zip(tr.moe_wraps, plan):
+        if w is not None:
+            w.hot = set(h) if hot else set()
+    T.clear()
+    s0 = tr.moe_stats()
+    mlp0 = tr.host_mlp.seconds
+    m = tr.train_step(next(it))
+    s1 = tr.moe_stats()
+    mlp_s = tr.host_mlp.seconds - mlp0
+    log(event="step", cfg=name, step=m["step"], loss=m["loss"], seconds=m["seconds"],
+        device_attn_head_s=round(m["seconds"] - mlp_s, 3), **{k: round(v, 3) for k, v in T.items()},
+        device_calls=s1["device_calls"] - s0["device_calls"], uploads=s1.get("weight_uploads", 0) - s0.get("weight_uploads", 0))
+
+
+plan = [set() for _ in tr.moe_wraps]
+for _ in range(4):                                    # warm-up host-only: routing counts for the hot set
+    step("warmup", False, True)
+MD.replan(tr.moe_wraps, tr.moe_dev)
+plan = [set(w.hot) if w is not None else None for w in tr.moe_wraps]
+cfgs = [("host", False, False), ("host+keep", False, True), ("hot", True, False), ("hot+keep", True, True)]
+for r in range(rounds):                               # round-robin: load drift hits every config alike
+    for name, hot, keep in (cfgs if r % 2 == 0 else cfgs[::-1]):
+        step(name, hot, keep)