Repository navigation
[MoE] Hot experts on RNGD during training (host dispatch, device expert graphs) #14
Description
Problem
In MoE training on RNGD (#6) every expert runs on the host. In inference (OptimusPrime moe_offload) hot experts stay resident on RNGD and only cold ones run on the host; training should do the same.
Evidence
Inference: Qwen3-30B / Solar-Open-100B with RNGD expert tiers (OptimusPrime rngd-opt, 186–640 µs per token for k=1..8 experts). Training today: all S×top-k expert evaluations on the host CPU, forward + recomputed forward + backward per layer.
Proposal
Routing and token dispatch stay on the host (no device gather). Per (H, F, bucket of rows) compile expert forward and backward graphs (backward: dx; and dW for full FT); hot experts' weights resident on the device, rows gathered on the host, padded to the bucket, uploaded, computed, results scattered back on the host. Placement from a routing profile, re-planned between steps as routing drifts (the inference planner's adaptive re-placement).
Scope
LoRA (experts frozen) first, then full FT (dW for resident experts, host AdamW, re-upload).
Done when
Device-expert path == host path on CPU backend (outputs and gradients); on RNGD, hot fraction and step time vs all-host reported for Qwen3-30B-A3B.
Activity
MoE training: hot experts on the device, the rest on the host
MoE training: hot experts on the device, the rest on the host
MoE training: hot experts on the device, the rest on the host
MoE checkpoints: LoRA with attention on RNGD, MoE block on the host (#16
MoE training: hot experts on the device, the rest on the host
MoE training: hot experts on the device (#18)
MoE checkpoints: LoRA with attention on RNGD, MoE block on the host (#16
MoE training: hot experts on the device (#18)
MoE hot experts: build the expert graphs with grad mode on (#27)
npu-lora check: pin MoE routing to the NPU run; step-time split host …
gspark-etri commentedon Oct 21, 2026
Done-when met by #38 and #44:
test_moe_hot_experts_on_device_match_hostchecks outputs and LoRA gradients, plus the device-call count.test_host_mlp_keep_graph_equals_recomputechecks LoRA input gradients and full-FT parameter gradients, bitwise.