From 281dfe3416169f27fc8124fcb09a846b37898a71 Mon Sep 17 00:00:00 2001 From: gspark-etri Date: Wed, 21 Oct 2026 14:57:41 +0900 Subject: [PATCH 1/2] CPU tests: run every test (one was defined after the __main__ block); CI uses pytest test_routing_pins_record_and_force_experts (#32) sat below the 'if __name__ == "__main__"' runner, so tests/test_cpu.py never ran it. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 7 +++---- tests/test_cpu.py | 25 ++++++++++++------------- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ac0c3fe..2b7ab49 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,6 @@ jobs: python -m pip install --upgrade pip pip install torch --index-url https://download.pytorch.org/whl/cpu pip install -e . pytest - - name: CPU tests - run: python tests/test_cpu.py - - name: ai-rm adapter tests - run: python -m pytest -q tests/test_airm.py + - name: CPU and ai-rm adapter tests + # pytest collects every test_* function, wherever it sits in the file + run: python -m pytest -q tests diff --git a/tests/test_cpu.py b/tests/test_cpu.py index 1a8bea3..6bed8db 100644 --- a/tests/test_cpu.py +++ b/tests/test_cpu.py @@ -497,19 +497,6 @@ def capture(self, clip, lr, k): assert rel(a[n], b[n]) < 5e-2, (n, rel(a[n], b[n])) -if __name__ == "__main__": - failed = 0 - for name, fn in sorted(globals().items()): - if name.startswith("test_") and callable(fn): - try: - fn() - print("ok ", name) - except Exception as e: # noqa: BLE001 - failed += 1 - print("FAIL", name, repr(e)[:200]) - sys.exit(1 if failed else 0) - - def test_routing_pins_record_and_force_experts(): """npu-lora check on MoE: the reference run can be pinned to another run's experts (moe.RoutingPins)""" from npu_lora.moe import HostMLP, RoutingPins @@ -536,3 +523,15 @@ def test_routing_pins_record_and_force_experts(): flipped = RoutingPins(mlp.blocks) flipped.ids = other.ids assert rec.flips(flipped)[0] > 0 + +if __name__ == "__main__": + failed = 0 + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + try: + fn() + print("ok ", name) + except Exception as e: # noqa: BLE001 + failed += 1 + print("FAIL", name, repr(e)[:200]) + sys.exit(1 if failed else 0) From f0916cc41c3839b6624513d694db0547a6db0b74 Mon Sep 17 00:00:00 2001 From: gspark-etri Date: Wed, 21 Oct 2026 15:04:20 +0900 Subject: [PATCH 2/2] RoutingPins: take the softmax when the router returns logits (transformers > 5.1) transformers 5.1 returns the router softmax as the first output; later releases return the raw logits. Pinning gathered the logits as weights, so a pinned run did not reproduce its own output on transformers 5.20. Co-Authored-By: Claude Opus 5.5 --- npu_lora/moe.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/npu_lora/moe.py b/npu_lora/moe.py index 9d1a7f5..bac4458 100644 --- a/npu_lora/moe.py +++ b/npu_lora/moe.py @@ -140,7 +140,11 @@ def _route(self, i, router, hidden): self.ids.setdefault(i, ids.detach().clone()) # the forward pass (backward recomputes it) elif self.mode == "pin" and i in self.ids: ids = self.ids[i] - val = probs.gather(-1, ids) + # transformers 5.1 returns the softmax here, later releases return the raw logits + p = probs.float() + if (p < 0).any() or not torch.allclose(p.sum(-1), torch.ones_like(p[..., 0]), atol=1e-3): + p = torch.softmax(p, dim=-1) + val = p.gather(-1, ids) if getattr(router, "norm_topk_prob", True): val = val / val.sum(dim=-1, keepdim=True) scores = val.to(scores.dtype)