"""Small CPU checks for hybrid LoRA, memory states, and generation caching.""" import copy import json from pathlib import Path import sys import tempfile import unittest sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) import torch from omegaconf import OmegaConf from transformers import PreTrainedTokenizerFast, Qwen3_5Config, Qwen3_5TextConfig from tokenizers import Tokenizer from tokenizers.models import WordLevel from tokenizers.pre_tokenizers import Whitespace from parametric_skills.checkpoint import CONTROL_TOKENS, load_components, release_metadata, restore_special_embeddings from parametric_skills.hypernetwork import Metanetwork from parametric_skills.lora import LoraLinear from parametric_skills.qwen35 import LoraQwen35ForConditionalGeneration from parametric_skills.runtime import ParametricSkills, inference_context def tiny_model(): text = Qwen3_5TextConfig( vocab_size=64, hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, head_dim=16, linear_key_head_dim=8, linear_value_head_dim=8, linear_num_key_heads=2, linear_num_value_heads=2, linear_conv_kernel_dim=4, layer_types=["linear_attention", "full_attention"], max_position_embeddings=256, pad_token_id=0, eos_token_id=2, rope_parameters={"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.5, "mrope_section": [1, 1, 0]}, ) cfg = Qwen3_5Config(text_config=text.to_dict()) cfg._attn_implementation = "sdpa" cfg.text_config.num_mem_token = 32 model = LoraQwen35ForConditionalGeneration(cfg) rank = 2 width = max(model.lora_layer_param_widths(rank)) memory = (width + text.hidden_size - 1) // text.hidden_size cfg.text_config.num_mem_token = memory model.model.language_model.num_mem_token = memory model.model.language_model.mem_tokens = torch.nn.Parameter(torch.randn(memory, text.hidden_size) * 0.01) runtime_cfg = OmegaConf.create({ "num_layers": 2, "hidden_size": 32, "num_mem_token": memory, "model": {"lora_r": rank}, "metanetwork": {"type": "transformer", "method": "rl", "transformer_cfg": { "mean_pool_size": 1, "layer_transformer_first": True, "num_layers": 2, "couple_num_layers": 1, "scale": 0.01, "encoder_cfg": {"d_model": 32, "nhead": 2, "dim_feedforward": 64, "dropout": 0, "batch_first": True}, "couple_encoder_cfg": {"d_model": 8, "nhead": 2, "dim_feedforward": 16, "dropout": 0, "batch_first": True}, }}, }) network = Metanetwork(model, runtime_cfg, model.lora_params_numel(rank)) return model.eval(), network.eval(), runtime_cfg class RuntimeTests(unittest.TestCase): @classmethod def setUpClass(cls): torch.set_num_threads(1) torch.manual_seed(42) def test_dynamic_linear_matches_matrix_update_and_beams(self): linear = LoraLinear(5, 7, bias=False) x = torch.randn(4, 3, 5) a, b = torch.randn(2, 5, 2), torch.randn(2, 2, 7) adapter = {"A": a, "B": b, "C": None} expected = torch.nn.functional.linear(x, linear.weight) + torch.bmm(torch.bmm(x, a.repeat_interleave(2, 0)), b.repeat_interleave(2, 0)) torch.testing.assert_close(linear(x, adapter), expected) self.assertEqual(linear.lora_params_numel(2), 24) self.assertEqual(linear.lora_params_numel(4), 48) def test_rejects_other_variant_before_loading_tensors(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "checkpoint_config.json").write_text(json.dumps(release_metadata("Qwen3.5-9B"))) with self.assertRaisesRegex(ValueError, "metadata differs"): load_components("Qwen3.5-4B", root) def test_rejects_corrupted_weights_before_deserialization(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "checkpoint_config.json").write_text(json.dumps(release_metadata("Qwen3.5-4B"))) (root / "metanetwork.pth").write_bytes(b"not the published tensor file") with self.assertRaisesRegex(ValueError, "SHA256 mismatch"): load_components("Qwen3.5-4B", root) def test_skill_encoding_and_cached_generation(self): model, network, cfg = tiny_model() metalora = model.init_lora_dict(4, scale=0.01, device="cpu") # Nonzero A and B make the memory encoding depend on Meta-LoRA. for layer in metalora.values(): for group in layer.values(): for value in group.values(): value["B"] = torch.randn_like(value["B"]) * 0.01 ids = torch.tensor([[4, 5, 6]]) mask = torch.ones_like(ids) with inference_context(torch.device("cpu"), torch.float32): encoded = model(input_ids=ids, attention_mask=mask, loradict=metalora, use_cache=False, logits_to_keep=1) self.assertEqual(encoded.memory_states.shape, (1, 2, cfg.num_mem_token, 32)) adapter = network.generate_lora_dict(ids, mask, metalora) self.assertEqual(set(adapter), {0, 1}) base = model(input_ids=ids, attention_mask=mask, ignore_mem_token=True, use_cache=False).logits active = model(input_ids=ids, attention_mask=mask, loradict=adapter, ignore_mem_token=True, use_cache=False).logits self.assertGreater(float((base - active).abs().max()), 1e-7) # A new skill must change the generated adapter. changed = network.generate_lora_dict(ids + 1, mask, metalora) self.assertFalse(torch.equal(adapter[0]["mlp"]["up_proj"]["A"], changed[0]["mlp"]["up_proj"]["A"])) generated = model.generate(input_ids=ids, attention_mask=mask, loradict=adapter, ignore_mem_token=True, do_sample=False, max_new_tokens=3, eos_token_id=None, pad_token_id=0, use_cache=True) self.assertEqual(generated.shape, (1, 6)) # Compare cached continuation against full-prefix inference. prefix = model(input_ids=ids, attention_mask=mask, loradict=adapter, ignore_mem_token=True, use_cache=True) token = torch.tensor([[9]]) cached = model(input_ids=token, attention_mask=torch.ones(1, 4, dtype=torch.long), past_key_values=prefix.past_key_values, loradict=adapter, ignore_mem_token=True, use_cache=True).logits full = model(input_ids=torch.cat((ids, token), dim=1), attention_mask=torch.ones(1, 4, dtype=torch.long), loradict=adapter, ignore_mem_token=True, use_cache=False).logits torch.testing.assert_close(cached[:, -1], full[:, -1], rtol=2e-4, atol=2e-5) def test_encoder_matches_reference_padding_and_checks_padded_context(self): model, network, cfg = tiny_model() tokenizer_backend = Tokenizer(WordLevel( {"[PAD]": 0, "[UNK]": 1, "[BOS]": 2, "[EOS]": 3, "alpha": 4, "beta": 5, "gamma": 6}, unk_token="[UNK]", )) tokenizer_backend.pre_tokenizer = Whitespace() tokenizer = PreTrainedTokenizerFast( tokenizer_object=tokenizer_backend, pad_token="[PAD]", unk_token="[UNK]", padding_side="left", ) cfg.data = {"context_max_length": 3} runtime = ParametricSkills() runtime.config, runtime.tokenizer = cfg, tokenizer runtime.backbone, runtime.network = model, network runtime.device, runtime.dtype = torch.device("cpu"), torch.float32 runtime.metalora = model.init_lora_dict(4, scale=0.01, device="cpu") skill = "alpha beta gamma" canonical = tokenizer([skill], truncation=False, padding="longest", pad_to_multiple_of=8, return_tensors="pt") with inference_context(runtime.device, runtime.dtype): expected = network.generate_lora_dict( canonical.input_ids, canonical.attention_mask, runtime.metalora ) actual = runtime.encode_skill(skill) for layer in expected: for group in expected[layer]: for name in expected[layer][group]: for factor in ["A", "B"]: torch.testing.assert_close(actual[layer][group][name][factor], expected[layer][group][name][factor]) # The limit counts real skill tokens, while the model window must # accommodate all padding plus the appended learned memory tokens. with self.assertRaisesRegex(ValueError, "Skill has 3 tokens"): runtime.encode_skill(skill, max_skill_tokens=2) model.config.text_config.max_position_embeddings = cfg.num_mem_token + 7 with self.assertRaisesRegex(ValueError, "exceeds the base-model context"): runtime.encode_skill(skill) def test_special_embeddings_restore_by_token_text(self): model, _, _ = tiny_model() class Tokenizer: def get_vocab(self): return dict(zip(CONTROL_TOKENS, [61, 62, 63])) def convert_tokens_to_ids(self, token): return self.get_vocab()[token] rows = torch.randn(3, 32) payload = {"schema_version": 1, "tokens": list(CONTROL_TOKENS), "token_ids": [1, 2, 3], "input_embedding_rows": rows, "output_embedding_rows": rows + 1} restore_special_embeddings(model, Tokenizer(), payload) torch.testing.assert_close(model.get_input_embeddings().weight[61:64], rows) torch.testing.assert_close(model.get_output_embeddings().weight[61:64], rows + 1) bad = copy.copy(payload) bad["input_embedding_rows"] = rows[:, :4] with self.assertRaises(ValueError): restore_special_embeddings(model, Tokenizer(), bad) if __name__ == "__main__": unittest.main()