File size: 4,184 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 | """
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)
# Create causal mask
attention_mask = torch.triu(torch.ones(seq_len, seq_len) * float('-inf'), diagonal=1)
attention_mask = attention_mask.unsqueeze(0).unsqueeze(0) # [1, 1, seq_len, seq_len]
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 # cached sequence length
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)
# Forward pass
attn_out, _ = attn(hidden_states, position_ids=position_ids)
mlp_out = mlp(attn_out)
# Compute loss and backward
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()
|