mirror of
https://github.com/pese-git/llm-arch-research.git
synced 2026-01-23 13:00:54 +00:00
25caf69ced
- Renamed Decoder (and decoder.py) to GptDecoder (gpt_decoder.py) for clarity in GPT1 - Implemented support for cache and use_cache parameters in GptDecoder.forward (API unification) - Adapted all usages in GPT model to use new decoder structure and handle tuple output - Refactored core tests (test_gpt.py, test_gpt_decoder.py, test_basic.py) to correctly expect tuple or logits and ensure shape/device checks work as before - Improved clarity and future extensibility for autoregressive generation and benchmarking - No changes to architectural details or training loop; pure API and test modernization
283 lines
9.6 KiB
Python
283 lines
9.6 KiB
Python
"""
|
|
Tests for GPT model.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
from llm.models.gpt import GPT
|
|
|
|
|
|
class TestGPT:
|
|
"""Test cases for GPT model."""
|
|
|
|
def test_initialization(self, gpt_config):
|
|
"""Test that GPT can be initialized."""
|
|
model = GPT(gpt_config)
|
|
assert model is not None
|
|
|
|
# Check that model has required components
|
|
assert hasattr(model, "_token_embeddings")
|
|
assert hasattr(model, "_position_embeddings")
|
|
assert hasattr(model, "_decoders")
|
|
assert hasattr(model, "_linear")
|
|
assert hasattr(model, "_dropout")
|
|
|
|
# Check number of decoder layers
|
|
assert len(model._decoders) == gpt_config["num_layers"]
|
|
|
|
def test_forward_pass(self, gpt_config, random_inputs):
|
|
"""Test forward pass of GPT."""
|
|
model = GPT(gpt_config)
|
|
|
|
# Forward pass
|
|
logits, _ = model(random_inputs)
|
|
|
|
# Check output shape
|
|
batch_size, seq_len = random_inputs.shape
|
|
vocab_size = gpt_config["vocab_size"]
|
|
assert logits.shape == (batch_size, seq_len, vocab_size)
|
|
assert isinstance(logits, torch.Tensor)
|
|
|
|
def test_forward_with_attention_mask(
|
|
self, gpt_config, random_inputs, attention_mask
|
|
):
|
|
"""Test forward pass with attention mask."""
|
|
model = GPT(gpt_config)
|
|
|
|
# Forward pass with mask
|
|
logits, _ = model(random_inputs, attention_mask=attention_mask)
|
|
|
|
# Check output shape
|
|
batch_size, seq_len = random_inputs.shape
|
|
vocab_size = gpt_config["vocab_size"]
|
|
assert logits.shape == (batch_size, seq_len, vocab_size)
|
|
|
|
def test_generate_text(self, gpt_config):
|
|
"""Test text generation."""
|
|
model = GPT(gpt_config)
|
|
model.eval() # Set to evaluation mode for generation
|
|
|
|
# Create initial input
|
|
batch_size = 2
|
|
initial_seq_len = 5
|
|
input_ids = torch.randint(
|
|
0, gpt_config["vocab_size"], (batch_size, initial_seq_len)
|
|
)
|
|
|
|
# Generate text
|
|
with torch.no_grad():
|
|
generated = model.generate(
|
|
x=input_ids,
|
|
max_new_tokens=10,
|
|
do_sample=False, # Use greedy for deterministic testing
|
|
)
|
|
|
|
# Check output shape
|
|
expected_seq_len = initial_seq_len + 10
|
|
assert generated.shape == (batch_size, expected_seq_len)
|
|
|
|
# Check that initial sequence is preserved
|
|
assert torch.allclose(generated[:, :initial_seq_len], input_ids)
|
|
|
|
def test_generate_with_temperature(self, gpt_config):
|
|
"""Test text generation with temperature sampling."""
|
|
model = GPT(gpt_config)
|
|
model.eval()
|
|
|
|
# Create initial input
|
|
input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3))
|
|
|
|
# Generate with temperature
|
|
with torch.no_grad():
|
|
generated = model.generate(
|
|
x=input_ids, max_new_tokens=5, do_sample=True, temperature=0.8
|
|
)
|
|
|
|
assert generated.shape == (1, 8) # 3 initial + 5 new tokens
|
|
|
|
def test_generate_with_top_k(self, gpt_config):
|
|
"""Test text generation with top-k sampling."""
|
|
model = GPT(gpt_config)
|
|
model.eval()
|
|
|
|
# Create initial input
|
|
input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3))
|
|
|
|
# Generate with top-k
|
|
with torch.no_grad():
|
|
generated = model.generate(
|
|
x=input_ids, max_new_tokens=5, do_sample=True, top_k=10
|
|
)
|
|
|
|
assert generated.shape == (1, 8)
|
|
|
|
def test_generate_with_top_p(self, gpt_config):
|
|
"""Test text generation with top-p (nucleus) sampling."""
|
|
model = GPT(gpt_config)
|
|
model.eval()
|
|
|
|
# Create initial input
|
|
input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3))
|
|
|
|
# Generate with top-p
|
|
with torch.no_grad():
|
|
generated = model.generate(
|
|
x=input_ids, max_new_tokens=5, do_sample=True, top_p=0.9
|
|
)
|
|
|
|
assert generated.shape == (1, 8)
|
|
|
|
def test_gradient_flow(self, gpt_config, random_inputs):
|
|
"""Test that gradients flow through GPT."""
|
|
model = GPT(gpt_config)
|
|
|
|
# Forward pass
|
|
logits, _ = model(random_inputs)
|
|
|
|
# Create a dummy loss and backward pass
|
|
targets = torch.randint(0, gpt_config["vocab_size"], random_inputs.shape)
|
|
loss = torch.nn.functional.cross_entropy(
|
|
logits.view(-1, logits.size(-1)), targets.view(-1)
|
|
)
|
|
loss.backward()
|
|
|
|
# Check that gradients are computed for various components
|
|
assert model._token_embeddings._embedding.weight.grad is not None
|
|
assert model._linear.weight.grad is not None
|
|
if len(model._decoders) > 0:
|
|
# Проверяем через новый интерфейс attention оптимизации:
|
|
attn = model._decoders[0]._heads
|
|
assert attn._q.weight.grad is not None
|
|
assert attn._k.weight.grad is not None
|
|
assert attn._v.weight.grad is not None
|
|
|
|
def test_device_consistency(self, gpt_config, random_inputs, device):
|
|
"""Test that GPT works on correct device."""
|
|
model = GPT(gpt_config).to(device)
|
|
inputs = random_inputs.to(device)
|
|
|
|
# Forward pass
|
|
logits, _ = model(inputs)
|
|
|
|
# Check device consistency
|
|
assert logits.device == device
|
|
assert model._token_embeddings._embedding.weight.device == device
|
|
|
|
def test_different_configurations(self):
|
|
"""Test GPT with different configurations."""
|
|
test_configs = [
|
|
{
|
|
"vocab_size": 1000,
|
|
"embed_dim": 128,
|
|
"num_heads": 2,
|
|
"num_layers": 2,
|
|
"max_position_embeddings": 256,
|
|
"dropout": 0.1,
|
|
},
|
|
{
|
|
"vocab_size": 5000,
|
|
"embed_dim": 256,
|
|
"num_heads": 4,
|
|
"num_layers": 4,
|
|
"max_position_embeddings": 512,
|
|
"dropout": 0.1,
|
|
},
|
|
{
|
|
"vocab_size": 10000,
|
|
"embed_dim": 512,
|
|
"num_heads": 8,
|
|
"num_layers": 6,
|
|
"max_position_embeddings": 1024,
|
|
"dropout": 0.1,
|
|
},
|
|
]
|
|
|
|
for config in test_configs:
|
|
model = GPT(config)
|
|
batch_size, seq_len = 2, 16
|
|
inputs = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
|
|
|
|
logits, _ = model(inputs)
|
|
|
|
expected_shape = (batch_size, seq_len, config["vocab_size"])
|
|
assert logits.shape == expected_shape
|
|
|
|
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)])
|
|
def test_different_input_shapes(self, gpt_config, batch_size, seq_len):
|
|
"""Test GPT with different input shapes."""
|
|
model = GPT(gpt_config)
|
|
|
|
inputs = torch.randint(0, gpt_config["vocab_size"], (batch_size, seq_len))
|
|
logits, _ = model(inputs)
|
|
|
|
expected_shape = (batch_size, seq_len, gpt_config["vocab_size"])
|
|
assert logits.shape == expected_shape
|
|
|
|
def test_training_vs_evaluation(self, gpt_config, random_inputs):
|
|
"""Test that GPT behaves differently in train vs eval mode."""
|
|
model = GPT(gpt_config)
|
|
|
|
# Training mode
|
|
model.train()
|
|
output_train, _ = model(random_inputs)
|
|
|
|
# Evaluation mode
|
|
model.eval()
|
|
output_eval, _ = model(random_inputs)
|
|
|
|
# Outputs should be different due to dropout
|
|
assert not torch.allclose(output_train, output_eval)
|
|
|
|
def test_parameter_count(self, gpt_config):
|
|
"""Test that GPT has reasonable number of parameters."""
|
|
model = GPT(gpt_config)
|
|
|
|
total_params = sum(p.numel() for p in model.parameters())
|
|
|
|
# For a small GPT model, parameters should be in reasonable range
|
|
vocab_size = gpt_config["vocab_size"]
|
|
embed_dim = gpt_config["embed_dim"]
|
|
num_layers = gpt_config["num_layers"]
|
|
num_heads = gpt_config["num_heads"]
|
|
|
|
# Rough estimate: token_embeddings + output_layer + (attention + ff) * layers
|
|
expected_min = vocab_size * embed_dim * 2 # embeddings and output
|
|
expected_max = expected_min * 10 # Allow for decoder parameters
|
|
|
|
assert expected_min < total_params < expected_max
|
|
|
|
def test_causal_attention(self, gpt_config):
|
|
"""Test that GPT uses causal attention during generation."""
|
|
model = GPT(gpt_config)
|
|
model.eval()
|
|
|
|
# Create input with known pattern
|
|
input_ids = torch.tensor([[1, 2, 3]]).long()
|
|
|
|
with torch.no_grad():
|
|
# Get logits for next token prediction
|
|
logits = model(input_ids)
|
|
|
|
# The model should only attend to previous tokens (causal)
|
|
# We can't directly test attention masks in the public API,
|
|
# but we can verify the generation works correctly
|
|
|
|
generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False)
|
|
|
|
# Generated sequence should be longer than input
|
|
assert generated.shape[1] == input_ids.shape[1] + 3
|
|
|
|
def test_output_distribution(self, gpt_config, random_inputs):
|
|
"""Test that GPT output has proper distribution."""
|
|
model = GPT(gpt_config)
|
|
|
|
logits, _ = model(random_inputs)
|
|
|
|
# Logits should not have extreme values
|
|
assert logits.abs().max() < 100
|
|
|
|
# Softmax should produce valid probabilities
|
|
probs = torch.softmax(logits, dim=-1)
|
|
assert torch.allclose(probs.sum(dim=-1), torch.ones_like(probs.sum(dim=-1)))
|
|
assert (probs >= 0).all() and (probs <= 1).all()
|