| """ |
| Test Core Transformer Components |
| """ |
|
|
| import sys |
| sys.path.insert(0, '.') |
|
|
| import torch |
| from angstrom_nano import AngstromNanoConfig |
| from angstrom_nano.model import ( |
| AngstromNanoRMSNorm, |
| AngstromNanoRotaryEmbedding, |
| AngstromNanoAttention, |
| AngstromNanoMLP |
| ) |
|
|
|
|
| def test_rms_norm(): |
| print("\n[Testing RMSNorm]") |
| config = AngstromNanoConfig() |
| norm = AngstromNanoRMSNorm(config.hidden_size, config.rms_norm_eps) |
| |
| x = torch.randn(2, 512, config.hidden_size) |
| out = norm(x) |
| |
| print(f" Input shape: {x.shape}") |
| print(f" Output shape: {out.shape}") |
| print(f" Output mean: {out.mean():.6f}, std: {out.std():.6f}") |
| assert out.shape == x.shape |
| print(" [PASS]") |
|
|
|
|
| def test_rotary_embedding(): |
| print("\n[Testing Rotary Embedding]") |
| config = AngstromNanoConfig() |
| rope = AngstromNanoRotaryEmbedding( |
| config.head_dim, |
| config.max_position_embeddings, |
| config.rope_theta |
| ) |
| |
| batch_size, seq_len = 2, 512 |
| x = torch.randn(batch_size, config.num_attention_heads, seq_len, config.head_dim) |
| position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) |
| |
| cos, sin = rope(x, position_ids) |
| |
| print(f" Input shape: {x.shape}") |
| print(f" Position IDs shape: {position_ids.shape}") |
| print(f" Cos shape: {cos.shape}") |
| print(f" Sin shape: {sin.shape}") |
| assert cos.shape == (batch_size, seq_len, config.head_dim) |
| assert sin.shape == (batch_size, seq_len, config.head_dim) |
| print(" [PASS]") |
|
|
|
|
| def test_attention(): |
| print("\n[Testing Multi-Head Attention]") |
| config = AngstromNanoConfig() |
| attn = AngstromNanoAttention(config, layer_idx=0) |
| |
| batch_size, seq_len = 2, 512 |
| hidden_states = torch.randn(batch_size, seq_len, config.hidden_size) |
| position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) |
| |
| |
| attention_mask = torch.triu(torch.ones(seq_len, seq_len) * float('-inf'), diagonal=1) |
| attention_mask = attention_mask.unsqueeze(0).unsqueeze(0) |
| |
| output, past_kv = attn( |
| hidden_states, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| use_cache=True |
| ) |
| |
| print(f" Input shape: {hidden_states.shape}") |
| print(f" Output shape: {output.shape}") |
| print(f" KV cache shapes: {past_kv[0].shape}, {past_kv[1].shape}") |
| assert output.shape == hidden_states.shape |
| assert past_kv[0].shape[2] == seq_len |
| print(" [PASS]") |
|
|
|
|
| def test_mlp(): |
| print("\n[Testing Feed-Forward Network]") |
| config = AngstromNanoConfig() |
| mlp = AngstromNanoMLP(config) |
| |
| batch_size, seq_len = 2, 512 |
| x = torch.randn(batch_size, seq_len, config.hidden_size) |
| |
| output = mlp(x) |
| |
| print(f" Input shape: {x.shape}") |
| print(f" Output shape: {output.shape}") |
| assert output.shape == x.shape |
| print(" [PASS]") |
|
|
|
|
| def test_gradient_flow(): |
| print("\n[Testing Gradient Flow]") |
| config = AngstromNanoConfig() |
| attn = AngstromNanoAttention(config, layer_idx=0) |
| mlp = AngstromNanoMLP(config) |
| |
| batch_size, seq_len = 2, 128 |
| hidden_states = torch.randn(batch_size, seq_len, config.hidden_size, requires_grad=True) |
| position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) |
| |
| |
| attn_out, _ = attn(hidden_states, position_ids=position_ids) |
| mlp_out = mlp(attn_out) |
| |
| |
| loss = mlp_out.mean() |
| loss.backward() |
| |
| print(f" Loss: {loss.item():.6f}") |
| print(f" Input gradient exists: {hidden_states.grad is not None}") |
| print(f" Input gradient norm: {hidden_states.grad.norm().item():.6f}") |
| assert hidden_states.grad is not None |
| print(" [PASS]") |
|
|
|
|
| def main(): |
| print("=" * 80) |
| print("Testing Core Transformer Components") |
| print("=" * 80) |
| |
| torch.manual_seed(42) |
| |
| test_rms_norm() |
| test_rotary_embedding() |
| test_attention() |
| test_mlp() |
| test_gradient_flow() |
| |
| print("\n" + "=" * 80) |
| print("All tests passed!") |
| print("=" * 80) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
