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