fix(staging): derive klein-9b hidden/heads from context_embedder when img_in/x_embedder absent
The klein-9b-50x-lighting (FP4) export ships NEITHER img_in NOR x_embedder (only x_embedder._wscales, a 1-D scale) — unique among klein tiers (the 4B-50x, 9B-40x, 9B-base-50x, and all 30x keep x_embedder). So in derive_transformer_config step 2, _first_nonempty_shape(keys, hidden_from=[img_in.*, x_embedder.*]) returns None → num_attention_heads is NOT derived → the bundled 4B template's 24 sticks → hidden = 24*128 = 3072 → the engine builds a 4B-sized transformer → loading the 9B weights overflows: norm_out.norm.bias dst[3072] != src[4096] → SIGSEGV/crash at load (reported on RTX 5090, worker 0.0.12, method=prequant_lighting_separate). Fix: add a HIDDEN-ONLY fallback source `hidden_extra_from` (context_embedder / norm_out.norm.* dim-0), used ONLY when hidden_from is absent and ONLY for hidden → num_attention_heads — NEVER in_channels (context_embedder dim-1 is JOINT, not in_channels; corrupting in_channels would break the patch embed). in_channels stays from the bundled config (4B/9B share it = 128). Applied to Flux2Klein + QwenImage (shared key layout / same fallthrough class). Verified (GPU-free staging A/B on the real safetensors headers): klein-9b-50x : OLD final config heads=24 (hidden 3072) -> crash ; FIXED heads=32 (4096) PASS klein-9b-30x : heads=32 unchanged (has x_embedder) klein-4b-50x : heads=24 unchanged (has x_embedder) klein-4b-30x : heads=24 unchanged in_channels=128 uncorrupted in all cases; joint=12288/7680 correct. Engine loads a 9B with heads=32 (proven separately) and the user confirmed heads=32 loads+runs on their 5090 (the known-good target). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Y
yejianhui committed
93d2d0d475186f8494814136567962616ffb9714
Parent: 169c654