diff --git a/scripts/make-test-fixtures.py b/scripts/make-test-fixtures.py index 953f437..ee2e48c 100755 --- a/scripts/make-test-fixtures.py +++ b/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)) diff --git a/tests/fixtures/moe-nested/config.json b/tests/fixtures/moe-nested/config.json index bfee7eb..8dc58d1 100644 --- a/tests/fixtures/moe-nested/config.json +++ b/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 } } \ No newline at end of file diff --git a/tests/fixtures/moe-nested/model.safetensors b/tests/fixtures/moe-nested/model.safetensors index fab58e1..984e8d9 100644 Binary files a/tests/fixtures/moe-nested/model.safetensors and b/tests/fixtures/moe-nested/model.safetensors differ diff --git a/tests/test-fixtures.sh b/tests/test-fixtures.sh index 92f2248..7a26f2f 100755 --- a/tests/test-fixtures.sh +++ b/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