Repository navigation
npu-lora check: pin MoE routing to the NPU run; step-time split host MoE vs device - #32
Merged
Merged
Conversation
…MoE vs device (#6) On MoE checkpoints tiny numeric differences flip some tokens' top-k experts (Qwen3-30B-A3B, 2 layers, 128 tokens: 12 and 20 tokens), which changes their whole MoE output, so the gradient comparison failed for routing, not numerics. check now records the NPU run's experts and pins the CPU reference to them (moe.RoutingPins; router probabilities still flow), reports the unpinned flips, and keeps --free-routing for the old behaviour. train_step reports host_moe_s and device_and_transfer_s for MoE checkpoints. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
gspark-etri
added a commit
that referenced
this pull request
Oct 21, 2026
## Background #32 added `test_routing_pins_record_and_force_experts` below the `if __name__ == "__main__":` runner in `tests/test_cpu.py`, so running the file (what CI did after #37) never ran it: CI reported 21 tests, not 22. Once the test ran in CI (transformers 5.20.2), it **failed**: a run pinned to its own experts did not reproduce its own output. Cause: the Qwen-MoE router returns `(router_logits, scores, indices)`. In transformers 5.1 the first value is the softmax; in later releases (5.20) it is the raw logits. `RoutingPins` gathered the pinned experts' weights from that value, so with newer transformers the pinned weights were logits, not probabilities. That also affects `npu-lora check` on MoE checkpoints (#32's default), not only the test. ## Changes - `npu_lora/moe.py`: `RoutingPins` takes the softmax when the router's first output is not a probability distribution (negative values or rows not summing to 1). - `tests/test_cpu.py`: the runner block moves back to the end. - `.github/workflows/ci.yml`: one step, `python -m pytest -q tests`, which collects every `test_*` function regardless of position. ## Verification Local, Windows, torch 2.9.1+cpu: - transformers 5.20.2, before the fix: `pytest -k routing_pins` fails (926/1024 elements differ, max abs 1.23). - transformers 5.1.0 and 5.20.2, after the fix: `tests/test_cpu.py` 22 passed on both. (`tests/test_airm.py` needs `pip install -e .`, which CI does.) - CI on this PR: see checks. ## Compatibility and limits `npu-lora check` results for MoE taken with transformers > 5.1 before this fix are invalid for the pinned comparison; #32's table was measured on the node's transformers 5.1 stack, which returned probabilities. ## Related Issue Related to #32 and #36. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: gspark-etri <gspark-etri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
Author
|
Follow-up: on transformers > 5.1 the router returns logits, and RoutingPins used them as probabilities (fixed in #39). The new test also never ran in CI because it sat below the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
#6's remaining criteria are an RNGD vs CPU gradient check on Qwen3-30B-A3B (
npu-lora check, 2 layers) and a step-time split between NPU and host. With the current check, Qwen3-30B fails:For comparison, dense Qwen3-0.6B passes (worst 3.7%).
The cause is routing, not numerics. With identical inputs and adapters, the top-8 expert set differs between the RNGD and CPU runs for 12 / 128 tokens in layer 0 and 20 / 128 in layer 1. Each flip changes that token's whole MoE output and its gradient.
Changes
npu_lora/moe.py:RoutingPins.flips()counts the tokens whose expert set differs.npu-lora check. On MoE checkpoints the CPU reference is pinned to the NPU run's experts, and the result also reportsrouting_flips_unpinnedfrom an extra unpinned CPU pass.--free-routingkeeps the old comparison.LoRATrainer.train_step. On MoE checkpoints it also reportshost_moe_s(HostMLP forward + backward) anddevice_and_transfer_s.test_routing_pins_record_and_force_experts.Verification
RNGD 2026.4.0, furiosa-torch 2026.2.0, Qwen3-30B-A3B-Instruct-2507,
--layers 2 --seq-len 128:Step-time split, LoRA, seq 128, 3 steps on RNGD chip 0. The host load average was 35–43 from other users' jobs, so the host share is inflated:
The host MoE blocks dominate (97–98%), so the next lever is the host MoE path or hot experts on the device (#14).
CPU tests: the full suite passes.
Compatibility and limits
checkon dense models is unchanged.--free-routingrestores the old comparison.Related Issue
Closes #6: functional equivalence and the real-checkpoint run were already merged (#16), and this adds the RNGD gradient check and the step split.
🤖 Generated with Claude Code