[gspark-etri]

Background

After #38, hot experts on RNGD were at parity with host-only: 13.0 vs 13.5 s/step, alternating steps at host load 120–150. #35 is still open for "faster than host-only". A per-step profile (Qwen3-30B-A3B, LoRA r16 q·v, S=128, hot fraction 0.25) shows where the time went: HostMLP.backward recomputes the whole MoE block forward before its autograd pass. That is about a third of the host MoE time, and one of the three device calls per layer.

Changes

• npu_lora/moe.py: HostMLP(keep_graph=True). forward builds the autograd graph and keeps (x, y) per layer. backward reuses it when it receives the same input tensor, and otherwise recomputes as before. It costs the block's saved activations until that layer's backward, a few MB per layer at S=128.
• Tests:
  • test_host_mlp_keep_graph_equals_recompute: LoRA input gradients and full-FT parameter gradients are bitwise equal to the recompute path;
  • the hot-expert test now expects 2 device calls per layer instead of 3.
• Record: a follow-up section in experiments/2026-10-21-moe-batched-hot-experts.md, with raw data and the profiler in experiments/data/moe-batched/. STATUS updated.

Verification

CPU: python -m pytest -q tests (27 passed, torch 2.14.1+cpu, transformers 5.20.2).

RNGD setup:

• hardware: npu-node-1, one card in a run-and-exit k8s Job; furiosa-torch 2026.2.0 / furiosa-llm 2026.3.0, torch 2.10.0+cpu, transformers 5.1.0;
• model and data: Qwen3-30B-A3B-Instruct-2507, NSMC data, seed 0, hot set fixed after 4 host-only warm-up steps;
• method: one process; 6 rounds of 4 configurations, order reversed every other round; load average 53;
• commit: 1fee9f2.

config (median step)	step	attention + head (device)	HostMLP fwd	HostMLP bwd (of which recompute)
host-only (as main)	9.21 s	1.26	2.48	5.34 (2.48)
host-only + keep_graph	8.00 s	1.16	2.46	4.31 (0)
hot (as #38)	6.70 s	1.07	1.73	3.79 (1.66)
hot + keep_graph	5.19 s	1.10	1.81	2.21 (0)

Paired per round:

• hot + keep_graph minus host-only + keep_graph: median −2.66 s (−33%), 6 of 6 rounds negative;
• hot + keep_graph minus host-only as on main: median −4.17 s (−45%).

Reproduce: python3 -m npu_lora.moe_prof MODEL DATA 0.25 128 6 out.jsonl, with experiments/data/moe-batched/moe_prof.py copied into npu_lora/.

Compatibility and limits

• keep_graph is on by default. It speeds up host-only MoE too (9.21 → 8.00 s). Pass keep_graph=False to get the old memory profile.
• The gain depends on the shared host's load. At load 120–150 (MoE hot experts: one batched device call per layer #38's A/B) hot and host-only were at parity without keep_graph; at load 53, hot was already faster.
• Not tried: E=64 and S=256/512. Each needs about 40 min of compile per expert bucket, and S needs new attention graphs. They were not needed for the criterion.

Related Issue

Closes #35

🤖 Generated with Claude Code