mirror of
https://github.com/pese-git/llm-arch-research.git
synced 2026-01-23 21:10:54 +00:00
feat(gemma): initial implementation of Gemma model and configs
- Add core Gemma model (architecture, attention, GeGLU, RoPE, RMSNorm, etc) - Add configs for training and generation: gemma_train.json, gemma_generate.json - Add Gemma notebook for exploratory analysis and demonstration - Add __init__.py for Gemma submodule - Update run_llm_experiment.py to support Gemma experiment configs test(gemma): add comprehensive unit tests for Gemma - Test forward pass (with/without cache) - Test autoregressive generation (greedy, top-k, top-p) - Test shape correctness and max sequence length errors - Test multi-layer stack and token embeddings docs: add documentation notebook for Gemma usage and analysis Closes: #issue (if applicable)
This commit is contained in:
56
llm/tests/models/test_gemma.py
Normal file
56
llm/tests/models/test_gemma.py
Normal file
@@ -0,0 +1,56 @@
|
||||
# llm/tests/models/test_gemma.py
|
||||
|
||||
import torch
|
||||
import pytest
|
||||
from llm.models.gemma.gemma import Gemma
|
||||
|
||||
@pytest.fixture
|
||||
def config():
|
||||
return {
|
||||
"vocab_size": 100,
|
||||
"embed_dim": 32,
|
||||
"num_q_heads": 4,
|
||||
"num_layers": 2,
|
||||
"max_position_embeddings": 16,
|
||||
"dropout": 0.0,
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def model(config):
|
||||
return Gemma(config)
|
||||
|
||||
def test_forward_basic(model):
|
||||
x = torch.randint(0, 100, (2, 8))
|
||||
logits, cache = model(x)
|
||||
assert logits.shape == (2, 8, 100)
|
||||
assert isinstance(cache, list)
|
||||
assert len(cache) == model._decoders.__len__()
|
||||
|
||||
def test_forward_with_cache(model):
|
||||
x = torch.randint(0, 100, (2, 4))
|
||||
logits, cache = model(x, use_cache=True)
|
||||
# Второй проход с cache и одним новым токеном
|
||||
x2 = torch.randint(0, 100, (2, 1))
|
||||
logits2, cache2 = model(x2, use_cache=True, cache=cache)
|
||||
assert logits2.shape == (2, 1, 100)
|
||||
assert isinstance(cache2, list)
|
||||
|
||||
def test_generate_and_shape(model):
|
||||
x = torch.randint(0, 100, (1, 5))
|
||||
result = model.generate(x, max_new_tokens=3, do_sample=False)
|
||||
assert result.shape == (1, 8)
|
||||
|
||||
def test_forward_sequence_too_long(model, config):
|
||||
x = torch.randint(0, 100, (1, config["max_position_embeddings"] + 1))
|
||||
with pytest.raises(ValueError):
|
||||
model(x)
|
||||
|
||||
def test_generate_with_sampling_topk(model):
|
||||
x = torch.randint(0, 100, (1, 3))
|
||||
out = model.generate(x, max_new_tokens=2, do_sample=True, top_k=5)
|
||||
assert out.shape == (1, 5)
|
||||
|
||||
def test_generate_with_sampling_topp(model):
|
||||
x = torch.randint(0, 100, (1, 3))
|
||||
out = model.generate(x, max_new_tokens=2, do_sample=True, top_p=0.8)
|
||||
assert out.shape == (1, 5)
|
||||
Reference in New Issue
Block a user