CPU tests: run every test; fix RoutingPins on transformers > 5.1

.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

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)

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)