test: make the MoE fixture actually test nesting, and fix two harness faults

scripts/make-test-fixtures.py

@@ -23,6 +23,7 @@
 """
 import json
 import os
+import shutil
 import sys
 
 import numpy as np
@@ -230,77 +231,107 @@ def build_gemma4_kv_shared(out, vestigial):
 
 
 def build_moe_nested(out):
-    """A Qwen3-style MoE: per-expert tensors, nested config, decoy vision count.
-
-    Covers the MoE *load* path, which nothing else in CI touches — no MoE model is in
-    the matrix at all. Specifically: per-expert `mlp.experts.N.*` tensors being stacked
-    into `switch_mlp` by sanitize, and an expert count nested under `text_config`
-    rather than at top level.
-
-    What it does NOT cover, despite the nesting, is the #112 detection bug itself.
-    That bug was `num_local_experts` being the only spelling checked, so a config
-    declaring `num_experts` was called dense. But `modelTypeImpliesMoE` treats any
-    model_type containing "moe" as MoE regardless of keys, so `qwen3_moe` is caught by
-    that fallback whatever the config says — verified by reverting detection to the
-    single-key form and watching this fixture still enable streaming.
-
-    Reproducing #112 end to end needs a MoE architecture whose model_type does not
-    contain "moe" — deepseek_v3 (`n_routed_experts`) is the candidate. Its MLA
-    attention makes that a larger fixture; the profiler itself is covered directly by
-    ModelProfilerMoEDetectionTests in the meantime.
+    """Gemma 4 with a MoE text block, so the expert count is genuinely nested.
+
+    An earlier version of this fixture used `qwen3_moe`, whose Swift configuration is
+    decoded from the root of config.json — so the expert count had to sit at the top
+    level, and the "nested" copy underneath it was never reached. `findExpertCounts`
+    is breadth-first and returned on the root hit, which made both the nesting and the
+    vision decoy dead weight.
+
+    `Gemma4Configuration` decodes `text_config` and nothing else, so here the count
+    exists *only* one level down. That makes two things real rather than decorative:
+    the nested walk added in #112's review follow-up, and the rule that a count under
+    `vision_config` must not be mistaken for the language model's.
+
+    It also covers the fused-expert remap: real gemma4 checkpoints ship
+    `experts.gate_up_proj` as one tensor that sanitize splits in half into
+    `switch_glu.gate_proj` / `switch_glu.up_proj`. A wrong split axis, or swapped
+    halves, is a silent numerical fault no unit test here would see.
     """
     H, L, HEADS, KVH, HD = 64, 2, 4, 2, 16
     INTER, MOE_INTER, EXPERTS, TOPK = 128, 32, 4, 2
+    PLI, VPLI = 32, 16
     rng = np.random.default_rng(0)
 
     text_config = {
-        "model_type": "qwen3_moe",
+        "model_type": "gemma4_text",
         "hidden_size": H,
         "num_hidden_layers": L,
         "intermediate_size": INTER,
         "num_attention_heads": HEADS,
         "num_key_value_heads": KVH,
         "head_dim": HD,
-        "num_experts": EXPERTS,
-        "num_experts_per_tok": TOPK,
-        "moe_intermediate_size": MOE_INTER,
-        "decoder_sparse_step": 1,
-        "mlp_only_layers": [],
+        "global_head_dim": HD,
         "rms_norm_eps": 1e-6,
         "vocab_size": VOCAB,
+        "rope_traditional": False,
         "rope_theta": 10000.0,
-        "tie_word_embeddings": False,
+        "sliding_window": 128,
+        "sliding_window_pattern": 1,
         "max_position_embeddings": 512,
-        "norm_topk_prob": True,
+        "num_kv_shared_layers": 0,
+        "use_double_wide_mlp": False,
+        "tie_word_embeddings": True,
+        "hidden_size_per_layer_input": PLI,
+        "vocab_size_per_layer_input": VPLI,
+        "final_logit_softcapping": 30.0,
+        "attention_k_eq_v": False,
+        "enable_moe_block": True,
+        "num_experts": EXPERTS,
+        "top_k_experts": TOPK,
+        "moe_intermediate_size": MOE_INTER,
+    }
+    # num_experts appears here and nowhere else at the root — that is the point.
+    cfg = {
+        "model_type": "gemma4",
+        "architectures": ["Gemma4ForConditionalGeneration"],
+        "vocab_size": VOCAB,
+        "text_config": text_config,
+        # A decoy the language-model walk has to skip.
+        "vision_config": {"model_type": "gemma4_vision", "num_experts": 999},
     }
-    cfg = dict(text_config)
-    cfg["text_config"] = text_config
-    # Must not be mistaken for the language model's expert count.
-    cfg["vision_config"] = {"model_type": "qwen3_vl", "num_experts": 999}
     json.dump(cfg, open(os.path.join(out, "config.json"), "w"), indent=2)
 
     w = {
-        "model.embed_tokens.weight": rand(rng, VOCAB, H),
-        "model.norm.weight": ones(H),
-        "lm_head.weight": rand(rng, VOCAB, H),
+        "language_model.model.embed_tokens.weight": rand(rng, VOCAB, H),
+        "language_model.model.norm.weight": ones(H),
+        "language_model.model.embed_tokens_per_layer.weight": rand(rng, VPLI, L * PLI),
+        "language_model.model.per_layer_model_projection.weight": rand(rng, L * PLI, H),
+        "language_model.model.per_layer_projection_norm.weight": ones(PLI),
     }
     for i in range(L):
-        p_ = f"model.layers.{i}"
+        # A gemma4 wrapper checkpoint prefixes its text weights this way; matches
+        # what a real gemma-4-e2b ships.
+        p_ = f"language_model.model.layers.{i}"
         w[f"{p_}.self_attn.q_proj.weight"] = rand(rng, HEADS * HD, H)
         w[f"{p_}.self_attn.k_proj.weight"] = rand(rng, KVH * HD, H)
         w[f"{p_}.self_attn.v_proj.weight"] = rand(rng, KVH * HD, H)
         w[f"{p_}.self_attn.o_proj.weight"] = rand(rng, H, HEADS * HD)
         w[f"{p_}.self_attn.q_norm.weight"] = ones(HD)
         w[f"{p_}.self_attn.k_norm.weight"] = ones(HD)
+        w[f"{p_}.layer_scalar"] = np.ones(1, np.float16)
+        w[f"{p_}.mlp.gate_proj.weight"] = rand(rng, INTER, H)
+        w[f"{p_}.mlp.up_proj.weight"] = rand(rng, INTER, H)
+        w[f"{p_}.mlp.down_proj.weight"] = rand(rng, H, INTER)
         w[f"{p_}.input_layernorm.weight"] = ones(H)
         w[f"{p_}.post_attention_layernorm.weight"] = ones(H)
-        w[f"{p_}.mlp.gate.weight"] = rand(rng, EXPERTS, H)
-        # Per-expert tensors, the layout a real checkpoint ships; sanitize stacks
-        # these into switch_mlp.
-        for e in range(EXPERTS):
-            w[f"{p_}.mlp.experts.{e}.gate_proj.weight"] = rand(rng, MOE_INTER, H)
-            w[f"{p_}.mlp.experts.{e}.up_proj.weight"] = rand(rng, MOE_INTER, H)
-            w[f"{p_}.mlp.experts.{e}.down_proj.weight"] = rand(rng, H, MOE_INTER)
+        w[f"{p_}.pre_feedforward_layernorm.weight"] = ones(H)
+        w[f"{p_}.post_feedforward_layernorm.weight"] = ones(H)
+        w[f"{p_}.per_layer_input_gate.weight"] = rand(rng, PLI, H)
+        w[f"{p_}.per_layer_projection.weight"] = rand(rng, H, PLI)
+        w[f"{p_}.post_per_layer_input_norm.weight"] = ones(H)
+        # Router and experts are siblings of mlp on the decoder layer, not nested in it.
+        # Enabling the MoE block also builds a second pair of feedforward norms.
+        w[f"{p_}.pre_feedforward_layernorm_2.weight"] = ones(H)
+        w[f"{p_}.post_feedforward_layernorm_1.weight"] = ones(H)
+        w[f"{p_}.post_feedforward_layernorm_2.weight"] = ones(H)
+        w[f"{p_}.router.proj.weight"] = rand(rng, EXPERTS, H)
+        w[f"{p_}.router.scale"] = ones(H)
+        w[f"{p_}.router.per_expert_scale"] = ones(EXPERTS)
+        # Fused, as a real checkpoint ships it: sanitize splits dim -2 in half.
+        w[f"{p_}.experts.gate_up_proj"] = rand(rng, EXPERTS, 2 * MOE_INTER, H)
+        w[f"{p_}.experts.down_proj"] = rand(rng, EXPERTS, H, MOE_INTER)
 
     save_file(w, os.path.join(out, "model.safetensors"), metadata={"format": "pt"})
     return len(w)
@@ -319,7 +350,16 @@ def build_moe_nested(out):
     total = 0
     for name, (fn, kwargs) in FIXTURES.items():
         out = os.path.join(ROOT, name)
-        os.makedirs(out, exist_ok=True)
+        # Clear this fixture's own directory, and only its own. Writing over an existing
+        # one leaves behind files the current builder no longer emits — flip stray-shard
+        # off and its decoy shard and index survive, so the fixture keeps testing a shape
+        # the source no longer describes, and the printed size counts files that are not
+        # part of it. Scoping the removal to one known fixture directory is also what
+        # keeps regeneration from reaching siblings such as tests/fixtures/omni, whose
+        # assets belong to test-omni.sh and are not generated here.
+        if os.path.isdir(out):
+            shutil.rmtree(out)
+        os.makedirs(out)
         n = fn(out, **kwargs)
         write_tokenizer(out)
         size = sum(os.path.getsize(os.path.join(out, f)) for f in os.listdir(out))

tests/fixtures/moe-nested/config.json

@@ -1,44 +1,39 @@
 {
-  "model_type": "qwen3_moe",
-  "hidden_size": 64,
-  "num_hidden_layers": 2,
-  "intermediate_size": 128,
-  "num_attention_heads": 4,
-  "num_key_value_heads": 2,
-  "head_dim": 16,
-  "num_experts": 4,
-  "num_experts_per_tok": 2,
-  "moe_intermediate_size": 32,
-  "decoder_sparse_step": 1,
-  "mlp_only_layers": [],
-  "rms_norm_eps": 1e-06,
+  "model_type": "gemma4",
+  "architectures": [
+    "Gemma4ForConditionalGeneration"
+  ],
   "vocab_size": 288,
-  "rope_theta": 10000.0,
-  "tie_word_embeddings": false,
-  "max_position_embeddings": 512,
-  "norm_topk_prob": true,
   "text_config": {
-    "model_type": "qwen3_moe",
+    "model_type": "gemma4_text",
     "hidden_size": 64,
     "num_hidden_layers": 2,
     "intermediate_size": 128,
     "num_attention_heads": 4,
     "num_key_value_heads": 2,
     "head_dim": 16,
-    "num_experts": 4,
-    "num_experts_per_tok": 2,
-    "moe_intermediate_size": 32,
-    "decoder_sparse_step": 1,
-    "mlp_only_layers": [],
+    "global_head_dim": 16,
     "rms_norm_eps": 1e-06,
     "vocab_size": 288,
+    "rope_traditional": false,
     "rope_theta": 10000.0,
-    "tie_word_embeddings": false,
+    "sliding_window": 128,
+    "sliding_window_pattern": 1,
     "max_position_embeddings": 512,
-    "norm_topk_prob": true
+    "num_kv_shared_layers": 0,
+    "use_double_wide_mlp": false,
+    "tie_word_embeddings": true,
+    "hidden_size_per_layer_input": 32,
+    "vocab_size_per_layer_input": 16,
+    "final_logit_softcapping": 30.0,
+    "attention_k_eq_v": false,
+    "enable_moe_block": true,
+    "num_experts": 4,
+    "top_k_experts": 2,
+    "moe_intermediate_size": 32
   },
   "vision_config": {
-    "model_type": "qwen3_vl",
+    "model_type": "gemma4_vision",
     "num_experts": 999
   }
 }

tests/test-fixtures.sh

@@ -10,9 +10,9 @@
 #   stray-shard         #118: a .safetensors beside the index but absent from it
 #   kv-shared-absent    #120: gemma-4-e4b shape, shared layers ship no k/v
 #   kv-shared-present   b674:  gemma-4-e2b shape, shared layers ship k/v anyway
-#   moe-nested          MoE load path: per-expert tensors stacked into switch_mlp,
-#                       expert count nested under text_config. Note this does not
-#                       reproduce #112 — see scripts/make-test-fixtures.py for why.
+#   moe-nested          #112: expert count nested under text_config only, with a
+#                       decoy count under vision_config; plus the fused
+#                       experts.gate_up_proj split that sanitize performs
 #
 # The output is gibberish by construction — the weights are random. A fixture passes
 # when the server loads it and produces *a* token, which is what exercises config
@@ -35,7 +35,16 @@ pass() { PASS=$((PASS + 1)); echo -e "  ${GREEN}✅ PASS${NC}: $*"; }
 fail() { FAIL=$((FAIL + 1)); echo -e "  ${RED}❌ FAIL${NC}: $*"; }
 
 SERVER_PID=""
-cleanup() { [ -n "$SERVER_PID" ] && kill "$SERVER_PID" 2>/dev/null; SERVER_PID=""; }
+# `kill` only asks. Without waiting for the process to go, the next fixture binds the
+# same port while the previous server may still hold it — see the readiness loop below
+# for why that is worse than a flake.
+cleanup() {
+    if [ -n "$SERVER_PID" ]; then
+        kill "$SERVER_PID" 2>/dev/null
+        wait "$SERVER_PID" 2>/dev/null
+    fi
+    SERVER_PID=""
+}
 trap cleanup EXIT
 
 run_fixture() {
@@ -52,10 +61,14 @@ run_fixture() {
     "$BINARY" --model "$dir" --port "$PORT" --host "$HOST" > "$logfile" 2>&1 &
     SERVER_PID=$!
 
+    # Liveness is checked *before* the health probe on purpose. The other order lets a
+    # server that failed to bind (because the previous one still held the port) pass as
+    # ready on a probe answered by that previous server — the assertions then run
+    # against the wrong checkpoint and report a false pass rather than a failure.
     local ready=0
     for _ in $(seq 1 60); do
-        if curl -sf "$url/health" >/dev/null 2>&1; then ready=1; break; fi
         if ! kill -0 "$SERVER_PID" 2>/dev/null; then break; fi
+        if curl -sf "$url/health" >/dev/null 2>&1; then ready=1; break; fi
         sleep 1
     done
 
@@ -97,29 +110,36 @@ for name in dense stray-shard kv-shared-absent kv-shared-present moe-nested; do
     run_fixture "$name"
 done
 
-# The MoE fixture again, this time with --stream-experts. A model classified dense
-# has the flag silently dropped and gets materialised whole, which is how #112
-# reached an OOM kill rather than an error message; asserting the flag was honoured
-# is the closest end-to-end check available.
-log "Shape: moe-nested (--stream-experts honoured)"
+# The MoE fixture again with --stream-experts, to assert the *config-level* gate.
+#
+# #112: expert counts are spelled and nested differently per family, and a model
+# misread as dense had --stream-experts silently dropped and was then materialised
+# whole until the OS killed it. The check is that the config gate does not reject —
+# not that streaming actually engages. Those are two separate gates: gemma4 passes
+# detection but has no StreamableMoE conformance, so the second one legitimately
+# declines. Asserting on the first is what tracks #112.
+#
+# This is a live check rather than a decorative one: `gemma4` contains no "moe", so
+# the model_type fallback in modelTypeImpliesMoE cannot rescue it, and the count
+# exists only inside text_config. Reverting detection to a top-level single-key form
+# makes this fail.
+log "Shape: moe-nested (nested expert count is detected)"
 MOE_LOG="/tmp/SwiftLM-test-fixture-moe-stream.log"
 "$BINARY" --model "$FIXTURE_DIR/moe-nested" --port "$PORT" --host "$HOST" \
     --stream-experts > "$MOE_LOG" 2>&1 &
 SERVER_PID=$!
 ready=0
 for _ in $(seq 1 60); do
-    curl -sf "http://$HOST:$PORT/health" >/dev/null 2>&1 && { ready=1; break; }
     kill -0 "$SERVER_PID" 2>/dev/null || break
+    curl -sf "http://$HOST:$PORT/health" >/dev/null 2>&1 && { ready=1; break; }
     sleep 1
 done
 if [ "$ready" -ne 1 ]; then
     fail "moe-nested did not start with --stream-experts"
-elif grep -qi "is not MoE" "$MOE_LOG"; then
-    fail "moe-nested was classified dense: $(grep -i 'is not MoE' "$MOE_LOG" | head -1)"
-elif grep -q "SSD Expert Streaming enabled" "$MOE_LOG"; then
-    pass "moe-nested detected as MoE, --stream-experts honoured"
+elif grep -q "is not MoE" "$MOE_LOG"; then
+    fail "nested expert count missed: $(grep 'is not MoE' "$MOE_LOG" | head -1)"
 else
-    fail "moe-nested: no streaming confirmation in log"
+    pass "moe-nested: expert count found under text_config, vision decoy ignored"
 fi
 cleanup