| """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() |
| |
| |
| 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) |
|
|
