| """scripts/p2/check_qwen3_snapshot.py on a tiny fake HF snapshot (CPU, header-only checks). |
| |
| Run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_check_qwen3_snapshot.py -q |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
| from safetensors.torch import save_file |
|
|
| from align.model.qwen3 import Qwen3Config |
| from align.model.qwen3_checkpoint import _expected_shapes |
| from test_p2_support import load_script |
|
|
| CK = load_script("p2_check_qwen3_snapshot", "scripts/p2/check_qwen3_snapshot.py") |
| TINY = dict(model_type="qwen3", architectures=["Qwen3ForCausalLM"], vocab_size=32, hidden_size=16, intermediate_size=24, num_hidden_layers=2, |
| num_attention_heads=4, num_key_value_heads=2, head_dim=8, rms_norm_eps=1e-6, rope_theta=1e6, max_position_embeddings=64, |
| hidden_act="silu", attention_bias=False, tie_word_embeddings=False, sliding_window=None, use_sliding_window=False) |
|
|
|
|
| def make(d: Path, drop=None, extra=None, bad_shape=None, cfg_over=None, truncate=False): |
| d.mkdir(parents=True, exist_ok=True) |
| cfg = {**TINY, **(cfg_over or {})} |
| (d / "config.json").write_text(json.dumps(cfg)) |
| shapes = _expected_shapes(Qwen3Config.from_hf_config(TINY)) |
| state = {k: torch.zeros(s, dtype=torch.bfloat16) for k, s in shapes.items()} |
| if drop: |
| state.pop(drop) |
| if extra: |
| state[extra] = torch.zeros(2, dtype=torch.bfloat16) |
| if bad_shape: |
| state[bad_shape] = torch.zeros(3, 3, dtype=torch.bfloat16) |
| keys = sorted(state) |
| half = len(keys) // 2 |
| wm = {} |
| for i, part in enumerate((keys[:half], keys[half:])): |
| name = f"model-0000{i + 1}-of-00002.safetensors" |
| save_file({k: state[k] for k in part}, str(d / name)) |
| wm.update({k: name for k in part}) |
| (d / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": wm})) |
| if truncate: |
| p = d / "model-00002-of-00002.safetensors" |
| p.write_bytes(p.read_bytes()[:-10]) |
| return d |
|
|
|
|
| def test_a_complete_qwen3_snapshot_is_compatible(tmp_workdir): |
| rep = CK.check(str(make(Path(tmp_workdir) / "ok"))) |
| assert rep["compatible"] and rep["problems"] == [] and rep["dtypes"] == {"BF16": rep["index"]["tensors"]} |
|
|
|
|
| @pytest.mark.parametrize("kw,needle", [ |
| (dict(drop="model.layers.1.mlp.up_proj.weight"), "absent from the index"), |
| (dict(extra="model.layers.0.self_attn.q_proj.bias"), "cannot map"), |
| (dict(bad_shape="model.norm.weight"), "implied by the config"), |
| (dict(truncate=True), "truncated"), |
| ]) |
| def test_each_defect_is_named(tmp_workdir, kw, needle): |
| rep = CK.check(str(make(Path(tmp_workdir) / "bad", **kw))) |
| assert not rep["compatible"] and any(needle in p for p in rep["problems"]), rep["problems"] |
|
|
|
|
| def test_unsupported_config_variants_are_refused(tmp_workdir): |
| assert CK.main([str(make(Path(tmp_workdir) / "sw", cfg_over={"use_sliding_window": True, "sliding_window": 8}))]) == 1 |
| assert CK.main([str(make(Path(tmp_workdir) / "ok2"))]) == 0 |
|
|
