Overview

This PR introduces a modular and approach to model testing by separating forward pass and generation tests. It adds the handling of generative models (e.g., causal language models like LLaMA) while maintaining compatibility for non-generative models.

Changes

1. Test Structure Refactoring

• test_model_forward(): Tests only the forward pass signature (renamed from test_model_output).
• test_model_generation(): New dedicated test for generation testing (1 prefill + 1 decode).
• test_model_output(): Maintained for backward compatibility, delegates to test_model_forward().

2. Generation Testing

• _supports_generation: Class attribute to enable generation testing (set to True in test classes).
• _get_generation_signature():
  • Performs 2-token generation using generate() with greedy sampling.
  • Captures logits for prefill and decode steps.
  • Applies the same signature compression as get_signature() (max-min normalization).
• Generation signature: Concatenates prefill + decode logits, uses consistent signature pattern.

3. Independent Expectation Files

• Forward pass: models.test_llama.TestLlama2.test_model_forward.
• Generation: models.test_llama.TestLlama2.test_model_generation.
• Each test maintains separate expectations for cleaner separation of concerns.

Usage

Capture New Expectations

# Capture forward pass expectations
CUDA_VISIBLE_DEVICES=-1 python -m pytest tests/models/test_llama.py::TestLlama2::test_model_forward -v -s --capture_expectation

# Capture generation expectations  
CUDA_VISIBLE_DEVICES=-1 python -m pytest tests/models/test_llama.py::TestLlama2::test_model_generation -v -s --capture_expectation

Test

# Forward pass test
CUDA_VISIBLE_DEVICES=-1 python -m pytest tests/models/test_llama.py::TestLlama2::test_model_forward -v -s 

# Generation test
CUDA_VISIBLE_DEVICES=-1 python -m pytest tests/models/test_llama.py::TestLlama2::test_model_generation -v -s 

@spzala @JRosenkranz