code/tests/test_checkpoint_roundtrip.py

align / code /tests /test_checkpoint_roundtrip.py
Add files using upload-large-folder tool
201dc77 verified 16 days ago
33.2 kB
"""Trained-checkpoint round-trip gate: nmo -> safetensors -> load_hf_safetensors
-> align.model reload -> logits. Interp (section-08) and eval (section-10) load
checkpoints through this layout, and outside tools load the same files, so a
mapping defect here silently invalidates the entire traceability pillar.
No transformers in the default run: the reload goes through the repo's own
`hf_config.read_config` + `load_hf_safetensors` + `NmoMistralForCausalLM`.
The `oracle`-marked tests at the bottom re-tie the on-disk layout to real
transformers (legacy venv only) -- that is what "stays HF-loadable by outside
tools" means as a checkable claim.
Run: uv run pytest tests/test_checkpoint_roundtrip.py -v -m "not gpu and not oracle"
Oracle: PYTHONPATH=/apps/development/rlhf/align \\
        /apps/development/rlhf/nano-mistral-o1/.venv/bin/python -m pytest \\
        tests/test_checkpoint_roundtrip.py -v -m oracle
"""

from __future__ import annotations

import ast
import glob
import json
import os
import subprocess
import sys
from pathlib import Path

import pytest
import torch
from test_model_parity import tiny_config  # tests/ is on sys.path under pytest

from align.config import ModelConfig
from align.model import checkpoint as ckpt
from align.model import hf_config as hfc
from align.model.mistral import NmoMistralForCausalLM

REPO = Path(__file__).resolve().parents[1]

#: The architecture fields that must survive any translation unchanged.
ARCH_FIELDS = (
    "n_layers",
    "hidden_size",
    "intermediate_size",
    "n_heads",
    "n_kv_heads",
    "head_dim",
    "vocab_size",
    "max_position_embeddings",
    "rms_norm_eps",
    "rope_theta",
    "sliding_window",
    "attention_dropout",
    "tie_word_embeddings",
)

#: Recorded copy of Mistral-7B-v0.3's upstream config.json (the 3 shards'
#: sibling; the file the `V03_FACTS` table and `ModelConfig()` claim to be).
#: Kept in-repo because the local snapshot is not guaranteed to be on every
#: box that runs CI; when the snapshot IS present the test below checks the
#: recorded copy against it, so this cannot rot into a second source of truth.
#: Note what is ABSENT: head_dim and attention_dropout (v0.3 omits both).
V03_UPSTREAM_CONFIG = {
    "architectures": ["MistralForCausalLM"],
    "bos_token_id": 1,
    "eos_token_id": 2,
    "hidden_act": "silu",
    "hidden_size": 4096,
    "initializer_range": 0.02,
    "intermediate_size": 14336,
    "max_position_embeddings": 32768,
    "model_type": "mistral",
    "num_attention_heads": 32,
    "num_hidden_layers": 32,
    "num_key_value_heads": 8,
    "rms_norm_eps": 1e-05,
    "rope_theta": 1000000.0,
    "sliding_window": None,
    "tie_word_embeddings": False,
    "torch_dtype": "bfloat16",
    "transformers_version": "4.42.0.dev0",
    "use_cache": True,
    "vocab_size": 32768,
}

V03_SNAPSHOT_GLOB = (
    "/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/snapshots/*"
)
#: A config.json that transformers 5.x itself wrote from the real v0.3 weights
#: (SFT checkpoint of the legacy lineage). A second, independent real source.
LEGACY_EMITTED_GLOB = (
    "/apps/development/rlhf/nano-mistral-o1/checkpoints/sft/checkpoints/*/config.json"
)


@pytest.fixture
def tiny(tiny_model_constants):
    return tiny_config(tiny_model_constants)


@pytest.fixture
def trained(tiny, tiny_model_constants):
    """A tiny nmo model with every parameter perturbed away from init -- the
    stand-in for a trained model (no tensor left at its init value)."""
    torch.manual_seed(tiny_model_constants["seed"] + 9)
    model = NmoMistralForCausalLM(tiny).eval()
    with torch.no_grad():
        for p in model.parameters():
            p.add_(torch.randn_like(p) * 0.05)
    return model


@pytest.fixture
def hf_sd(trained):
    """An HF-layout state dict (what an outside tool's model.state_dict() is)."""
    return ckpt.nmo_to_hf(trained.state_dict())


def _reload(out_dir: str, attn: str = "eager") -> NmoMistralForCausalLM:
    """The consumer path: config.json + safetensors -> align model."""
    cfg = hfc.read_config(out_dir)
    return NmoMistralForCausalLM.from_state_dict(
        cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(out_dir), expected_n_layers=cfg.n_layers),
        attn_implementation=attn,
    ).eval()


def test_hf_to_nmo_to_safetensors_to_reload_preserves_all_tensors(tiny, hf_sd, tmp_workdir):
    """HF-layout sd -> nmo -> save safetensors in HF layout -> reload through
    load_hf_safetensors. EVERY tensor preserved: same key set (no missing, no
    unexpected), same shape, same dtype, bitwise-equal values."""
    out = os.path.join(tmp_workdir, "ckpt")
    ckpt.save_nmo_as_hf(ckpt.hf_to_nmo(hf_sd), tiny, out)

    got = ckpt.load_hf_safetensors(out)
    assert set(got) == set(hf_sd), (
        f"missing={sorted(set(hf_sd) - set(got))} unexpected={sorted(set(got) - set(hf_sd))}"
    )
    for k in hf_sd:
        assert got[k].shape == hf_sd[k].shape, k
        assert got[k].dtype == hf_sd[k].dtype, k
        assert torch.equal(got[k], hf_sd[k]), f"{k} not bitwise-equal"


def test_roundtrip_logits_match(tiny, trained, hf_sd, tmp_workdir):
    """The small-scale trained-checkpoint round-trip gate, both spellings: with
    an explicit reference model, and with the default self-built reference
    (an in-memory model that never touched disk)."""
    nmo_sd = ckpt.hf_to_nmo(hf_sd)

    # Same backend as the gate's reload (eager): a bitwise comparison across
    # sdpa and eager would fail by ~1e-7 for reasons that are not a defect.
    same_backend = NmoMistralForCausalLM.from_state_dict(
        tiny, {k: v.clone() for k, v in trained.state_dict().items()}, "eager"
    ).eval()
    report = ckpt.roundtrip_gate(
        nmo_sd, tiny, os.path.join(tmp_workdir, "with_ref"), reference_model=same_backend
    )
    assert report["tensors_compared"] == len(hf_sd)
    assert report["max_abs_dlogit"] == 0.0, report
    assert report["config_mismatch"] == []
    assert report["passed"] is True

    report = ckpt.roundtrip_gate(nmo_sd, tiny, os.path.join(tmp_workdir, "no_ref"))
    assert report["max_abs_dlogit"] == 0.0, report
    assert report["passed"] is True


def test_roundtrip_gate_has_teeth(tiny, hf_sd, tmp_workdir, monkeypatch):
    """A gate that cannot fail is not a gate. Break each thing it claims to
    check, one at a time, and require passed=False with the right field set:
    (a) a tensor corrupted in serialisation, (b) a config.json that states the
    wrong rope_theta (reloads cleanly and is a different model), (c) a
    reference model that disagrees."""
    nmo_sd = ckpt.hf_to_nmo(hf_sd)

    # (a) tensor corrupted between save and load
    real_load = ckpt.load_hf_safetensors

    def corrupting_load(path):
        sd = dict(real_load(path))
        sd["model.layers.1.mlp.down_proj.weight"] = sd["model.layers.1.mlp.down_proj.weight"] + 1.0
        return sd

    monkeypatch.setattr(ckpt, "load_hf_safetensors", corrupting_load)
    rep = ckpt.roundtrip_gate(nmo_sd, tiny, os.path.join(tmp_workdir, "a"))
    assert rep["passed"] is False
    assert rep["mismatched_tensors"] == ["model.layers.1.mlp.down_proj.weight"]
    assert rep["max_abs_dlogit"] > 0.0
    monkeypatch.undo()

    # (b) config.json writer states the wrong theta
    real_to_dict = hfc.to_hf_config_dict

    def wrong_theta(cfg, **kw):
        d = real_to_dict(cfg, **kw)
        d["rope_theta"] = 10000.0
        d["rope_parameters"] = {"rope_theta": 10000.0, "rope_type": "default"}
        return d

    monkeypatch.setattr(hfc, "to_hf_config_dict", wrong_theta)
    rep = ckpt.roundtrip_gate(nmo_sd, tiny, os.path.join(tmp_workdir, "b"))
    assert rep["passed"] is False
    assert rep["config_mismatch"] == ["rope_theta"]
    assert rep["max_abs_dlogit"] > 0.0, "wrong theta must move the logits, not just the config"
    monkeypatch.undo()

    # (c) a reference model that is not the saved one
    other = NmoMistralForCausalLM(tiny).eval()  # fresh init, different weights
    rep = ckpt.roundtrip_gate(nmo_sd, tiny, os.path.join(tmp_workdir, "c"), reference_model=other)
    assert rep["passed"] is False and rep["max_abs_dlogit"] > 0.0


def test_nmo_model_weights_survive_the_roundtrip(tiny, trained, tmp_workdir):
    """The direction that actually ships: a TRAINED nmo model is saved and
    reloaded through the consumer path (read_config + load_hf_safetensors),
    and the reloaded model reproduces the original's own logits -- on both
    attention backends."""
    out = os.path.join(tmp_workdir, "trained")
    ckpt.save_nmo_as_hf(trained.state_dict(), tiny, out)

    ids = torch.randint(0, tiny.vocab_size, (2, 13))
    for attn in ("eager", "sdpa"):
        back = _reload(out, attn)
        with torch.no_grad():
            mine = NmoMistralForCausalLM.from_state_dict(
                tiny, {k: v.clone() for k, v in trained.state_dict().items()}, attn
            )(input_ids=ids)
            theirs = back(input_ids=ids)
        d = (mine - theirs).abs().max().item()
        assert d <= 2e-5, f"{attn}: reloaded vs original max|dlogit| {d:.3e}"


def test_mapping_is_total_and_loud(tiny, hf_sd):
    """Both mapping directions raise on an unknown/missing key rather than
    skipping it (the silent-partial-load failure class: a missing tensor
    silently left at random init is the worst possible outcome).
    Both directions of the guard: the complete state dict maps cleanly, and
    each of the four ways to be incomplete raises.
    """
    nmo_sd = ckpt.hf_to_nmo(hf_sd)
    assert ckpt.nmo_to_hf(nmo_sd).keys() == hf_sd.keys()
    for k, v in ckpt.nmo_to_hf(nmo_sd).items():
        assert torch.equal(v, hf_sd[k])

    # 1. unknown source key, HF -> nmo
    bad = dict(hf_sd)
    bad["model.layers.0.self_attn.rotary_emb.inv_freq"] = torch.zeros(8)
    with pytest.raises(KeyError, match="unmapped"):
        ckpt.hf_to_nmo(bad)

    # 2. missing source key, HF -> nmo (the silent-partial-load shape)
    bad = dict(hf_sd)
    del bad["model.layers.1.mlp.down_proj.weight"]
    with pytest.raises(KeyError, match="missing"):
        ckpt.hf_to_nmo(bad)

    # 3. unknown source key, nmo -> HF
    bad = dict(nmo_sd)
    bad["layers.0.attn.q_norm.weight"] = torch.zeros(8)
    with pytest.raises(KeyError, match="unmapped"):
        ckpt.nmo_to_hf(bad)

    # 4. missing source key, nmo -> HF
    bad = dict(nmo_sd)
    del bad["lm_head.weight"]
    with pytest.raises(KeyError, match="missing"):
        ckpt.nmo_to_hf(bad)


def test_mapping_refuses_a_truncated_layer_stack(tiny, hf_sd):
    """A whole layer vanishing is the failure a per-key table cannot see: every
    key present is mapped, the state dict just describes a shallower model.
    Completeness is asserted against the layer indices actually present, and
    against the config when one is supplied."""
    trunc = {k: v for k, v in hf_sd.items() if not k.startswith("model.layers.1.")}
    with pytest.raises(KeyError, match="missing"):
        ckpt.hf_to_nmo(trunc, expected_n_layers=tiny.n_layers)


def test_hf_key_table_is_pinned_to_the_hf_layout():
    """The key names are the interface to every outside tool. Pin a few
    literals so a 'cleanup' of the table cannot pass by being self-consistent
    (an identity mapping would also round-trip)."""
    hf = ckpt.nmo_to_hf(
        {
            "model.embed_tokens.weight": 0,
            "model.norm.weight": 0,
            "lm_head.weight": 0,
            **{
                f"model.layers.0.{n}.weight": 0
                for n in (
                    "attn.q_proj",
                    "attn.k_proj",
                    "attn.v_proj",
                    "attn.o_proj",
                    "mlp.gate_proj",
                    "mlp.up_proj",
                    "mlp.down_proj",
                    "attn_norm",
                    "mlp_norm",
                )
            },
        }
    )
    assert sorted(hf) == sorted(
        [
            "model.embed_tokens.weight",
            "model.norm.weight",
            "lm_head.weight",
            "model.layers.0.self_attn.q_proj.weight",
            "model.layers.0.self_attn.k_proj.weight",
            "model.layers.0.self_attn.v_proj.weight",
            "model.layers.0.self_attn.o_proj.weight",
            "model.layers.0.mlp.gate_proj.weight",
            "model.layers.0.mlp.up_proj.weight",
            "model.layers.0.mlp.down_proj.weight",
            "model.layers.0.input_layernorm.weight",
            "model.layers.0.post_attention_layernorm.weight",
        ]
    )


def test_save_refuses_paths_off_apps(tiny, hf_sd):
    """Checkpoints on this box go to /apps and nowhere else. / has ~13GB free
    and /tmp is a noexec tmpfs; a 13.5GB checkpoint written there takes the
    machine down, and the guard must fire BEFORE the first byte is written."""
    with pytest.raises(RuntimeError, match="/apps"):
        ckpt.save_nmo_as_hf(ckpt.hf_to_nmo(hf_sd), tiny, "/tmp/nmo_ckpt")


def test_save_leaves_no_partial_file_when_serialisation_fails(
    tiny, hf_sd, tmp_workdir, monkeypatch
):
    """A half-written checkpoint is worse than none -- it looks resumable and
    is not. The write goes to a sibling temp path and is renamed into place
    only once complete, so a mid-write failure leaves nothing that resembles a
    checkpoint."""
    import safetensors.torch as st

    def boom(*a, **k):
        raise OSError("simulated ENOSPC mid-write")

    monkeypatch.setattr(st, "save_file", boom)
    out = os.path.join(tmp_workdir, "partial")
    with pytest.raises(OSError, match="ENOSPC"):
        ckpt.save_nmo_as_hf(ckpt.hf_to_nmo(hf_sd), tiny, out)
    leftovers = os.listdir(out) if os.path.isdir(out) else []
    assert not any(f.endswith((".safetensors", ".partial", "config.json")) for f in leftovers), (
        leftovers
    )


def test_config_write_is_atomic_and_leaves_no_partial(tiny, tmp_workdir, monkeypatch):
    """Same contract for config.json: a failed write leaves neither a
    config.json nor a `.partial` next to good weights."""
    out = os.path.join(tmp_workdir, "cfgfail")

    def boom(*a, **k):
        raise OSError("simulated ENOSPC writing config")

    monkeypatch.setattr(hfc, "config_to_json", boom)
    with pytest.raises(OSError, match="ENOSPC"):
        hfc.write_config(tiny, out)
    assert not any(f.startswith("config.json") for f in os.listdir(out))


# ---------------------------------------------------------------------------
# hf_config: the transformers-free config translation
# ---------------------------------------------------------------------------


def test_hf_config_roundtrips_through_nmo_config(tiny, tmp_workdir):
    """to_hf_config / from_hf_config are inverses on every architecture field,
    through the dict AND through a config.json on disk. The 7B default must
    survive it too: a config that drifts in translation reloads a checkpoint
    into the wrong-shaped model."""
    for cfg in (tiny, ModelConfig()):
        d = ckpt.to_hf_config(cfg)
        for back in (
            ckpt.from_hf_config(d),
            ckpt.from_hf_config(json.loads(json.dumps(d))),
        ):
            for f in ARCH_FIELDS:
                assert getattr(back, f) == getattr(cfg, f), f

        hfc.write_config(cfg, tmp_workdir, dtype=torch.bfloat16)
        on_disk = hfc.read_config(tmp_workdir)
        for f in ARCH_FIELDS:
            assert getattr(on_disk, f) == getattr(cfg, f), f
        # also reachable by the config.json path itself, and via to_dict() duck-typing
        assert hfc.read_config(os.path.join(tmp_workdir, "config.json")).n_layers == cfg.n_layers

        class Duck:  # what an oracle's transformers config looks like to us
            def to_dict(self):
                return d

        assert ckpt.from_hf_config(Duck()).hidden_size == cfg.hidden_size


def test_n_params_is_derived_not_inherited(tiny):
    """The pre-port translation left n_params at the 7B default for every
    config it touched. It is derived from the architecture now: the v0.3
    default reproduces ModelConfig's own constant, and the tiny model reports
    its real size."""
    assert hfc.count_params(ModelConfig()) == ModelConfig().n_params == 7_248_023_552
    back = ckpt.from_hf_config(ckpt.to_hf_config(tiny))
    assert back.n_params == hfc.count_params(tiny)
    assert back.n_params < 10_000_000
    n_actual = sum(p.numel() for p in NmoMistralForCausalLM(tiny).parameters())
    assert back.n_params == n_actual, "count_params disagrees with the real model"


def test_config_defaults_match_the_real_v03_facts():
    """ModelConfig's defaults claim to BE Mistral-7B-v0.3. Check that claim
    field by field against (1) the in-repo V03_FACTS table, (2) the recorded
    upstream config.json read through the production reader, (3) the real
    local snapshot when present, and (4) a config.json that transformers
    itself wrote from the real weights when present. A comment is not
    evidence: this fails if someone 'fixes' rms_norm_eps to the 1e-6 library
    default or lets sliding_window back in."""
    defaults = ModelConfig()
    for f in ARCH_FIELDS:
        assert getattr(defaults, f) == hfc.V03_FACTS[f], f

    # (2) upstream v0.3 omits head_dim and attention_dropout: both must come out
    # as the facts say (128 inferred, 0.0 defaulted), not as anything inherited.
    assert "head_dim" not in V03_UPSTREAM_CONFIG
    from_upstream = hfc.from_hf_config_dict(V03_UPSTREAM_CONFIG)
    for f in ARCH_FIELDS:
        assert getattr(from_upstream, f) == getattr(defaults, f), f
    assert from_upstream.n_params == defaults.n_params

    # (3) the recorded copy must still equal the real snapshot's file
    snaps = sorted(glob.glob(V03_SNAPSHOT_GLOB))
    if snaps:
        real = hfc.read_config_dict(snaps[-1])
        for k in (
            "num_hidden_layers",
            "hidden_size",
            "intermediate_size",
            "num_attention_heads",
            "num_key_value_heads",
            "vocab_size",
            "max_position_embeddings",
            "rms_norm_eps",
            "sliding_window",
            "tie_word_embeddings",
        ):
            assert real[k] == V03_UPSTREAM_CONFIG[k], k
        snap_cfg = hfc.read_config(snaps[-1])
        for f in ARCH_FIELDS:
            assert getattr(snap_cfg, f) == getattr(defaults, f), f

    # (4) independent real artifact: transformers-5-written config of the real model
    emitted = sorted(glob.glob(LEGACY_EMITTED_GLOB))
    if emitted:
        e_cfg = hfc.read_config(emitted[-1])
        for f in ARCH_FIELDS:
            assert getattr(e_cfg, f) == getattr(defaults, f), f

    if not snaps and not emitted:
        pytest.skip("recorded copy verified; no local snapshot or legacy checkpoint to cross-check")


def test_saved_config_carries_verified_v03_facts(tiny, hf_sd, tmp_workdir):
    """The emitted config.json must state sliding_window=null and
    attention_dropout=0.0 explicitly. v0.3 REMOVED sliding-window attention;
    a reader that fills defaults resolves an omitted field to 4096 and the
    checkpoint reloads as a sliding-window model, silently changing the
    attention pattern of every checkpoint we ship. Same for the rope theta
    spellings and dtype keys that older/newer outside readers look for."""
    out = os.path.join(tmp_workdir, "ckpt")
    ckpt.save_nmo_as_hf(ckpt.hf_to_nmo(hf_sd), tiny, out)
    raw = (Path(out) / "config.json").read_text()
    emitted = json.loads(raw)

    assert "sliding_window" in emitted and emitted["sliding_window"] is None
    assert emitted["attention_dropout"] == 0.0
    assert emitted["architectures"] == ["MistralForCausalLM"]
    assert emitted["model_type"] == "mistral"
    assert emitted["tie_word_embeddings"] is False
    assert emitted["rms_norm_eps"] == tiny.rms_norm_eps
    assert emitted["head_dim"] == tiny.head_dim
    assert emitted["hidden_act"] == "silu"
    assert emitted["bos_token_id"] == 1 and emitted["eos_token_id"] == 2
    assert isinstance(emitted["transformers_version"], str)
    # both RoPE spellings, in agreement
    assert emitted["rope_theta"] == tiny.rope_theta
    assert emitted["rope_parameters"] == {"rope_theta": tiny.rope_theta, "rope_type": "default"}
    # both dtype spellings, equal to what the tensors actually are
    assert emitted["dtype"] == emitted["torch_dtype"] == "float32"
    # transformers' formatting: sorted keys, 2-space indent, trailing newline
    assert raw == json.dumps(emitted, indent=2, sort_keys=True) + "\n"


def test_saved_dtype_follows_the_cast(tiny, hf_sd, tmp_workdir):
    """save_nmo_as_hf(dtype=bf16) must say bfloat16 in config.json AND write
    bf16 tensors: a config that claims fp32 over bf16 weights makes an outside
    loader upcast silently (or refuse), and the reverse hides a precision loss."""
    out = os.path.join(tmp_workdir, "bf16")
    ckpt.save_nmo_as_hf(ckpt.hf_to_nmo(hf_sd), tiny, out, dtype=torch.bfloat16)
    assert json.loads((Path(out) / "config.json").read_text())["dtype"] == "bfloat16"
    assert all(v.dtype == torch.bfloat16 for v in ckpt.load_hf_safetensors(out).values())


def _mut(**changes):
    """V03_UPSTREAM_CONFIG with keys changed; a value of ... deletes the key."""
    d = json.loads(json.dumps(V03_UPSTREAM_CONFIG))
    for k, v in changes.items():
        if v is ...:
            d.pop(k)
        else:
            d[k] = v
    return d


@pytest.mark.parametrize(
    "bad, match",
    [
        (_mut(sliding_window=...), "sliding_window"),
        (_mut(rms_norm_eps=...), "rms_norm_eps"),
        (_mut(rope_theta=...), "rope_theta"),
        (_mut(tie_word_embeddings=...), "tie_word_embeddings"),
        (_mut(num_key_value_heads=...), "num_key_value_heads"),
        (_mut(vocab_size=...), "vocab_size"),
        (_mut(model_type="llama"), "model_type"),
        (_mut(architectures=["LlamaForCausalLM"]), "architectures"),
        (_mut(hidden_act="gelu"), "hidden_act"),
        (_mut(num_key_value_heads=7), "not divisible"),
        (_mut(rope_parameters={"rope_type": "yarn", "rope_theta": 1e6}), "rope_type"),
        (_mut(rope_parameters={"rope_type": "default", "rope_theta": 1e4}), "conflicting"),
        (_mut(rms_norm_eps="1e-5"), "rms_norm_eps"),
        (_mut(vocab_size=True), "vocab_size"),
        (_mut(tie_word_embeddings=0), "tie_word_embeddings"),
    ],
)
def test_reader_refuses_underspecified_or_foreign_configs(bad, match):
    """The reader's whole job is refusing to guess. Each of these is a config
    a default-filling reader would accept and turn into a different model
    (sliding_window=4096, eps=1e-6, theta=1e4) or a different architecture."""
    with pytest.raises(hfc.HFConfigError, match=match):
        hfc.from_hf_config_dict(bad)
    # the silent direction: the unmodified recorded config parses
    hfc.from_hf_config_dict(V03_UPSTREAM_CONFIG)


def test_reader_accepts_either_rope_spelling_alone():
    """transformers <=4.x wrote top-level rope_theta; >=5 writes
    rope_parameters. Real checkpoints exist in both; both must read, to the
    same theta."""
    only_top = _mut()
    only_new = _mut(rope_theta=..., rope_parameters={"rope_type": "default", "rope_theta": 1e6})
    assert "rope_parameters" not in only_top
    assert "rope_theta" not in only_new
    assert hfc.from_hf_config_dict(only_top).rope_theta == 1e6
    assert hfc.from_hf_config_dict(only_new).rope_theta == 1e6


def test_reader_missing_file_is_a_clear_error(tmp_workdir):
    with pytest.raises(FileNotFoundError, match="config.json"):
        hfc.read_config(tmp_workdir)


def test_the_model_still_refuses_what_the_reader_passes_through():
    """The reader maps a nonzero attention_dropout / a sliding window through
    unchanged on purpose; the MODEL owns the refusal (with the better
    message). Pin the hand-off so neither side assumes the other checks."""
    swa = hfc.from_hf_config_dict(_mut(sliding_window=4096))
    assert swa.sliding_window == 4096
    with pytest.raises(ValueError, match="sliding_window"):
        NmoMistralForCausalLM(swa, init_weights=False)
    drop = hfc.from_hf_config_dict(_mut(attention_dropout=0.1))
    with pytest.raises(ValueError, match="attention_dropout"):
        NmoMistralForCausalLM(drop, init_weights=False)


# ---------------------------------------------------------------------------
# The firewall, from this subsystem's side
# ---------------------------------------------------------------------------


def test_align_model_never_names_transformers_as_an_import():
    """AST scan of align/model (not a grep: docstrings legitimately say
    'transformers'). Covers lazy in-function imports, which the dependency
    firewall also counts."""
    banned = {"transformers", "flash_attn", "safetensors_hf", "huggingface_hub"}
    offenders = []
    for py in sorted((REPO / "align" / "model").rglob("*.py")):
        for node in ast.walk(ast.parse(py.read_text())):
            names = []
            if isinstance(node, ast.Import):
                names = [a.name.split(".")[0] for a in node.names]
            elif isinstance(node, ast.ImportFrom) and not node.level and node.module:
                names = [node.module.split(".")[0]]
            offenders += [f"{py.name}:{node.lineno}: {n}" for n in names if n in banned]
    assert not offenders, offenders


def test_importing_the_checkpoint_stack_does_not_load_transformers():
    """Behavioural twin of the AST scan: a fresh interpreter imports every
    module in align.model and runs a save+reload, and transformers must not
    be in sys.modules afterwards (it would be, via the old roundtrip_gate)."""
    code = (
        "import sys, os, torch\n"
        "from align.config import ModelConfig\n"
        "from align.model import checkpoint as c, hf_config, attention, mistral, gates\n"
        "cfg = ModelConfig(name='t', n_layers=2, hidden_size=32, intermediate_size=64,\n"
        "    n_heads=4, n_kv_heads=2, head_dim=8, vocab_size=64, max_position_embeddings=32)\n"
        "m = mistral.NmoMistralForCausalLM(cfg).eval()\n"
        "out = os.path.join(os.environ['ALIGN_TEST_OUT'], 'fresh')\n"
        "rep = c.roundtrip_gate(m.state_dict(), cfg, out)\n"
        "assert rep['passed'], rep\n"
        "bad = [n for n in sys.modules if n.split('.')[0] in ('transformers', 'flash_attn')]\n"
        "assert not bad, bad\n"
    )
    out = REPO / ".tmp" / "tests" / "fresh_interp"
    out.mkdir(parents=True, exist_ok=True)
    env = dict(os.environ, ALIGN_TEST_OUT=str(out), PYTHONPATH=str(REPO), TMPDIR=str(REPO / ".tmp"))
    proc = subprocess.run(
        [sys.executable, "-c", code], capture_output=True, text=True, timeout=300, env=env
    )
    assert proc.returncode == 0, proc.stderr[-2000:]


# ---------------------------------------------------------------------------
# Oracle tier: the on-disk layout vs real transformers (legacy venv only)
# ---------------------------------------------------------------------------


@pytest.mark.oracle
def test_oracle_saved_checkpoint_loads_in_transformers_and_matches(tiny, trained, tmp_workdir):
    """'Stays HF-loadable by outside tools', as a checkable claim: a checkpoint
    saved by save_nmo_as_hf is loaded by real transformers `from_pretrained`
    with every tensor bitwise-preserved, and its logits equal align's own."""
    transformers = pytest.importorskip("transformers")
    out = os.path.join(tmp_workdir, "oracle_ckpt")
    ckpt.save_nmo_as_hf(trained.state_dict(), tiny, out)

    back = transformers.MistralForCausalLM.from_pretrained(out, dtype=torch.float32).eval()
    got = back.state_dict()
    want = ckpt.nmo_to_hf(trained.state_dict())
    assert set(got) == set(want), (sorted(set(want) - set(got)), sorted(set(got) - set(want)))
    for k in want:
        assert got[k].dtype == want[k].dtype and torch.equal(got[k], want[k]), k

    # transformers must have READ our architecture fields, not defaulted them
    c = back.config
    assert c.sliding_window is None
    assert c.rms_norm_eps == tiny.rms_norm_eps
    assert c.num_key_value_heads == tiny.n_kv_heads and c.head_dim == tiny.head_dim
    assert ckpt.from_hf_config(c).rope_theta == tiny.rope_theta

    ids = torch.randint(0, tiny.vocab_size, (2, 13))
    with torch.no_grad():
        mine = trained(input_ids=ids)
        theirs = back(input_ids=ids).logits
    assert (mine - theirs).abs().max().item() <= 2e-5


@pytest.mark.oracle
def test_oracle_transformers_written_checkpoint_loads_through_align(tiny, trained, tmp_workdir):
    """The other direction: a checkpoint transformers itself saves (its own
    config.json, its own key set) is read by hf_config.read_config +
    load_hf_safetensors + hf_to_nmo into an align model with identical logits."""
    transformers = pytest.importorskip("transformers")
    hf_cfg = transformers.MistralConfig.from_dict(ckpt.to_hf_config(tiny))
    hf_cfg._attn_implementation = "eager"
    hf = transformers.MistralForCausalLM(hf_cfg).to(torch.float32).eval()
    hf.load_state_dict(ckpt.nmo_to_hf(trained.state_dict()), strict=True)
    out = os.path.join(tmp_workdir, "hf_written")
    hf.save_pretrained(out)

    back = _reload(out)
    ids = torch.randint(0, tiny.vocab_size, (2, 11))
    with torch.no_grad():
        theirs = hf(input_ids=ids).logits
        mine = back(input_ids=ids)
    assert (mine - theirs).abs().max().item() <= 2e-5
    for f in ARCH_FIELDS:
        assert getattr(hfc.read_config(out), f) == getattr(tiny, f), f


def test_load_hf_safetensors_reads_sharded_base_snapshots_and_checks_the_index(tmp_workdir):
    """THE BASE-SNAPSHOT LAYOUT GAP: the pinned Mistral-7B-v0.3 snapshot is THREE
    shards plus model.safetensors.index.json, but `load_hf_safetensors` only knew
    the single-file layout trainer checkpoints use (transformers resolved shards
    for us before the port). Both directions: shards merge (with an optional
    per-shard dtype cast that leaves integer tensors alone), the single-file
    layout and a bare file path still work, and an index naming an absent shard
    or an uncarried tensor raises instead of returning a partial model."""
    import json

    from safetensors.torch import save_file

    from align.model import checkpoint as ckpt

    a = {"x": torch.randn(3), "i": torch.arange(3)}
    b = {"z": torch.randn(4)}
    save_file(a, os.path.join(tmp_workdir, "s1.safetensors"))
    save_file(b, os.path.join(tmp_workdir, "s2.safetensors"))
    wm = {**{k: "s1.safetensors" for k in a}, **{k: "s2.safetensors" for k in b}}
    idx = os.path.join(tmp_workdir, "model.safetensors.index.json")
    with open(idx, "w") as f:
        json.dump({"weight_map": wm}, f)

    out = ckpt.load_hf_safetensors(tmp_workdir)
    assert set(out) == {"x", "i", "z"} and out["x"].dtype == torch.float32
    cast = ckpt.load_hf_safetensors(tmp_workdir, dtype=torch.bfloat16)
    assert cast["x"].dtype == cast["z"].dtype == torch.bfloat16 and cast["i"].dtype == torch.int64

    with open(idx, "w") as f:  # index claims a tensor no shard carries
        json.dump({"weight_map": {**wm, "ghost": "s1.safetensors"}}, f)
    with pytest.raises(RuntimeError, match="ghost"):
        ckpt.load_hf_safetensors(tmp_workdir)

    os.remove(os.path.join(tmp_workdir, "s2.safetensors"))
    with pytest.raises(FileNotFoundError, match="s2"):
        ckpt.load_hf_safetensors(tmp_workdir)

    single = os.path.join(tmp_workdir, "single")
    os.makedirs(single)
    save_file(a, os.path.join(single, "model.safetensors"))
    assert set(ckpt.load_hf_safetensors(single)) == set(a)
    assert set(ckpt.load_hf_safetensors(os.path.join(single, "model.safetensors"))) == set(a)
    with pytest.raises(FileNotFoundError):
        ckpt.load_hf_safetensors(os.path.join(tmp_workdir, "nowhere"))


def _write_tok(d, names):
    os.makedirs(d, exist_ok=True)
    for n in names:
        Path(d, n).write_text(f"snap-{n}")


def test_partial_tokenizer_set_falls_back_to_the_whole_vendored_set(tmp_path, caplog):
    """A snapshot with only some tokenizer files used to yield a partial copy
    with no warning. The set is all-or-nothing: the vendored set is copied
    whole, and a warning names the missing files. Mutant (copy what exists)
    leaves snapshot-sourced content / a short set and fails."""
    snap = tmp_path / "snap"
    _write_tok(snap, ["tokenizer.json", "tokenizer_config.json"])
    (snap / "config.json").write_text("{}")
    dst = tmp_path / "dst"
    dst.mkdir()
    with caplog.at_level("WARNING"):
        copied = ckpt.copy_tokenizer_files(str(snap), str(dst))
    assert sorted(copied) == sorted(ckpt.TOKENIZER_FILES)
    assert sorted(os.listdir(dst)) == sorted(ckpt.TOKENIZER_FILES)
    assert not (dst / "tokenizer.json").read_text().startswith("snap-")
    assert any("PARTIAL" in r.message for r in caplog.records)


def test_complete_tokenizer_set_is_copied_from_the_snapshot(tmp_path):
    snap = tmp_path / "snap"
    _write_tok(snap, ckpt.TOKENIZER_FILES)
    (snap / "config.json").write_text("{}")
    dst = tmp_path / "dst"
    dst.mkdir()
    assert sorted(ckpt.copy_tokenizer_files(str(snap), str(dst))) == sorted(ckpt.TOKENIZER_FILES)
    assert (dst / "tokenizer.json").read_text() == "snap-tokenizer.json"