mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-19 12:16:13 -05:00
Vijay Janapa Reddi
ec7168dc90
Added all module development files to modules/XX_name/ directories:
Module notebooks and scripts:
- 18 modules with .ipynb and .py files (01-20, excluding some gaps)
- Moved from modules/source/ to direct module directories
- Includes tensor, autograd, layers, transformers, optimization modules
Module README files:
- Added README.md for modules with additional documentation
- Complements ABOUT.md files added earlier
This completes the module restructuring:
- Before: modules/source/XX_name/*_dev.{py,ipynb}
- After: modules/XX_name/*_dev.{py,ipynb}
All development happens directly in numbered module directories now.
100 KiB
100 KiB
In [ ]:
#| default_exp models.transformerIn [ ]:
#| export
import numpy as np
from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Linear
from tinytorch.core.attention import MultiHeadAttention
from tinytorch.core.activations import GELUIn [ ]:
import numpy as np
import math
from typing import Optional, List
# Import from previous modules - following proper dependency chain
# Note: Actual imports happen in try/except blocks below with fallback implementations
from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Linear
# MultiHeadAttention import happens in try/except below
# For development, we'll use minimal implementations if imports fail
try:
from tinytorch.core.tensor import Tensor
except ImportError:
print("Warning: Using minimal Tensor implementation for development")
class Tensor:
"""Minimal Tensor class for transformer development."""
def __init__(self, data, requires_grad=False):
self.data = np.array(data)
self.shape = self.data.shape
self.size = self.data.size
self.requires_grad = requires_grad
self.grad = None
def __add__(self, other):
if isinstance(other, Tensor):
return Tensor(self.data + other.data)
return Tensor(self.data + other)
def __mul__(self, other):
if isinstance(other, Tensor):
return Tensor(self.data * other.data)
return Tensor(self.data * other)
def matmul(self, other):
return Tensor(np.dot(self.data, other.data))
def sum(self, axis=None, keepdims=False):
return Tensor(self.data.sum(axis=axis, keepdims=keepdims))
def mean(self, axis=None, keepdims=False):
return Tensor(self.data.mean(axis=axis, keepdims=keepdims))
def reshape(self, *shape):
return Tensor(self.data.reshape(shape))
def __repr__(self):
return f"Tensor(data={self.data}, shape={self.shape})"
try:
from tinytorch.core.layers import Linear
except ImportError:
class Linear:
"""Minimal Linear layer for development."""
def __init__(self, in_features, out_features, bias=True):
std = math.sqrt(2.0 / (in_features + out_features))
self.weight = Tensor(np.random.normal(0, std, (in_features, out_features)))
self.bias = Tensor(np.zeros(out_features)) if bias else None
def forward(self, x):
output = x.matmul(self.weight)
if self.bias is not None:
output = output + self.bias
return output
def parameters(self):
params = [self.weight]
if self.bias is not None:
params.append(self.bias)
return params
try:
from tinytorch.core.attention import MultiHeadAttention
except ImportError:
class MultiHeadAttention:
"""Minimal MultiHeadAttention for development."""
def __init__(self, embed_dim, num_heads):
assert embed_dim % num_heads == 0
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.q_proj = Linear(embed_dim, embed_dim)
self.k_proj = Linear(embed_dim, embed_dim)
self.v_proj = Linear(embed_dim, embed_dim)
self.out_proj = Linear(embed_dim, embed_dim)
def forward(self, query, key, value, mask=None):
batch_size, seq_len, embed_dim = query.shape
# Linear projections
Q = self.q_proj.forward(query)
K = self.k_proj.forward(key)
V = self.v_proj.forward(value)
# Reshape for multi-head attention
Q = Q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
K = K.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
V = V.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
# Transpose to (batch_size, num_heads, seq_len, head_dim)
Q = Tensor(np.transpose(Q.data, (0, 2, 1, 3)))
K = Tensor(np.transpose(K.data, (0, 2, 1, 3)))
V = Tensor(np.transpose(V.data, (0, 2, 1, 3)))
# Scaled dot-product attention
scores = Tensor(np.matmul(Q.data, np.transpose(K.data, (0, 1, 3, 2))))
scores = scores * (1.0 / math.sqrt(self.head_dim))
# Apply causal mask for autoregressive generation
if mask is not None:
scores = Tensor(scores.data + mask.data)
# Softmax
attention_weights = self._softmax(scores)
# Apply attention to values
out = Tensor(np.matmul(attention_weights.data, V.data))
# Transpose back and reshape
out = Tensor(np.transpose(out.data, (0, 2, 1, 3)))
out = out.reshape(batch_size, seq_len, embed_dim)
# Final linear projection
return self.out_proj.forward(out)
def _softmax(self, x):
"""Numerically stable softmax."""
exp_x = Tensor(np.exp(x.data - np.max(x.data, axis=-1, keepdims=True)))
return Tensor(exp_x.data / np.sum(exp_x.data, axis=-1, keepdims=True))
def parameters(self):
params = []
params.extend(self.q_proj.parameters())
params.extend(self.k_proj.parameters())
params.extend(self.v_proj.parameters())
params.extend(self.out_proj.parameters())
return params
try:
from tinytorch.core.embeddings import Embedding
except ImportError:
class Embedding:
"""Minimal Embedding layer for development."""
def __init__(self, vocab_size, embed_dim):
self.vocab_size = vocab_size
self.embed_dim = embed_dim
self.weight = Tensor(np.random.normal(0, 0.02, (vocab_size, embed_dim)))
def forward(self, indices):
return Tensor(self.weight.data[indices.data.astype(int)])
def parameters(self):
return [self.weight]
def gelu(x):
"""GELU activation function."""
return Tensor(0.5 * x.data * (1 + np.tanh(np.sqrt(2 / np.pi) * (x.data + 0.044715 * x.data**3))))In [ ]:
#| export
class LayerNorm:
"""
Layer Normalization for transformer blocks.
Normalizes across the feature dimension (last axis) for each sample independently,
unlike batch normalization which normalizes across the batch dimension.
"""
def __init__(self, normalized_shape, eps=1e-5):
"""
Initialize LayerNorm with learnable parameters.
TODO: Set up normalization parameters
APPROACH:
1. Store the shape to normalize over (usually embed_dim)
2. Initialize learnable scale (gamma) and shift (beta) parameters
3. Set small epsilon for numerical stability
EXAMPLE:
>>> ln = LayerNorm(512) # For 512-dimensional embeddings
>>> x = Tensor(np.random.randn(2, 10, 512)) # (batch, seq, features)
>>> normalized = ln.forward(x)
>>> # Each (2, 10) sample normalized independently across 512 features
HINTS:
- gamma should start at 1.0 (identity scaling)
- beta should start at 0.0 (no shift)
- eps prevents division by zero in variance calculation
"""
### BEGIN SOLUTION
self.normalized_shape = normalized_shape
self.eps = eps
# Learnable parameters: scale and shift
# CRITICAL: requires_grad=True so optimizer can train these!
self.gamma = Tensor(np.ones(normalized_shape), requires_grad=True) # Scale parameter
self.beta = Tensor(np.zeros(normalized_shape), requires_grad=True) # Shift parameter
### END SOLUTION
def forward(self, x):
"""
Apply layer normalization.
TODO: Implement layer normalization formula
APPROACH:
1. Compute mean and variance across the last dimension
2. Normalize: (x - mean) / sqrt(variance + eps)
3. Apply learnable scale and shift: gamma * normalized + beta
MATHEMATICAL FORMULA:
y = (x - μ) / σ * γ + β
where μ = mean(x), σ = sqrt(var(x) + ε)
HINT: Use keepdims=True to maintain tensor dimensions for broadcasting
"""
### BEGIN SOLUTION
# CRITICAL: Use Tensor operations (not .data) to maintain gradient flow!
# Compute statistics across last dimension (features)
mean = x.mean(axis=-1, keepdims=True)
# Compute variance: E[(x - μ)²]
diff = x - mean # Tensor subtraction maintains gradient
variance = (diff * diff).mean(axis=-1, keepdims=True) # Tensor ops maintain gradient
# Normalize: (x - mean) / sqrt(variance + eps)
# Note: sqrt and division need to preserve gradient flow
std_data = np.sqrt(variance.data + self.eps)
normalized = diff * Tensor(1.0 / std_data) # Scale by reciprocal to maintain gradient
# Apply learnable transformation
output = normalized * self.gamma + self.beta
return output
### END SOLUTION
def parameters(self):
"""Return learnable parameters."""
return [self.gamma, self.beta]In [ ]:
def test_unit_layer_norm():
"""🔬 Test LayerNorm implementation."""
print("🔬 Unit Test: Layer Normalization...")
# Test basic normalization
ln = LayerNorm(4)
x = Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]) # (2, 4)
normalized = ln.forward(x)
# Check output shape
assert normalized.shape == (2, 4)
# Check normalization properties (approximately)
# For each sample, mean should be close to 0, std close to 1
for i in range(2):
sample_mean = np.mean(normalized.data[i])
sample_std = np.std(normalized.data[i])
assert abs(sample_mean) < 1e-5, f"Mean should be ~0, got {sample_mean}"
assert abs(sample_std - 1.0) < 1e-4, f"Std should be ~1, got {sample_std}"
# Test parameter shapes
params = ln.parameters()
assert len(params) == 2
assert params[0].shape == (4,) # gamma
assert params[1].shape == (4,) # beta
print("✅ LayerNorm works correctly!")
# Run test immediately when developing this module
if __name__ == "__main__":
test_unit_layer_norm()In [ ]:
#| export
class MLP:
"""
Multi-Layer Perceptron (Feed-Forward Network) for transformer blocks.
Standard pattern: Linear -> GELU -> Linear with expansion ratio of 4:1.
This provides the non-linear transformation in each transformer block.
"""
def __init__(self, embed_dim, hidden_dim=None, dropout_prob=0.1):
"""
Initialize MLP with two linear layers.
TODO: Set up the feed-forward network layers
APPROACH:
1. First layer expands from embed_dim to hidden_dim (usually 4x larger)
2. Second layer projects back to embed_dim
3. Use GELU activation (smoother than ReLU, preferred in transformers)
EXAMPLE:
>>> mlp = MLP(512) # Will create 512 -> 2048 -> 512 network
>>> x = Tensor(np.random.randn(2, 10, 512))
>>> output = mlp.forward(x)
>>> assert output.shape == (2, 10, 512)
HINT: Standard transformer MLP uses 4x expansion (hidden_dim = 4 * embed_dim)
"""
### BEGIN SOLUTION
if hidden_dim is None:
hidden_dim = 4 * embed_dim # Standard 4x expansion
self.embed_dim = embed_dim
self.hidden_dim = hidden_dim
# Two-layer feed-forward network
self.linear1 = Linear(embed_dim, hidden_dim)
self.linear2 = Linear(hidden_dim, embed_dim)
### END SOLUTION
def forward(self, x):
"""
Forward pass through MLP.
TODO: Implement the feed-forward computation
APPROACH:
1. First linear transformation: embed_dim -> hidden_dim
2. Apply GELU activation (smooth, differentiable)
3. Second linear transformation: hidden_dim -> embed_dim
COMPUTATION FLOW:
x -> Linear -> GELU -> Linear -> output
HINT: GELU activation is implemented above as a function
"""
### BEGIN SOLUTION
# First linear layer with expansion
hidden = self.linear1.forward(x)
# GELU activation
hidden = gelu(hidden)
# Second linear layer back to original size
output = self.linear2.forward(hidden)
return output
### END SOLUTION
def parameters(self):
"""Return all learnable parameters."""
params = []
params.extend(self.linear1.parameters())
params.extend(self.linear2.parameters())
return paramsIn [ ]:
def test_unit_mlp():
"""🔬 Test MLP implementation."""
print("🔬 Unit Test: MLP (Feed-Forward Network)...")
# Test MLP with standard 4x expansion
embed_dim = 64
mlp = MLP(embed_dim)
# Test forward pass
batch_size, seq_len = 2, 10
x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))
output = mlp.forward(x)
# Check shape preservation
assert output.shape == (batch_size, seq_len, embed_dim)
# Check hidden dimension is 4x
assert mlp.hidden_dim == 4 * embed_dim
# Test parameter counting
params = mlp.parameters()
expected_params = 4 # 2 weights + 2 biases
assert len(params) == expected_params
# Test custom hidden dimension
custom_mlp = MLP(embed_dim, hidden_dim=128)
assert custom_mlp.hidden_dim == 128
print("✅ MLP works correctly!")
# Run test immediately when developing this module
if __name__ == "__main__":
test_unit_mlp()In [ ]:
#| export
class TransformerBlock:
"""
Complete Transformer Block with self-attention, MLP, and residual connections.
This is the core building block of GPT and other transformer models.
Each block processes the input sequence and passes it to the next block.
"""
def __init__(self, embed_dim, num_heads, mlp_ratio=4, dropout_prob=0.1):
"""
Initialize a complete transformer block.
TODO: Set up all components of the transformer block
APPROACH:
1. Multi-head self-attention for sequence modeling
2. First layer normalization (pre-norm architecture)
3. MLP with specified expansion ratio
4. Second layer normalization
TRANSFORMER BLOCK ARCHITECTURE:
x → LayerNorm → MultiHeadAttention → + (residual) →
LayerNorm → MLP → + (residual) → output
EXAMPLE:
>>> block = TransformerBlock(embed_dim=512, num_heads=8)
>>> x = Tensor(np.random.randn(2, 10, 512)) # (batch, seq, embed)
>>> output = block.forward(x)
>>> assert output.shape == (2, 10, 512)
HINT: We use pre-norm architecture (LayerNorm before attention/MLP)
"""
### BEGIN SOLUTION
self.embed_dim = embed_dim
self.num_heads = num_heads
# Multi-head self-attention
self.attention = MultiHeadAttention(embed_dim, num_heads)
# Layer normalizations (pre-norm architecture)
self.ln1 = LayerNorm(embed_dim) # Before attention
self.ln2 = LayerNorm(embed_dim) # Before MLP
# Feed-forward network
hidden_dim = int(embed_dim * mlp_ratio)
self.mlp = MLP(embed_dim, hidden_dim)
### END SOLUTION
def forward(self, x, mask=None):
"""
Forward pass through transformer block.
TODO: Implement the complete transformer block computation
APPROACH:
1. Apply layer norm, then self-attention, then add residual
2. Apply layer norm, then MLP, then add residual
3. Return the transformed sequence
COMPUTATION FLOW:
x → ln1 → attention → + x → ln2 → mlp → + → output
RESIDUAL CONNECTIONS:
These are crucial for training deep networks - they allow gradients
to flow directly through the network during backpropagation.
HINT: Store intermediate results to add residual connections properly
"""
### BEGIN SOLUTION
# First sub-layer: Multi-head self-attention with residual connection
# Pre-norm: LayerNorm before attention
normed1 = self.ln1.forward(x)
# Self-attention: query, key, value are all the same (normed1)
attention_out = self.attention.forward(normed1, normed1, normed1, mask)
# Residual connection
x = x + attention_out
# Second sub-layer: MLP with residual connection
# Pre-norm: LayerNorm before MLP
normed2 = self.ln2.forward(x)
mlp_out = self.mlp.forward(normed2)
# Residual connection
output = x + mlp_out
return output
### END SOLUTION
def parameters(self):
"""Return all learnable parameters."""
params = []
params.extend(self.attention.parameters())
params.extend(self.ln1.parameters())
params.extend(self.ln2.parameters())
params.extend(self.mlp.parameters())
return paramsIn [ ]:
def test_unit_transformer_block():
"""🔬 Test TransformerBlock implementation."""
print("🔬 Unit Test: Transformer Block...")
# Test transformer block
embed_dim = 64
num_heads = 4
block = TransformerBlock(embed_dim, num_heads)
# Test forward pass
batch_size, seq_len = 2, 8
x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))
output = block.forward(x)
# Check shape preservation
assert output.shape == (batch_size, seq_len, embed_dim)
# Test with causal mask (for autoregressive generation)
mask = Tensor(np.triu(np.ones((seq_len, seq_len)) * -np.inf, k=1))
masked_output = block.forward(x, mask)
assert masked_output.shape == (batch_size, seq_len, embed_dim)
# Test parameter counting
params = block.parameters()
expected_components = 4 # attention, ln1, ln2, mlp parameters
assert len(params) > expected_components # Should have parameters from all components
# Test different configurations
large_block = TransformerBlock(embed_dim=128, num_heads=8, mlp_ratio=2)
assert large_block.mlp.hidden_dim == 256 # 128 * 2
print("✅ TransformerBlock works correctly!")
# Run test immediately when developing this module
if __name__ == "__main__":
test_unit_transformer_block()In [ ]:
#| export
class GPT:
"""
Complete GPT (Generative Pre-trained Transformer) model.
This combines embeddings, positional encoding, multiple transformer blocks,
and a language modeling head for text generation.
"""
def __init__(self, vocab_size, embed_dim, num_layers, num_heads, max_seq_len=1024):
"""
Initialize complete GPT model.
TODO: Set up all components of the GPT architecture
APPROACH:
1. Token embedding layer to convert tokens to vectors
2. Positional embedding to add position information
3. Stack of transformer blocks (the main computation)
4. Final layer norm and language modeling head
GPT ARCHITECTURE:
tokens → embedding → + pos_embedding →
transformer_blocks → layer_norm → lm_head → logits
EXAMPLE:
>>> model = GPT(vocab_size=1000, embed_dim=256, num_layers=6, num_heads=8)
>>> tokens = Tensor(np.random.randint(0, 1000, (2, 10))) # (batch, seq)
>>> logits = model.forward(tokens)
>>> assert logits.shape == (2, 10, 1000) # (batch, seq, vocab)
HINTS:
- Positional embeddings are learned, not fixed sinusoidal
- Final layer norm stabilizes training
- Language modeling head shares weights with token embedding (tie_weights)
"""
### BEGIN SOLUTION
self.vocab_size = vocab_size
self.embed_dim = embed_dim
self.num_layers = num_layers
self.num_heads = num_heads
self.max_seq_len = max_seq_len
# Token and positional embeddings
self.token_embedding = Embedding(vocab_size, embed_dim)
self.position_embedding = Embedding(max_seq_len, embed_dim)
# Stack of transformer blocks
self.blocks = []
for _ in range(num_layers):
block = TransformerBlock(embed_dim, num_heads)
self.blocks.append(block)
# Final layer normalization
self.ln_f = LayerNorm(embed_dim)
# Language modeling head (projects to vocabulary)
self.lm_head = Linear(embed_dim, vocab_size, bias=False)
### END SOLUTION
def forward(self, tokens):
"""
Forward pass through GPT model.
TODO: Implement the complete GPT forward pass
APPROACH:
1. Get token embeddings and positional embeddings
2. Add them together (broadcasting handles different shapes)
3. Pass through all transformer blocks sequentially
4. Apply final layer norm and language modeling head
COMPUTATION FLOW:
tokens → embed + pos_embed → blocks → ln_f → lm_head → logits
CAUSAL MASKING:
For autoregressive generation, we need to prevent tokens from
seeing future tokens. This is handled by the attention mask.
HINT: Create position indices as range(seq_len) for positional embedding
"""
### BEGIN SOLUTION
batch_size, seq_len = tokens.shape
# Token embeddings
token_emb = self.token_embedding.forward(tokens)
# Positional embeddings
positions = Tensor(np.arange(seq_len).reshape(1, seq_len))
pos_emb = self.position_embedding.forward(positions)
# Combine embeddings
x = token_emb + pos_emb
# Create causal mask for autoregressive generation
mask = self._create_causal_mask(seq_len)
# Pass through transformer blocks
for block in self.blocks:
x = block.forward(x, mask)
# Final layer normalization
x = self.ln_f.forward(x)
# Language modeling head
logits = self.lm_head.forward(x)
return logits
### END SOLUTION
def _create_causal_mask(self, seq_len):
"""Create causal mask to prevent attending to future positions."""
### BEGIN SOLUTION
# Upper triangular matrix filled with -inf
mask = np.triu(np.ones((seq_len, seq_len)) * -np.inf, k=1)
return Tensor(mask)
### END SOLUTION
def generate(self, prompt_tokens, max_new_tokens=50, temperature=1.0):
"""
Generate text autoregressively.
TODO: Implement autoregressive text generation
APPROACH:
1. Start with prompt tokens
2. For each new position:
- Run forward pass to get logits
- Sample next token from logits
- Append to sequence
3. Return generated sequence
AUTOREGRESSIVE GENERATION:
At each step, the model predicts the next token based on all
previous tokens. This is how GPT generates coherent text.
EXAMPLE:
>>> model = GPT(vocab_size=100, embed_dim=64, num_layers=2, num_heads=4)
>>> prompt = Tensor([[1, 2, 3]]) # Some token sequence
>>> generated = model.generate(prompt, max_new_tokens=5)
>>> assert generated.shape[1] == 3 + 5 # original + new tokens
HINT: Use np.random.choice with temperature for sampling
"""
### BEGIN SOLUTION
current_tokens = Tensor(prompt_tokens.data.copy())
for _ in range(max_new_tokens):
# Get logits for current sequence
logits = self.forward(current_tokens)
# Get logits for last position (next token prediction)
last_logits = logits.data[:, -1, :] # (batch_size, vocab_size)
# Apply temperature scaling
scaled_logits = last_logits / temperature
# Convert to probabilities (softmax)
exp_logits = np.exp(scaled_logits - np.max(scaled_logits, axis=-1, keepdims=True))
probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
# Sample next token
next_token = np.array([[np.random.choice(self.vocab_size, p=probs[0])]])
# Append to sequence
current_tokens = Tensor(np.concatenate([current_tokens.data, next_token], axis=1))
return current_tokens
### END SOLUTION
def parameters(self):
"""Return all learnable parameters."""
params = []
params.extend(self.token_embedding.parameters())
params.extend(self.position_embedding.parameters())
for block in self.blocks:
params.extend(block.parameters())
params.extend(self.ln_f.parameters())
params.extend(self.lm_head.parameters())
return paramsIn [ ]:
def test_unit_gpt():
"""🔬 Test GPT model implementation."""
print("🔬 Unit Test: GPT Model...")
# Test small GPT model
vocab_size = 100
embed_dim = 64
num_layers = 2
num_heads = 4
model = GPT(vocab_size, embed_dim, num_layers, num_heads)
# Test forward pass
batch_size, seq_len = 2, 8
tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))
logits = model.forward(tokens)
# Check output shape
expected_shape = (batch_size, seq_len, vocab_size)
assert logits.shape == expected_shape
# Test generation
prompt = Tensor(np.random.randint(0, vocab_size, (1, 5)))
generated = model.generate(prompt, max_new_tokens=3)
# Check generation shape
assert generated.shape == (1, 8) # 5 prompt + 3 new tokens
# Test parameter counting
params = model.parameters()
assert len(params) > 10 # Should have many parameters from all components
# Test different model sizes
larger_model = GPT(vocab_size=200, embed_dim=128, num_layers=4, num_heads=8)
test_tokens = Tensor(np.random.randint(0, 200, (1, 10)))
larger_logits = larger_model.forward(test_tokens)
assert larger_logits.shape == (1, 10, 200)
print("✅ GPT model works correctly!")
# Run test immediately when developing this module
if __name__ == "__main__":
test_unit_gpt()In [ ]:
def demonstrate_transformer_integration():
"""
Demonstrate complete transformer pipeline.
This simulates training a small language model on a simple vocabulary.
"""
print("🔗 Integration Demo: Complete Language Model Pipeline")
print("Building a mini-GPT for character-level text generation")
# Create a small vocabulary (character-level)
vocab = list("abcdefghijklmnopqrstuvwxyz .")
vocab_size = len(vocab)
char_to_idx = {char: i for i, char in enumerate(vocab)}
idx_to_char = {i: char for i, char in enumerate(vocab)}
print(f"Vocabulary size: {vocab_size}")
print(f"Characters: {''.join(vocab)}")
# Create model
model = GPT(
vocab_size=vocab_size,
embed_dim=64,
num_layers=2,
num_heads=4,
max_seq_len=32
)
# Sample text encoding
text = "hello world."
tokens = [char_to_idx[char] for char in text]
input_tokens = Tensor(np.array([tokens]))
print(f"\nOriginal text: '{text}'")
print(f"Tokenized: {tokens}")
print(f"Input shape: {input_tokens.shape}")
# Forward pass
logits = model.forward(input_tokens)
print(f"Output logits shape: {logits.shape}")
print(f"Each position predicts next token from {vocab_size} possibilities")
# Generation demo
prompt_text = "hello"
prompt_tokens = [char_to_idx[char] for char in prompt_text]
prompt = Tensor(np.array([prompt_tokens]))
print(f"\nGeneration demo:")
print(f"Prompt: '{prompt_text}'")
generated = model.generate(prompt, max_new_tokens=8, temperature=1.0)
generated_text = ''.join([idx_to_char[idx] for idx in generated.data[0]])
print(f"Generated: '{generated_text}'")
print("(Note: Untrained model produces random text)")
return model
demonstrate_transformer_integration()In [ ]:
def analyze_parameter_scaling():
"""📊 Analyze how parameter count scales with model dimensions."""
print("📊 Analyzing Parameter Scaling in Transformers...")
print("Understanding why model size affects performance and cost\n")
# Test different model sizes
configs = [
{"name": "Tiny", "embed_dim": 64, "num_layers": 2, "num_heads": 4},
{"name": "Small", "embed_dim": 128, "num_layers": 4, "num_heads": 8},
{"name": "Medium", "embed_dim": 256, "num_layers": 8, "num_heads": 16},
{"name": "Large", "embed_dim": 512, "num_layers": 12, "num_heads": 16},
]
vocab_size = 50000 # Typical vocabulary size
for config in configs:
model = GPT(
vocab_size=vocab_size,
embed_dim=config["embed_dim"],
num_layers=config["num_layers"],
num_heads=config["num_heads"]
)
# Count parameters
total_params = 0
for param in model.parameters():
total_params += param.size
# Calculate memory requirements (4 bytes per float32 parameter)
memory_mb = (total_params * 4) / (1024 * 1024)
print(f"{config['name']} Model:")
print(f" Parameters: {total_params:,}")
print(f" Memory: {memory_mb:.1f} MB")
print(f" Embed dim: {config['embed_dim']}, Layers: {config['num_layers']}")
print()
print("💡 Parameter scaling is roughly quadratic with embedding dimension")
print("🚀 Real GPT-3 has 175B parameters, requiring ~350GB memory!")
analyze_parameter_scaling()In [ ]:
def analyze_attention_memory():
"""📊 Analyze attention memory complexity with sequence length."""
print("📊 Analyzing Attention Memory Complexity...")
print("Why long context is expensive and how it scales\n")
embed_dim = 512
num_heads = 8
batch_size = 4
# Test different sequence lengths
sequence_lengths = [128, 256, 512, 1024, 2048]
print("Attention Matrix Memory Usage:")
print("Seq Len | Attention Matrix Size | Memory (MB)")
print("-" * 45)
for seq_len in sequence_lengths:
# Attention matrix is (batch_size, num_heads, seq_len, seq_len)
attention_elements = batch_size * num_heads * seq_len * seq_len
# 4 bytes per float32
memory_bytes = attention_elements * 4
memory_mb = memory_bytes / (1024 * 1024)
print(f"{seq_len:6d} | {seq_len}×{seq_len} × {batch_size}×{num_heads} | {memory_mb:8.1f}")
print()
print("💡 Attention memory grows quadratically with sequence length")
print("🚀 This is why techniques like FlashAttention are crucial for long sequences")
analyze_attention_memory()In [ ]:
def test_module():
"""
Comprehensive test of entire module functionality.
This final test runs before module summary to ensure:
- All unit tests pass
- Functions work together correctly
- Module is ready for integration with TinyTorch
"""
print("🧪 RUNNING MODULE INTEGRATION TEST")
print("=" * 50)
# Run all unit tests
print("Running unit tests...")
test_unit_layer_norm()
test_unit_mlp()
test_unit_transformer_block()
test_unit_gpt()
print("\nRunning integration scenarios...")
# Test complete transformer training scenario
print("🔬 Integration Test: Full Training Pipeline...")
# Create model and data
vocab_size = 50
embed_dim = 64
num_layers = 2
num_heads = 4
model = GPT(vocab_size, embed_dim, num_layers, num_heads)
# Test batch processing
batch_size = 3
seq_len = 16
tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))
# Forward pass
logits = model.forward(tokens)
assert logits.shape == (batch_size, seq_len, vocab_size)
# Test generation with different temperatures
prompt = Tensor(np.random.randint(0, vocab_size, (1, 8)))
# Conservative generation
conservative = model.generate(prompt, max_new_tokens=5, temperature=0.1)
assert conservative.shape == (1, 13)
# Creative generation
creative = model.generate(prompt, max_new_tokens=5, temperature=2.0)
assert creative.shape == (1, 13)
# Test parameter counting consistency
total_params = sum(param.size for param in model.parameters())
assert total_params > 1000 # Should have substantial parameters
print("✅ Full transformer pipeline works!")
print("\n" + "=" * 50)
print("🎉 ALL TESTS PASSED! Module ready for export.")
print("Run: tito module complete 13")
# Call the comprehensive test
test_module()In [ ]:
if __name__ == "__main__":
print("🚀 Running Transformers module...")
test_module()
print("✅ Module validation complete!")