diff --git a/tests/unit/test_part_a_decode_diagnostics.py b/tests/unit/test_part_a_decode_diagnostics.py index b7c8ed3f..5990b6d1 100644 --- a/tests/unit/test_part_a_decode_diagnostics.py +++ b/tests/unit/test_part_a_decode_diagnostics.py @@ -195,6 +195,13 @@ def main(): dev = f"cuda:{li % NUM_GPUS}" if li in compressors: compressors[li].load(all_w, pfx, dev=dev) if li in indexers: indexers[li].load(all_w, f"{pfx}.indexer", dev=dev) + # Verify compressor kv_norm_w loaded correctly + for li in range(TEST_LAYERS): + if li in compressors and compressors[li].kv_norm_w is not None: + n = compressors[li].kv_norm_w + print(f" L{li} compressor kv_norm_w: shape={tuple(n.shape)} |w|={n.abs().max().item():.4f}", flush=True) + elif li in compressors: + print(f" L{li} compressor kv_norm_w: MISSING!", flush=True) print("Production components built") # Embedding + tokenizer