| """ |
| 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) |
| |
| |
| prompt = torch.tensor([[1, 100, 200, 300]]) |
| |
| 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)) |
| |
| |
| 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}") |
| |
| |
| 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() |
| |
| |
| 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) |
| |
| |
| 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_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 |
| 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() |
|
|
