import shutil import tempfile import unittest from types import SimpleNamespace import torch import transformers from transformers import LlamaConfig, LlamaForCausalLM from models.decoder import Decoder # Small, fixed dimensions for a tiny randomly-initialized decoder. Kept minimal so # the whole suite builds one model once and runs offline in well under a second. VOCAB = 100 HIDDEN = 32 LAYERS = 2 HEADS = 4 BATCH = 2 SEQ = 5 # Module-level fixtures populated by setUpModule and shared by all test classes. MODEL_DIR = None # path to a saved tiny LlamaForCausalLM DECODER = None # a Decoder wrapping that model (read-only tests reuse this) def make_config(model_dir, lm_hidden_dim=HIDDEN, lm_use_tokens=False): """ Build a minimal config exposing exactly the three attributes Decoder reads. Parameters: * model_dir (str) : filesystem path passed to AutoModelForCausalLM.from_pretrained; requires it to contain a saved HF causal LM * lm_hidden_dim (int) : the hidden size Decoder asserts against the loaded model; requires lm_hidden_dim >= 1. Defaults to the fixture's true hidden size so the assertion passes; override it to provoke the mismatch guard. * lm_use_tokens (bool) : value stored on the Decoder; defaults to False (backbone mode) Returns: A SimpleNamespace with lm_model_type, lm_hidden_dim, and lm_use_tokens. """ return SimpleNamespace( lm_model_type=model_dir, lm_hidden_dim=lm_hidden_dim, lm_use_tokens=lm_use_tokens, ) def setUpModule(): # Build and save one tiny tied-weight Llama, then wrap it once. Done at module scope # so the (relatively) expensive from_pretrained happens a single time. global MODEL_DIR, DECODER transformers.logging.set_verbosity_error() torch.manual_seed(0) llama_cfg = LlamaConfig( vocab_size=VOCAB, hidden_size=HIDDEN, intermediate_size=2 * HIDDEN, num_hidden_layers=LAYERS, num_attention_heads=HEADS, num_key_value_heads=HEADS, max_position_embeddings=128, tie_word_embeddings=True, ) model = LlamaForCausalLM(llama_cfg).eval() MODEL_DIR = tempfile.mkdtemp(prefix="decoder_test_") model.save_pretrained(MODEL_DIR) DECODER = Decoder(make_config(MODEL_DIR), load_backbone=True).eval() def tearDownModule(): if MODEL_DIR is not None: shutil.rmtree(MODEL_DIR, ignore_errors=True) def make_ids(batch=BATCH, seq=SEQ, seed=0): """ Draw a (batch, seq) tensor of valid token ids. Parameters: * batch (int) : number of sequences; requires batch >= 1 * seq (int) : sequence length; requires seq >= 1 * seed (int) : seed for the local generator; requires seed >= 0 Returns: A (batch, seq) int64 tensor with every entry in [0, VOCAB). """ generator = torch.Generator().manual_seed(seed) return torch.randint(0, VOCAB, (batch, seq), generator=generator) # --- Construction ------------------------------------------------------------ class TestDecoderConstruction(unittest.TestCase): """Behavioral spec for constructing a Decoder around a transformers model.""" def test_construction_returns_nn_module(self): # A Decoder is itself an nn.Module so it can be a child of VisionLanguageModel. self.assertIsInstance(DECODER, torch.nn.Module) def test_hidden_size_reflects_model_config(self): # hidden_size is read straight from the loaded model and is what the VLM uses to # size its modality projector. self.assertEqual(DECODER.hidden_size, HIDDEN) def test_asserts_on_hidden_size_mismatch(self): # The guard exists so swapping decoders without updating cfg.lm_hidden_size fails # loudly at construction rather than silently misshaping image embeddings. with self.assertRaises(AssertionError): Decoder(make_config(MODEL_DIR, lm_hidden_dim=HIDDEN + 1), load_backbone=True) def test_lm_use_tokens_is_stored(self): # The flag is carried verbatim from cfg; the VLM reads it to decide whether the # head still needs applying. for flag in (True, False): with self.subTest(flag=flag): decoder = Decoder(make_config(MODEL_DIR, lm_use_tokens=flag), load_backbone=True) self.assertEqual(decoder.lm_use_tokens, flag) # --- Accessors (token_embedding / head / base) ------------------------------- class TestDecoderAccessors(unittest.TestCase): """Behavioral spec for the property accessors that expose the model's pieces.""" def test_token_embedding_is_the_input_embedding(self): # token_embedding must be the model's own input embedding, not a fresh layer. self.assertIsInstance(DECODER.token_embedding, torch.nn.Embedding) self.assertIs(DECODER.token_embedding, DECODER.model.get_input_embeddings()) def test_token_embedding_applies_like_the_model(self): # Calling the accessor must embed ids identically to the underlying model. ids = make_ids() self.assertTrue(torch.equal( DECODER.token_embedding(ids), DECODER.model.get_input_embeddings()(ids))) def test_head_projects_hidden_to_vocab(self): # head is the output projection: hidden_size -> vocab_size. self.assertIsInstance(DECODER.head, torch.nn.Linear) self.assertEqual(DECODER.head.in_features, HIDDEN) self.assertEqual(DECODER.head.out_features, VOCAB) def test_head_produces_logits_shape(self): # Applying head to a hidden-state tensor yields per-token vocab logits. hidden = torch.randn(BATCH, SEQ, HIDDEN) logits = DECODER.head(hidden) self.assertEqual(logits.shape, (BATCH, SEQ, VOCAB)) def test_base_is_the_headless_model(self): # base must be the decoder stack obtained via get_decoder (no LM head attached). self.assertIs(DECODER.base, DECODER.model.get_decoder()) def test_base_returns_hidden_states_not_logits(self): # The base model emits hidden states (width HIDDEN), distinct from logits (VOCAB); # this is why the VLM can apply the head separately. embeds = DECODER.token_embedding(make_ids()) out = DECODER.base(inputs_embeds=embeds) self.assertEqual(out.last_hidden_state.shape[-1], HIDDEN) self.assertNotEqual(HIDDEN, VOCAB) # guards the test's own premise # --- Forward ----------------------------------------------------------------- class TestDecoderForward(unittest.TestCase): """Behavioral spec for the training-path forward.""" def test_forward_returns_two_tuple(self): # The VLM unpacks `logits, _ = self.decoder(...)`, so forward must return a pair. embeds = DECODER.token_embedding(make_ids()) out = DECODER(embeds) self.assertIsInstance(out, tuple) self.assertEqual(len(out), 2) def test_forward_second_element_is_none(self): # No KV cache during training: the second slot is always None. embeds = DECODER.token_embedding(make_ids()) _, cache = DECODER(embeds) self.assertIsNone(cache) def test_forward_output_shape(self): # Hidden states preserve (batch, seq) and carry hidden_size features. embeds = DECODER.token_embedding(make_ids()) hidden, _ = DECODER(embeds) self.assertEqual(hidden.shape, (BATCH, SEQ, HIDDEN)) def test_forward_consumes_embeddings_matching_direct_base_call(self): # forward is a thin shim over base(inputs_embeds=...); its hidden states must # equal calling the base model directly with the same embeddings. embeds = DECODER.token_embedding(make_ids()) with torch.no_grad(): hidden, _ = DECODER(embeds) direct = DECODER.base(inputs_embeds=embeds).last_hidden_state self.assertTrue(torch.allclose(hidden, direct, atol=1e-6)) def test_forward_output_dtype_tracks_model(self): # The shim adds no casting, so hidden states come back in the model's dtype. embeds = DECODER.token_embedding(make_ids()) hidden, _ = DECODER(embeds) self.assertEqual(hidden.dtype, DECODER.model.dtype) def test_forward_respects_attention_mask(self): # The mask must actually reach the model: masking out trailing positions changes # the hidden states of the positions that can attend to them. embeds = DECODER.token_embedding(make_ids()) full_mask = torch.ones(BATCH, SEQ, dtype=torch.long) partial_mask = full_mask.clone() partial_mask[:, -2:] = 0 # forbid attending to the last two tokens with torch.no_grad(): hidden_full, _ = DECODER(embeds, attention_mask=full_mask) hidden_partial, _ = DECODER(embeds, attention_mask=partial_mask) self.assertEqual(hidden_full.shape, hidden_partial.shape) self.assertFalse(torch.allclose(hidden_full, hidden_partial, atol=1e-5)) # --- Equivalence and parameter registration ---------------------------------- class TestDecoderEquivalenceAndRegistration(unittest.TestCase): """ Behavioral spec tying the split (token_embedding -> base -> head) back to the full CausalLM, and verifying the property accessors do not duplicate weights. """ def test_embed_base_head_reconstructs_full_logits(self): # The defining invariant: embedding the ids, running the base, then applying the # head must reproduce the full AutoModelForCausalLM logits byte-for-byte (up to # float tolerance). This is exactly how the VLM rebuilds logits from hidden states. ids = make_ids() with torch.no_grad(): full_logits = DECODER.model(input_ids=ids).logits hidden, _ = DECODER(DECODER.token_embedding(ids)) reconstructed = DECODER.head(hidden) self.assertEqual(reconstructed.shape, full_logits.shape) self.assertTrue(torch.allclose(reconstructed, full_logits, atol=1e-4)) def test_only_the_model_is_a_registered_submodule(self): # The accessors are @property, not stored modules, so the Decoder's only direct # child is `model` — this is what prevents duplicate state_dict entries. self.assertEqual(list(DECODER._modules.keys()), ["model"]) def test_no_duplicate_parameters(self): # Because base/token_embedding/head are not separately registered, the Decoder # owns exactly the underlying model's parameters and nothing more. decoder_params = sum(p.numel() for p in DECODER.parameters()) model_params = sum(p.numel() for p in DECODER.model.parameters()) self.assertEqual(decoder_params, model_params) # Every parameter name is reached through the single `model.` child. self.assertTrue(all(name.startswith("model.") for name, _ in DECODER.named_parameters())) def test_weight_tying_is_preserved_through_accessors(self): # The fixture ties input/output embeddings; the accessors must expose the same # shared parameter (rather than copies), confirming they return the real layers. self.assertIs(DECODER.head.weight, DECODER.token_embedding.weight) if __name__ == "__main__": unittest.main()