"""CPU proof: exported half-depth uses original layers, norm, and readout.""" import importlib.util import json from pathlib import Path import pytest import torch from safetensors import safe_open from safetensors.torch import save_file from transformers import ( Qwen3_5Config, Qwen3_5ForConditionalGeneration, Qwen3_5TextConfig, ) spec = importlib.util.spec_from_file_location( "export_low", Path(__file__).parents[1] / "export_low.py" ) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) def test_actual_32_to_16_prefix_logits_and_saved_tensors(tmp_path): torch.set_num_threads(1) text = Qwen3_5TextConfig( vocab_size=64, hidden_size=32, intermediate_size=48, num_hidden_layers=32, num_attention_heads=2, num_key_value_heads=1, head_dim=16, linear_key_head_dim=8, linear_value_head_dim=8, linear_num_key_heads=2, linear_num_value_heads=2, layer_types=(["linear_attention"] * 3 + ["full_attention"]) * 8, rope_parameters={ "rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 1.0, "mrope_section": [2, 3, 3], }, ) config = Qwen3_5Config( text_config=text, vision_config={ "depth": 1, "hidden_size": 32, "intermediate_size": 48, "num_heads": 2, "out_hidden_size": 32, "num_position_embeddings": 16, }, ) config._attn_implementation = "eager" config.text_config._attn_implementation = "eager" config.text_config.num_hidden_layers = 32 config.text_config.layer_types = (["linear_attention"] * 3 + ["full_attention"]) * 8 config.tie_word_embeddings = False config.text_config.tie_word_embeddings = False torch.manual_seed(47) full = Qwen3_5ForConditionalGeneration(config).eval() state = full.state_dict() source, output = tmp_path / "source", tmp_path / "low" source.mkdir() output.mkdir() save_file(state, source / "model.safetensors", metadata={"format": "pt"}) weight_map = {k: "model.safetensors" for k in state} kept, removed = module.select_keys(weight_map, 16, 32) assert len(removed) > 0 assert module.READOUT <= set(kept) assert all( not module.LAYER.match(k) or int(module.LAYER.match(k).group(1)) < 16 for k in kept ) files, tensors = module.copy_weights(source, output, weight_map, kept) short_config = module.prefix_config(config.to_dict()) assert short_config["vision_config"] == config.to_dict()["vision_config"] assert short_config["text_config"]["num_hidden_layers"] == 16 assert len(short_config["text_config"]["layer_types"]) == 16 loaded = {} for f in files: with safe_open(output / f["file"], framework="pt") as handle: loaded.update({key: handle.get_tensor(key) for key in handle.keys()}) low_config = Qwen3_5Config.from_dict(short_config) low_config._attn_implementation = "eager" low_config.text_config._attn_implementation = "eager" low = Qwen3_5ForConditionalGeneration(low_config).eval() low.load_state_dict(loaded, strict=True) assert len(low.model.language_model.layers) == 16 assert len(tensors) == len(loaded) for key in module.READOUT: torch.testing.assert_close(loaded[key], state[key], rtol=0, atol=0) ids, candidates = torch.tensor([[4, 5, 6, 7]]), torch.tensor([10, 11, 12]) with torch.no_grad(): captured = [] hook = full.model.language_model.layers[15].register_forward_hook( lambda _module, _inputs, output: captured.append(output.detach().clone()) ) try: full(input_ids=ids, use_cache=False) finally: hook.remove() # Independently capture the full model's layer-16 residual, then apply # its original final norm and head. No monorepo runtime dependency. native = full.lm_head( full.model.language_model.norm(captured[0][:, -1]) ).index_select(-1, candidates) exported = ( low(input_ids=ids, use_cache=False) .logits[:, -1] .index_select(-1, candidates) ) torch.testing.assert_close(native, exported, atol=1e-6, rtol=1e-5) assert ( json.loads((output / "model.safetensors.index.json").read_text())["metadata"][ "total_size" ] > 0 ) @pytest.mark.parametrize("missing", sorted(module.READOUT)) def test_readout_cannot_be_dropped(missing): keys = (module.READOUT - {missing}) | { f"model.language_model.layers.{i}.weight" for i in range(32) } with pytest.raises(ValueError, match="missing"): module.select_keys(keys, 16, 32) def test_unknown_adapter_namespace_rejected(): keys = ( module.READOUT | {f"model.language_model.layers.{i}.weight" for i in range(32)} | {"readout.adapter.weight"} ) with pytest.raises(ValueError, match="unexpected"): module.select_keys(keys, 16, 32) def test_existing_output_never_overwritten(tmp_path): with pytest.raises(ValueError, match="new directory"): module.export(tmp_path / "unused", tmp_path) def test_source_stat_change_rejects_verified_publication(tmp_path): path = tmp_path / "weights" path.write_bytes(b"original") snapshot = module.source_snapshot(tmp_path) module.assert_source_unchanged(tmp_path, snapshot) path.write_bytes(b"changed-longer") with pytest.raises(ValueError, match="changed during export"): module.assert_source_unchanged(tmp_path, snapshot) def test_launcher_full_low_weight_hash_detects_tampering(tmp_path): config = { "model_type": "qwen3_5", "text_config": { "hidden_size": 4096, "vocab_size": 248320, "num_hidden_layers": 16, "layer_types": ["linear_attention"] * 16, }, } (tmp_path / "config.json").write_text(json.dumps(config)) keys = module.READOUT | { f"model.language_model.layers.{i}.weight" for i in range(16) } (tmp_path / "model.safetensors.index.json").write_text( json.dumps({"weight_map": {key: "model.safetensors" for key in keys}}) ) shard = tmp_path / "model.safetensors" shard.write_bytes(b"original-weights") manifest = { "status": "verified", "actual_depth": 16, "all_retained_tensors_bit_identical": True, "export_config_sha256": module.file_sha(tmp_path / "config.json"), "export_index_sha256": module.file_sha( tmp_path / "model.safetensors.index.json" ), "export_weight_files": [ { "file": shard.name, "bytes": shard.stat().st_size, "sha256": module.file_sha(shard), } ], } (tmp_path / "export-low-provenance.json").write_text(json.dumps(manifest)) module.verify_artifact(tmp_path, 16) shard.write_bytes(b"tampered-weights") with pytest.raises(ValueError, match="weight SHA mismatch"): module.verify_artifact(tmp_path, 16) def test_launcher_high_release_weight_hash_detects_tampering(tmp_path): config = { "model_type": "qwen3_5", "architectures": ["Qwen3_5ForConditionalGeneration"], "tie_word_embeddings": False, "text_config": { "hidden_size": 4096, "vocab_size": 248320, "intermediate_size": 12288, "num_hidden_layers": 32, "tie_word_embeddings": False, "layer_types": (["linear_attention"] * 3 + ["full_attention"]) * 8, }, } (tmp_path / "config.json").write_text(json.dumps(config)) (tmp_path / "depth_config.json").write_text( json.dumps({"exit_depth": 16, "full_depth": 32}) ) for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): (tmp_path / name).write_text("mock tokenization asset") keys = module.READOUT | { f"model.language_model.layers.{i}.weight" for i in range(32) } (tmp_path / "model.safetensors.index.json").write_text( json.dumps({"weight_map": {key: "model.safetensors" for key in keys}}) ) shard = tmp_path / "model.safetensors" shard.write_bytes(b"original-weights") entries = [ {"path": p.name, "sha256": module.file_sha(p)} for p in tmp_path.iterdir() ] (tmp_path / "release-manifest.json").write_text(json.dumps({"files": entries})) module.verify_artifact(tmp_path, 32) shard.write_bytes(b"tampered-weights") with pytest.raises(ValueError, match="weight SHA mismatch"): module.verify_artifact(tmp_path, 32)