diff --git a/.research/events.jsonl b/.research/events.jsonl index 69d5621..9832e4e 100644 --- a/.research/events.jsonl +++ b/.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"} diff --git a/STATUS.md b/STATUS.md index e39fccb..a7fd770 100644 --- a/STATUS.md +++ b/STATUS.md @@ -1,12 +1,12 @@ # STATUS — RNGD NPU 파인튜닝 학습 -- **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 @@ -42,6 +42,13 @@ +### 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 diff --git a/experiments/2026-10-21-moe-batched-hot-experts.md b/experiments/2026-10-21-moe-batched-hot-experts.md index 3a7941f..bc83a36 100644 --- a/experiments/2026-10-21-moe-batched-hot-experts.md +++ b/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. diff --git a/experiments/data/moe-batched/moe_prof.py b/experiments/data/moe-batched/moe_prof.py new file mode 100644 index 0000000..514a139 --- /dev/null +++ b/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) diff --git a/experiments/data/moe-batched/prof-s128-roundrobin.jsonl b/experiments/data/moe-batched/prof-s128-roundrobin.jsonl new file mode 100644 index 0000000..7b6207f --- /dev/null +++ b/experiments/data/moe-batched/prof-s128-roundrobin.jsonl @@ -0,0 +1,29 @@ +{"event": "setup", "seconds": 731.4, "seq": 128, "backend": "furiosa", "device": "rngd:0", "sdk": "2026.2.0", "tactic": "ForMiscModel", "max_shapes": 8, "t": 1792588781.6, "load1": 37.5908203125} +{"event": "step", "cfg": "warmup", "step": 1, "loss": 1.0967254638671875e-05, "seconds": 27.78, "device_attn_head_s": 4.284, "router_fwd": 0.037, "cold_fwd": 5.69, "mlp_fwd": 7.097, "mlp_bwd": 16.399, "device_calls": 0, "uploads": 0, "t": 1792588811.6, "load1": 43.63427734375} +{"event": "step", "cfg": "warmup", "step": 2, "loss": 2.384185791015625e-06, "seconds": 17.8, "device_attn_head_s": 1.531, "router_fwd": 0.035, "cold_fwd": 1.907, "mlp_fwd": 3.171, "mlp_bwd": 13.098, "device_calls": 0, "uploads": 0, "t": 1792588829.4, "load1": 43.9384765625} +{"event": "step", "cfg": "warmup", "step": 3, "loss": 0.0, "seconds": 28.4, "device_attn_head_s": 1.923, "router_fwd": 0.025, "cold_fwd": 1.958, "mlp_fwd": 3.209, "mlp_bwd": 23.268, "device_calls": 0, "uploads": 0, "t": 1792588857.8, "load1": 48.35107421875} +{"event": "step", "cfg": "warmup", "step": 4, "loss": 5.2130208015441895, "seconds": 5.889, "device_attn_head_s": 1.19, "router_fwd": 0.023, "cold_fwd": 1.254, "mlp_fwd": 2.308, "mlp_bwd": 2.391, "device_calls": 0, "uploads": 0, "t": 1792588863.7, "load1": 49.3642578125} +{"event": "step", "cfg": "host", "step": 5, "loss": 0.0, "seconds": 9.729, "device_attn_head_s": 1.37, "router_fwd": 0.055, "cold_fwd": 2.591, "mlp_fwd": 2.808, "bwd_recompute": 2.589, "mlp_bwd": 5.551, "device_calls": 0, "uploads": 0, "t": 1792588873.4, "load1": 49.91748046875} +{"event": "step", "cfg": "host+keep", "step": 6, "loss": 2.384185791015625e-06, "seconds": 6.356, "device_attn_head_s": 1.191, "router_fwd": 0.024, "cold_fwd": 1.311, "mlp_fwd": 2.487, "mlp_bwd": 2.677, "device_calls": 0, "uploads": 0, "t": 1792588879.8, "load1": 50.64501953125} +{"event": "step", "cfg": "hot", "step": 7, "loss": 0.0614926815032959, "seconds": 12.541, "device_attn_head_s": 1.352, "router_fwd": 0.053, "hot_dev_fwd": 6.294, "cold_fwd": 1.307, "mlp_fwd": 7.578, "bwd_recompute": 1.581, "hot_dev_bwd": 0.512, "mlp_bwd": 3.61, "device_calls": 144, "uploads": 48, "t": 1792588892.3, "load1": 51.08740234375} +{"event": "step", "cfg": "hot+keep", "step": 8, "loss": 3.0279159545898438e-05, "seconds": 5.021, "device_attn_head_s": 1.051, "router_fwd": 0.024, "hot_dev_fwd": 0.474, "cold_fwd": 0.645, "mlp_fwd": 1.817, "hot_dev_bwd": 0.633, "mlp_bwd": 2.152, "device_calls": 96, "uploads": 0, "t": 1792588897.4, "load1": 51.16064453125} +{"event": "step", "cfg": "hot+keep", "step": 9, "loss": 1.430511474609375e-06, "seconds": 5.32, "device_attn_head_s": 1.044, "router_fwd": 0.024, "hot_dev_fwd": 0.698, "cold_fwd": 0.631, "mlp_fwd": 2.001, "hot_dev_bwd": 0.634, "mlp_bwd": 2.275, "device_calls": 96, "uploads": 0, "t": 1792588902.7, "load1": 51.1474609375} +{"event": "step", "cfg": "hot", "step": 10, "loss": 4.76837158203125e-06, "seconds": 6.46, "device_attn_head_s": 1.06, "router_fwd": 0.053, "hot_dev_fwd": 0.535, "cold_fwd": 1.235, "mlp_fwd": 1.613, "bwd_recompute": 1.663, "hot_dev_bwd": 0.515, "mlp_bwd": 3.787, "device_calls": 144, "uploads": 0, "t": 1792588909.1, "load1": 51.2158203125} +{"event": "step", "cfg": "host+keep", "step": 11, "loss": 0.002410888671875, "seconds": 8.448, "device_attn_head_s": 1.1, "router_fwd": 0.024, "cold_fwd": 1.33, "mlp_fwd": 2.53, "mlp_bwd": 4.818, "device_calls": 0, "uploads": 0, "t": 1792588917.6, "load1": 50.14111328125} +{"event": "step", "cfg": "host", "step": 12, "loss": -2.384185791015625e-07, "seconds": 10.352, "device_attn_head_s": 1.318, "router_fwd": 0.058, "cold_fwd": 2.655, "mlp_fwd": 2.511, "bwd_recompute": 3.01, "mlp_bwd": 6.522, "device_calls": 0, "uploads": 0, "t": 1792588927.9, "load1": 52.65966796875} +{"event": "step", "cfg": "host", "step": 13, "loss": 0.01467132568359375, "seconds": 11.044, "device_attn_head_s": 1.408, "router_fwd": 0.061, "cold_fwd": 2.98, "mlp_fwd": 3.005, "bwd_recompute": 3.09, "mlp_bwd": 6.631, "device_calls": 0, "uploads": 0, "t": 1792588939.0, "load1": 55.0986328125} +{"event": "step", "cfg": "host+keep", "step": 14, "loss": 3.337860107421875e-06, "seconds": 10.284, "device_attn_head_s": 1.289, "router_fwd": 0.025, "cold_fwd": 1.238, "mlp_fwd": 2.368, "mlp_bwd": 6.627, "device_calls": 0, "uploads": 0, "t": 1792588949.3, "load1": 57.4443359375} +{"event": "step", "cfg": "hot", "step": 15, "loss": 4.0531158447265625e-06, "seconds": 7.533, "device_attn_head_s": 1.247, "router_fwd": 0.063, "hot_dev_fwd": 0.616, "cold_fwd": 1.56, "mlp_fwd": 2.038, "bwd_recompute": 1.887, "hot_dev_bwd": 0.577, "mlp_bwd": 4.247, "device_calls": 144, "uploads": 0, "t": 1792588956.8, "load1": 55.61767578125} +{"event": "step", "cfg": "hot+keep", "step": 16, "loss": 0.00040841102600097656, "seconds": 6.442, "device_attn_head_s": 1.262, "router_fwd": 0.028, "hot_dev_fwd": 0.425, "cold_fwd": 0.661, "mlp_fwd": 1.988, "hot_dev_bwd": 0.578, "mlp_bwd": 3.192, "device_calls": 96, "uploads": 0, "t": 1792588963.3, "load1": 56.529296875} +{"event": "step", "cfg": "hot+keep", "step": 17, "loss": -4.76837158203125e-07, "seconds": 5.93, "device_attn_head_s": 1.254, "router_fwd": 0.028, "hot_dev_fwd": 0.405, "cold_fwd": 0.636, "mlp_fwd": 1.705, "hot_dev_bwd": 0.541, "mlp_bwd": 2.971, "device_calls": 96, "uploads": 0, "t": 1792588969.2, "load1": 57.4482421875} +{"event": "step", "cfg": "hot", "step": 18, "loss": 3.0994415283203125e-06, "seconds": 6.695, "device_attn_head_s": 1.166, "router_fwd": 0.055, "hot_dev_fwd": 0.597, "cold_fwd": 1.274, "mlp_fwd": 1.734, "bwd_recompute": 1.662, "hot_dev_bwd": 0.539, "mlp_bwd": 3.795, "device_calls": 144, "uploads": 0, "t": 1792588975.9, "load1": 57.869140625} +{"event": "step", "cfg": "host+keep", "step": 19, "loss": 1.1920928955078125e-06, "seconds": 8.433, "device_attn_head_s": 1.111, "router_fwd": 0.023, "cold_fwd": 1.318, "mlp_fwd": 2.429, "mlp_bwd": 4.892, "device_calls": 0, "uploads": 0, "t": 1792588984.3, "load1": 57.71923828125} +{"event": "step", "cfg": "host", "step": 20, "loss": 2.384185791015625e-06, "seconds": 8.509, "device_attn_head_s": 1.121, "router_fwd": 0.047, "cold_fwd": 2.337, "mlp_fwd": 2.373, "bwd_recompute": 2.319, "mlp_bwd": 5.015, "device_calls": 0, "uploads": 0, "t": 1792588992.8, "load1": 56.91357421875} +{"event": "step", "cfg": "host", "step": 21, "loss": 3.337860107421875e-06, "seconds": 8.498, "device_attn_head_s": 1.204, "router_fwd": 0.045, "cold_fwd": 2.281, "mlp_fwd": 2.292, "bwd_recompute": 2.379, "mlp_bwd": 5.001, "device_calls": 0, "uploads": 0, "t": 1792589001.3, "load1": 56.76025390625} +{"event": "step", "cfg": "host+keep", "step": 22, "loss": 1.3113021850585938e-05, "seconds": 7.566, "device_attn_head_s": 1.122, "router_fwd": 0.024, "cold_fwd": 1.416, "mlp_fwd": 2.64, "mlp_bwd": 3.804, "device_calls": 0, "uploads": 0, "t": 1792589008.9, "load1": 56.5390625} +{"event": "step", "cfg": "hot", "step": 23, "loss": 0.0, "seconds": 5.694, "device_attn_head_s": 1.065, "router_fwd": 0.049, "hot_dev_fwd": 0.506, "cold_fwd": 1.048, "mlp_fwd": 1.437, "bwd_recompute": 1.382, "hot_dev_bwd": 0.493, "mlp_bwd": 3.192, "device_calls": 144, "uploads": 0, "t": 1792589014.6, "load1": 56.25537109375} +{"event": "step", "cfg": "hot+keep", "step": 24, "loss": 1.430511474609375e-06, "seconds": 4.756, "device_attn_head_s": 1.085, "router_fwd": 0.026, "hot_dev_fwd": 0.396, "cold_fwd": 0.606, "mlp_fwd": 1.6, "hot_dev_bwd": 0.595, "mlp_bwd": 2.071, "device_calls": 96, "uploads": 0, "t": 1792589019.4, "load1": 52.23095703125} +{"event": "step", "cfg": "hot+keep", "step": 25, "loss": 4.76837158203125e-06, "seconds": 5.052, "device_attn_head_s": 1.108, "router_fwd": 0.024, "hot_dev_fwd": 0.352, "cold_fwd": 0.728, "mlp_fwd": 1.795, "hot_dev_bwd": 0.554, "mlp_bwd": 2.149, "device_calls": 96, "uploads": 0, "t": 1792589024.4, "load1": 52.29296875} +{"event": "step", "cfg": "hot", "step": 26, "loss": 4.0663676261901855, "seconds": 6.876, "device_attn_head_s": 1.069, "router_fwd": 0.052, "hot_dev_fwd": 0.546, "cold_fwd": 1.369, "mlp_fwd": 1.846, "bwd_recompute": 1.7, "hot_dev_bwd": 0.522, "mlp_bwd": 3.961, "device_calls": 144, "uploads": 0, "t": 1792589031.3, "load1": 52.40185546875} +{"event": "step", "cfg": "host+keep", "step": 27, "loss": 2.86102294921875e-06, "seconds": 6.304, "device_attn_head_s": 1.245, "router_fwd": 0.021, "cold_fwd": 1.221, "mlp_fwd": 2.268, "mlp_bwd": 2.791, "device_calls": 0, "uploads": 0, "t": 1792589037.6, "load1": 52.6904296875} +{"event": "step", "cfg": "host", "step": 28, "loss": 0.0, "seconds": 8.687, "device_attn_head_s": 1.11, "router_fwd": 0.045, "cold_fwd": 2.275, "mlp_fwd": 2.448, "bwd_recompute": 2.278, "mlp_bwd": 5.129, "device_calls": 0, "uploads": 0, "t": 1792589046.3, "load1": 52.89208984375} diff --git a/npu_lora/moe.py b/npu_lora/moe.py index bac4458..024ead3 100644 --- a/npu_lora/moe.py +++ b/npu_lora/moe.py @@ -23,7 +23,8 @@ class HostMLP: 메모리는 파라미터당 bf16 2 B + master 4 B + AdamW 8 B (+ 누적 4 B) 다. """ - def __init__(self, path, dtype=BF, trainable=False, lr=1e-5, betas=(0.9, 0.999), weight_decay=0.0): + def __init__(self, path, dtype=BF, trainable=False, lr=1e-5, betas=(0.9, 0.999), weight_decay=0.0, + keep_graph=True): from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained(path, dtype=dtype) self.blocks = [] @@ -35,6 +36,11 @@ def __init__(self, path, dtype=BF, trainable=False, lr=1e-5, betas=(0.9, 0.999), del model # embeddings, attention and lm_head are freed; the blocks stay self.trainable = trainable self.seconds = 0.0 # host time in forward/backward (step-time split: host MoE vs the rest) + # keep_graph: forward builds the autograd graph and backward reuses it instead of recomputing the forward + # (one of the three passes; with hot experts on the device also one device call per layer). Costs the + # block's saved activations until that layer's backward (S=128: a few MB per layer). + self.keep_graph = keep_graph + self._saved = {} # layer -> (x, y) of the last forward, consumed by its backward if trainable: self.params = [[p for p in (*norm.parameters(), *mlp.parameters())] for norm, mlp in self.blocks] self.master = [[p.detach().float().clone().requires_grad_(True) for p in ps] for ps in self.params] @@ -59,8 +65,15 @@ def _f(self, i, a): def forward(self, i, a): t0 = time.perf_counter() - with torch.no_grad(): - y = self._f(i, a) + if self.keep_graph: + x = a.detach().requires_grad_(True) + with torch.enable_grad(): + y = self._f(i, x) + self._saved[i] = (x, y) + y = y.detach() + else: + with torch.no_grad(): + y = self._f(i, a) self.seconds += time.perf_counter() - t0 return y @@ -72,9 +85,14 @@ def backward(self, i, a, g): self.seconds += time.perf_counter() - t0 def _backward(self, i, a, g): - x = a.detach().requires_grad_(True) + saved = self._saved.pop(i, None) + if saved is not None and saved[0].data_ptr() == a.data_ptr() and saved[0].shape == a.shape: + x, y = saved # the forward's graph (same input tensor) + else: + x = a.detach().requires_grad_(True) + with torch.enable_grad(): + y = self._f(i, x) with torch.enable_grad(): - y = self._f(i, x) if not self.trainable: (gx,) = torch.autograd.grad(y, x, g.to(y.dtype)) return gx diff --git a/tests/test_cpu.py b/tests/test_cpu.py index 5003146..dc94203 100644 --- a/tests/test_cpu.py +++ b/tests/test_cpu.py @@ -484,6 +484,27 @@ def capture(self, clip, lr, k): assert abs(now - hf_now) / hf_now < 2e-2, (now, hf_now) +def test_host_mlp_keep_graph_equals_recompute(): + """HostMLP.keep_graph reuses the forward's autograd graph in backward: same input gradients (LoRA) and the + same host-part parameter gradients (full fine-tuning) as recomputing the forward.""" + from npu_lora import moe as M + d = tempfile.mkdtemp() + _tiny_moe(d) + g = torch.Generator().manual_seed(2) + a = (torch.randn(32, 64, generator=g)).to(BF) + gy = (torch.randn(32, 64, generator=g) * 0.1).to(BF) + for trainable in (False, True): + out = [] + for keep in (False, True): + h = M.HostMLP(d, trainable=trainable, keep_graph=keep) + y = h.forward(0, a) + gx = h.backward(0, a, gy) + out.append((y, gx, [x.clone() for x in h.acc[0]] if trainable else [])) + (y0, g0, p0), (y1, g1, p1) = out + assert torch.equal(y0, y1) and torch.equal(g0, g1) + assert all(torch.equal(u, v) for u, v in zip(p0, p1)) + + def test_hot_expert_packing(): """pack(): rows grouped by hot expert, each at slot * nb + rank; nb = bucket of the busiest expert; the batched expert graph equals the per-expert host computation on the packed rows""" @@ -541,9 +562,9 @@ def rel(a, b): host._upload_adapters(), hyb._upload_adapters() _, gh, _ = host.forward_backward(ids, lab) _, gd, _ = hyb.forward_backward(ids, lab) - # batched, per layer (2 layers): forward, the forward that HostMLP.backward recomputes, backward — not one - # call per hot expert - assert hyb.moe_stats()["device_calls"] == 6, hyb.moe_stats() + # batched, per layer (2 layers): forward and backward (HostMLP keeps the forward's graph, no recompute) — + # not one call per hot expert + assert hyb.moe_stats()["device_calls"] == 4, hyb.moe_stats() for a, b in zip((x for gl in gd for x in gl), (x for gl in gh for x in gl)): assert rel(a, b) < 5e-2, rel(a, b) # bf16: host autograd vs compiled vjp # full fine-tuning: host-part gradients at the first step