File size: 5,482 Bytes
2d1ceea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """
Test Complete AngstromNano Model
"""
import sys
sys.path.insert(0, '.')
import torch
from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM
def test_model_initialization():
print("\n[Testing Model Initialization]")
config = AngstromNanoConfig()
model = AngstromNanoForCausalLM(config)
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f" Total parameters: {total_params:,} ({total_params/1e9:.2f}B)")
print(f" Trainable parameters: {trainable_params:,}")
print(f" Config estimate: {config.estimate_parameters()['total_billions']:.2f}B")
print(" [PASS]")
return model
def test_forward_pass():
print("\n[Testing Forward Pass]")
config = AngstromNanoConfig()
model = AngstromNanoForCausalLM(config)
batch_size, seq_len = 2, 128
input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
outputs = model(input_ids, labels=labels)
print(f" Input shape: {input_ids.shape}")
print(f" Logits shape: {outputs['logits'].shape}")
print(f" Loss: {outputs['loss'].item():.4f}")
print(f" Aux loss: {outputs['aux_loss'].item():.6f}")
assert outputs['logits'].shape == (batch_size, seq_len, config.vocab_size)
assert outputs['loss'] is not None
print(" [PASS]")
def test_generation():
print("\n[Testing Text Generation]")
config = AngstromNanoConfig()
model = AngstromNanoForCausalLM(config)
# Simple prompt
prompt = torch.tensor([[1, 100, 200, 300]]) # BOS + 3 tokens
generated = model.generate(
prompt,
max_new_tokens=20,
temperature=1.0,
)
print(f" Prompt length: {prompt.shape[1]}")
print(f" Generated length: {generated.shape[1]}")
print(f" Generated tokens: {generated[0].tolist()[:10]}...")
assert generated.shape[1] > prompt.shape[1]
print(" [PASS]")
def test_kv_cache():
print("\n[Testing KV Cache]")
config = AngstromNanoConfig()
model = AngstromNanoForCausalLM(config)
batch_size, seq_len = 1, 16
input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
# First forward pass with cache
outputs1 = model(input_ids, use_cache=True)
past_kv = outputs1['past_key_values']
print(f" Cached layers: {len(past_kv)}")
print(f" Cached K shape: {past_kv[0][0].shape}")
print(f" Cached V shape: {past_kv[0][1].shape}")
# Second forward pass using cache
next_token = torch.randint(0, config.vocab_size, (batch_size, 1))
outputs2 = model(next_token, past_key_values=past_kv, use_cache=True)
print(f" New cached K shape: {outputs2['past_key_values'][0][0].shape}")
assert len(past_kv) == config.num_hidden_layers
assert outputs2['past_key_values'][0][0].shape[2] == seq_len + 1
print(" [PASS]")
def test_gradient_flow():
print("\n[Testing Gradient Flow]")
config = AngstromNanoConfig()
model = AngstromNanoForCausalLM(config)
batch_size, seq_len = 2, 32
input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
outputs = model(input_ids, labels=labels)
loss = outputs['loss']
loss.backward()
# Check gradients in different parts
embed_grad = model.model.embed_tokens.weight.grad
first_layer_grad = model.model.layers[0].self_attn.q_proj.weight.grad
last_layer_grad = model.model.layers[-1].self_attn.q_proj.weight.grad
lm_head_grad = model.lm_head.weight.grad
print(f" Loss: {loss.item():.4f}")
print(f" Embedding grad norm: {embed_grad.norm().item():.6f}")
print(f" First layer grad norm: {first_layer_grad.norm().item():.6f}")
print(f" Last layer grad norm: {last_layer_grad.norm().item():.6f}")
print(f" LM head grad norm: {lm_head_grad.norm().item():.6f}")
assert all(g is not None for g in [embed_grad, first_layer_grad, last_layer_grad, lm_head_grad])
print(" [PASS]")
def test_memory_usage():
print("\n[Testing Memory Usage]")
config = AngstromNanoConfig()
model = AngstromNanoForCausalLM(config)
# Calculate model size
param_size = sum(p.numel() * p.element_size() for p in model.parameters())
buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
total_size = param_size + buffer_size
print(f" Parameter memory: {param_size / 1e9:.2f} GB")
print(f" Buffer memory: {buffer_size / 1e6:.2f} MB")
print(f" Total model size: {total_size / 1e9:.2f} GB")
# Expected: ~8GB for FP32, ~4GB for FP16
expected_fp32 = config.estimate_parameters()['estimated_size_fp32_gb']
print(f" Expected FP32 size: {expected_fp32:.2f} GB")
assert abs(total_size / 1e9 - expected_fp32) < 1.0 # Within 1GB tolerance
print(" [PASS]")
def main():
print("=" * 80)
print("Testing Complete AngstromNano Model")
print("=" * 80)
torch.manual_seed(42)
test_model_initialization()
test_forward_pass()
test_generation()
test_kv_cache()
test_gradient_flow()
test_memory_usage()
print("\n" + "=" * 80)
print("All model tests passed!")
print("=" * 80)
if __name__ == "__main__":
main()
|