Spaces:
Sleeping
Sleeping
| """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") | |
| 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() | |
| 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 | |
| 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) | |
| 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) | |