"""Exercise every engine against real, local BF16 Transformers models. These random, tiny models check execution and restoration, not scientific model quality. Their fast tokenizer is built in memory; no checkpoint or network is needed. Architectures cover the four decoder layouts currently used by engines. """ import json import pytest import torch from tokenizers import Tokenizer, models, pre_tokenizers from transformers import ( GPT2Config, GPT2LMHeadModel, GPTNeoXConfig, GPTNeoXForCausalLM, Gemma4ForCausalLM, Gemma4TextConfig, LlamaConfig, LlamaForCausalLM, OPTConfig, OPTForCausalLM, PreTrainedTokenizerFast, Qwen3Config, Qwen3ForCausalLM, ) from bol_space.inspection import run_inspection from bol_space.layers import find_decoder_layers from bol_space.printout import format_printout from bol_space.surgery import run_surgery INSPECTIONS = ( "fingerprint", "silhouette", "semantic", "cka", "attention", "gradient", "butterfly", "tugofwar", ) SURGERIES = ("layer_sweep", "component", "trisect", "wormhole", "noise", "duets") ARCHITECTURES = ("llama", "qwen3", "gpt2", "gpt_neox", "opt", "gemma4_text") @pytest.fixture(scope="module") def tokenizer(): words = ( "cat dog walks home river ocean wolf lion code python car truck vehicle " "happy joyful cold frozen king queen run sprint democracy algebra justice " "concrete philosophy equation banana silence orange MAPLE RAVEN The the " "access is . A small bird watched from a branch" ).split() vocabulary = {word: index for index, word in enumerate(dict.fromkeys( ["[PAD]", "[UNK]", "[EOS]"] + words ))} backend = Tokenizer(models.WordLevel(vocabulary, unk_token="[UNK]")) backend.pre_tokenizer = pre_tokenizers.Whitespace() return PreTrainedTokenizerFast( tokenizer_object=backend, pad_token="[PAD]", unk_token="[UNK]", eos_token="[EOS]", model_max_length=256, ) def make_model(architecture, vocabulary_size): common = dict(vocab_size=vocabulary_size, bos_token_id=2, eos_token_id=2, pad_token_id=0, attn_implementation="eager") dense = dict(hidden_size=16, intermediate_size=24, num_hidden_layers=4, num_attention_heads=2, max_position_embeddings=256) if architecture == "llama": model = LlamaForCausalLM(LlamaConfig(**common, **dense, num_key_value_heads=1)) elif architecture == "qwen3": model = Qwen3ForCausalLM(Qwen3Config( **common, **dense, num_key_value_heads=1, head_dim=8, )) elif architecture == "gpt2": model = GPT2LMHeadModel(GPT2Config( **common, n_embd=16, n_inner=24, n_layer=4, n_head=2, n_positions=256, resid_pdrop=0., embd_pdrop=0., attn_pdrop=0., )) elif architecture == "gpt_neox": model = GPTNeoXForCausalLM(GPTNeoXConfig(**common, **dense, rotary_pct=0.5)) elif architecture == "opt": model = OPTForCausalLM(OPTConfig( **common, hidden_size=16, ffn_dim=24, num_hidden_layers=4, num_attention_heads=2, max_position_embeddings=256, dropout=0., attention_dropout=0., )) elif architecture == "gemma4_text": model = Gemma4ForCausalLM(Gemma4TextConfig( **common, **dense, num_key_value_heads=1, head_dim=8, global_head_dim=8, layer_types=["sliding_attention", "full_attention"] * 2, num_kv_shared_layers=2, use_double_wide_mlp=True, hidden_size_per_layer_input=4, vocab_size_per_layer_input=vocabulary_size, )) assert model.model.layers[2].self_attn.is_kv_shared_layer assert not hasattr(model.model.layers[2].self_attn, "k_proj") assert model.model.layers[2].mlp.gate_proj.out_features == 48 else: raise AssertionError(architecture) return model.to(dtype=torch.bfloat16).eval() @pytest.fixture(scope="module", params=ARCHITECTURES) def model(request, tokenizer): with torch.random.fork_rng(): torch.manual_seed(23) instance = make_model(request.param, len(tokenizer)) # Existing caller state must survive gradient and surgery experiments too. parameters = list(instance.parameters()) parameters[0].requires_grad_(False) parameters[-1].grad = torch.full_like(parameters[-1], 0.125) # Transformers 5 installs its permanent, context-controlled capture hooks on # first request. Warm them once so restoration checks detect subsequent hook # leaks without mistaking the library's expected initialization for a leak. with torch.no_grad(): instance(**tokenizer("cat walks home", return_tensors="pt"), use_cache=False, output_hidden_states=True, output_attentions=True) return instance def options(): return { "text": "cat walks home\nA small bird watched from a branch", "compare_text": "dog walks home\nA small bird watched from a branch", "words": "cat, dog\nriver, ocean", "choice_a": "MAPLE", "choice_b": "RAVEN", "prompt": "cat walks", "layer_start": 0, "layer_end": 3, "sample_layers": 4, "max_tokens": 64, "max_new_tokens": 2, "seed": 17, "noise_scale": 0.1, } def snapshot(model): return { "weights": {name: value.detach().clone() for name, value in model.state_dict().items()}, "modes": [(module, module.training) for module in model.modules()], "parameters": [(parameter, parameter.requires_grad, parameter.grad, None if parameter.grad is None else parameter.grad.clone()) for parameter in model.parameters()], "hooks": [(module, tuple(module._forward_hooks), tuple(module._forward_pre_hooks), tuple(module._backward_hooks)) for module in model.modules()], } def assert_restored(model, previous): assert set(model.state_dict()) == set(previous["weights"]) for name, value in model.state_dict().items(): assert torch.equal(value, previous["weights"][name]), name for module, mode in previous["modes"]: assert module.training is mode for parameter, flag, gradient, copy in previous["parameters"]: assert parameter.requires_grad is flag assert parameter.grad is gradient if copy is not None: assert torch.equal(parameter.grad, copy) for module, forward, before_forward, backward in previous["hooks"]: assert tuple(module._forward_hooks) == forward assert tuple(module._forward_pre_hooks) == before_forward assert tuple(module._backward_hooks) == backward @pytest.mark.parametrize("tool", INSPECTIONS + SURGERIES) def test_every_engine_on_real_bf16_transformers(model, tokenizer, tool): before = snapshot(model) probe = tokenizer("cat walks home", return_tensors="pt") with torch.no_grad(): original_logits = model(**probe, use_cache=False).logits.clone() engine = run_inspection if tool in INSPECTIONS else run_surgery # Engines also support callers using no_grad; gradient enables its own graph. with torch.no_grad(): result = engine(model, tokenizer, tool, options()) assert result["rows"] json.dumps(result, allow_nan=False) result["run"] = {"tool": tool} result["model"] = {"repo_id": type(model).__name__, "dtype": "bfloat16"} assert result["title"].upper() in format_printout(result) if tool in SURGERIES: assert result["details"]["integrity_verified"] is True assert result["details"]["generation_comparisons"] == 2 assert len(result["text_outputs"]) == 2 if tool == "gradient": assert all(row["gradient_norm"] > 0 for row in result["rows"]) if tool == "fingerprint": assert {row["layer"] for row in result["rows"]} == {0, 1, 2, 3} assert_restored(model, before) with torch.no_grad(): restored_logits = model(**probe, use_cache=False).logits assert torch.equal(original_logits, restored_logits) @pytest.mark.parametrize("tool", ("layer_sweep", "noise")) def test_real_forward_failure_during_edit_restores_every_tensor(model, tokenizer, tool): before = snapshot(model) reached_edited_forward = False def fail_after_edit(module, args, output): nonlocal reached_edited_forward if any(not torch.equal(value, before["weights"][name]) for name, value in model.state_dict().items()): reached_edited_forward = True raise RuntimeError("injected real-transformer edited forward failure") # Fail after a real decoder block executes with the temporary weight edit # still in place, including the model's own permanent capture hooks. _, layers = find_decoder_layers(model) handle = layers[-1].register_forward_hook(fail_after_edit) try: with pytest.raises(RuntimeError, match="edited forward failure"): run_surgery(model, tokenizer, tool, options()) finally: handle.remove() assert reached_edited_forward assert_restored(model, before) def test_real_gradient_supports_outer_inference_mode(model, tokenizer): before = snapshot(model) with torch.inference_mode(): result = run_inspection(model, tokenizer, "gradient", options()) assert all(row["gradient_norm"] > 0 for row in result["rows"]) assert_restored(model, before)