| """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") |
| |
| 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) |
| |
| 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)) |
| |
| 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]) |
| |
| |
| 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() |
|
|
