mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-31 10:16:50 -05:00
Vijay Janapa Reddi
4aec4ba297
- Removed 01_setup module (archived to archive/setup_module) - Renumbered all modules: tensor is now 01, activations is 02, etc. - Added tito setup command for environment setup and package installation - Added numeric shortcuts: tito 01, tito 02, etc. for quick module access - Fixed view command to find dev files correctly - Updated module dependencies and references - Improved user experience: immediate ML learning instead of boring setup
124 KiB
124 KiB
In [ ]:
#| default_exp core.transformers
#| export
import math
import numpy as np
import os
import sys
from typing import Union, List, Optional, Tuple, Dict
# Import our Tensor class - try from package first, then from local module
try:
from tinytorch.core.tensor import Tensor
except ImportError:
# For development, import from local tensor module
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '02_tensor'))
from tensor_dev import Tensor
# Try to import attention classes
try:
from tinytorch.core.attention import ScaledDotProductAttention, MultiHeadAttention, KVCache
except ImportError:
# For development, import from local module
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '13_attention'))
try:
from attention_dev import ScaledDotProductAttention, MultiHeadAttention, KVCache
except ImportError:
# Create minimal mock classes if not available
class MultiHeadAttention:
def __init__(self, embed_dim, num_heads):
self.embed_dim = embed_dim
self.num_heads = num_heads
def forward(self, q, k, v, mask=None):
return q # Mock implementation
class ScaledDotProductAttention:
def __init__(self):
pass
class KVCache:
def __init__(self, *args, **kwargs):
pass
# Try to import embedding classes
try:
from tinytorch.core.embeddings import Embedding, PositionalEncoding
except ImportError:
# For development, import from local module
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '12_embeddings'))
try:
from embeddings_dev import Embedding, PositionalEncoding
except ImportError:
# Create minimal mock classes if not available
class Embedding:
def __init__(self, vocab_size, embedding_dim):
self.vocab_size = vocab_size
self.embedding_dim = embedding_dim
class PositionalEncoding:
def __init__(self, embedding_dim, max_seq_length=5000):
self.embedding_dim = embedding_dimIn [ ]:
print("🏗️ TinyTorch Transformers Module")
print(f"NumPy version: {np.__version__}")
print("Ready to build complete transformer architectures!")In [ ]:
#| export
class LayerNorm:
"""
Layer Normalization for transformers.
Normalizes across the feature dimension (last axis) for each sample,
making training more stable and enabling deeper networks.
"""
def __init__(self, normalized_shape: Union[int, Tuple[int]], eps: float = 1e-5):
"""
Initialize layer normalization with learnable parameters.
TODO: Implement layer normalization initialization.
STEP-BY-STEP IMPLEMENTATION:
1. Store normalization configuration
2. Initialize learnable scale (gamma) and shift (beta) parameters
3. Set epsilon for numerical stability
4. Set up parameter tracking for optimization
MATHEMATICAL FOUNDATION:
LayerNorm(x) = γ * (x - μ) / σ + β
Where:
- μ = mean across feature dimensions
- σ = std across feature dimensions
- γ = learnable scale parameter
- β = learnable shift parameter
Args:
normalized_shape: Shape of features to normalize (e.g., embedding_dim)
eps: Small value for numerical stability
"""
### BEGIN SOLUTION
if isinstance(normalized_shape, int):
self.normalized_shape = (normalized_shape,)
else:
self.normalized_shape = normalized_shape
self.eps = eps
# Initialize learnable parameters
# Gamma (scale): initialized to ones
# Beta (bias): initialized to zeros
self.gamma = Tensor(np.ones(self.normalized_shape))
self.beta = Tensor(np.zeros(self.normalized_shape))
# Track parameters for optimization
self.parameters = [self.gamma, self.beta]
### END SOLUTION
def forward(self, x: Tensor) -> Tensor:
"""
Apply layer normalization to input tensor.
TODO: Implement layer normalization forward pass.
STEP-BY-STEP IMPLEMENTATION:
1. Calculate mean across feature dimensions
2. Calculate standard deviation across feature dimensions
3. Normalize: (x - mean) / (std + eps)
4. Apply learnable scale and shift: gamma * normalized + beta
NUMERICAL STABILITY:
- Add eps to variance before taking sqrt
- Use unbiased variance calculation
EXAMPLE:
layer_norm = LayerNorm(256)
x = Tensor(np.random.randn(32, 128, 256)) # (batch, seq, features)
normalized = layer_norm.forward(x) # Same shape as input
Args:
x: Input tensor with shape (..., *normalized_shape)
Returns:
Normalized tensor with same shape as input
"""
### BEGIN SOLUTION
# Calculate mean and variance across the feature dimensions (last axes)
# For shape (..., *normalized_shape), we want to normalize over the last len(normalized_shape) axes
# Determine axes to normalize over
axes_to_normalize = tuple(range(len(x.shape) - len(self.normalized_shape), len(x.shape)))
# Calculate mean
mean = np.mean(x.data, axis=axes_to_normalize, keepdims=True)
# Calculate variance
variance = np.var(x.data, axis=axes_to_normalize, keepdims=True)
# Normalize
normalized = (x.data - mean) / np.sqrt(variance + self.eps)
# Apply learnable scale and shift
# Reshape gamma and beta to be broadcastable
gamma_broadcasted = self.gamma.data.reshape([1] * (len(x.shape) - len(self.normalized_shape)) + list(self.normalized_shape))
beta_broadcasted = self.beta.data.reshape([1] * (len(x.shape) - len(self.normalized_shape)) + list(self.normalized_shape))
output = gamma_broadcasted * normalized + beta_broadcasted
return Tensor(output)
### END SOLUTION
def __call__(self, x: Tensor) -> Tensor:
"""Make the class callable."""
return self.forward(x)
def get_memory_usage(self) -> Dict[str, float]:
"""
Calculate memory usage of layer normalization parameters.
This function is PROVIDED to show memory analysis.
"""
# Parameter memory
param_memory_mb = sum(param.data.nbytes for param in self.parameters) / (1024 * 1024)
return {
'parameter_memory_mb': param_memory_mb,
'total_parameters': sum(param.data.size for param in self.parameters),
'normalized_shape': self.normalized_shape
}In [ ]:
def test_unit_layer_norm():
"""Unit test for layer normalization."""
print("🔬 Unit Test: Layer Normalization...")
# Test 1: Basic functionality
embed_dim = 256
layer_norm = LayerNorm(embed_dim)
# Verify initialization
assert layer_norm.normalized_shape == (embed_dim,), "Should store normalized shape"
assert len(layer_norm.parameters) == 2, "Should have gamma and beta parameters"
assert layer_norm.gamma.shape == (embed_dim,), "Gamma should match normalized shape"
assert layer_norm.beta.shape == (embed_dim,), "Beta should match normalized shape"
# Verify parameter initialization
assert np.allclose(layer_norm.gamma.data, 1.0), "Gamma should be initialized to ones"
assert np.allclose(layer_norm.beta.data, 0.0), "Beta should be initialized to zeros"
# Test 2: Forward pass with 2D input
batch_size = 16
x_2d = Tensor(np.random.randn(batch_size, embed_dim))
output_2d = layer_norm.forward(x_2d)
assert output_2d.shape == x_2d.shape, "Output shape should match input shape"
# Test 3: Forward pass with 3D input (typical transformer use)
seq_length = 32
x_3d = Tensor(np.random.randn(batch_size, seq_length, embed_dim))
output_3d = layer_norm.forward(x_3d)
assert output_3d.shape == x_3d.shape, "3D output shape should match input shape"
# Test 4: Normalization properties
# For each sample, the normalized features should have ~zero mean and ~unit variance
for i in range(batch_size):
for j in range(seq_length):
sample_output = output_3d.data[i, j, :]
sample_mean = np.mean(sample_output)
sample_var = np.var(sample_output)
assert abs(sample_mean) < 1e-6, f"Normalized mean should be ~0, got {sample_mean}"
assert abs(sample_var - 1.0) < 1e-6, f"Normalized variance should be ~1, got {sample_var}"
# Test 5: Different normalized shapes
multi_dim_shape = (64, 4) # Multi-dimensional normalization
layer_norm_multi = LayerNorm(multi_dim_shape)
x_multi = Tensor(np.random.randn(8, 32, 64, 4))
output_multi = layer_norm_multi.forward(x_multi)
assert output_multi.shape == x_multi.shape, "Multi-dim normalization should preserve shape"
# Test 6: Callable interface
output_callable = layer_norm(x_3d)
assert np.allclose(output_callable.data, output_3d.data), "Callable interface should work"
# Test 7: Numerical stability with extreme values
extreme_x = Tensor(np.ones((4, embed_dim)) * 1e6) # Very large values
extreme_output = layer_norm.forward(extreme_x)
assert not np.any(np.isnan(extreme_output.data)), "Should handle extreme values without NaN"
assert not np.any(np.isinf(extreme_output.data)), "Should handle extreme values without inf"
# Test 8: Memory usage calculation
memory_stats = layer_norm.get_memory_usage()
assert 'parameter_memory_mb' in memory_stats, "Should provide memory statistics"
assert memory_stats['total_parameters'] == 2 * embed_dim, "Should count gamma and beta parameters"
print("✅ Layer normalization tests passed!")
print(f"✅ Properly normalizes across feature dimensions")
print(f"✅ Handles 2D and 3D inputs correctly")
print(f"✅ Maintains ~0 mean and ~1 variance after normalization")
print(f"✅ Parameter memory: {memory_stats['parameter_memory_mb']:.4f}MB")
# Test function defined (called in main block)In [ ]:
#| export
class PositionwiseFeedForward:
"""
Position-wise feed-forward network used in transformer blocks.
Applies the same feed-forward network to each position in the sequence:
FFN(x) = max(0, xW₁ + b₁)W₂ + b₂
"""
def __init__(self, embed_dim: int, hidden_dim: int, dropout: float = 0.0):
"""
Initialize position-wise feed-forward network.
TODO: Implement feed-forward network initialization.
STEP-BY-STEP IMPLEMENTATION:
1. Store network configuration
2. Initialize weight matrices and bias vectors for two linear layers
3. Set up parameter tracking for optimization
4. Store dropout rate for training
ARCHITECTURE:
- Input: (batch, seq_len, embed_dim)
- Linear 1: embed_dim → hidden_dim
- ReLU activation
- Linear 2: hidden_dim → embed_dim
- Output: (batch, seq_len, embed_dim)
PARAMETER INITIALIZATION:
Use Xavier/Glorot initialization for stable training
Args:
embed_dim: Embedding dimension (input and output size)
hidden_dim: Hidden layer dimension (typically 4 * embed_dim)
dropout: Dropout rate for regularization
"""
### BEGIN SOLUTION
self.embed_dim = embed_dim
self.hidden_dim = hidden_dim
self.dropout = dropout
# Initialize weights using Xavier initialization
# W1: embed_dim → hidden_dim
xavier_bound_1 = math.sqrt(6.0 / (embed_dim + hidden_dim))
self.w1 = Tensor(np.random.uniform(-xavier_bound_1, xavier_bound_1, (embed_dim, hidden_dim)))
self.b1 = Tensor(np.zeros(hidden_dim))
# W2: hidden_dim → embed_dim
xavier_bound_2 = math.sqrt(6.0 / (hidden_dim + embed_dim))
self.w2 = Tensor(np.random.uniform(-xavier_bound_2, xavier_bound_2, (hidden_dim, embed_dim)))
self.b2 = Tensor(np.zeros(embed_dim))
# Track parameters for optimization
self.parameters = [self.w1, self.b1, self.w2, self.b2]
### END SOLUTION
def forward(self, x: Tensor) -> Tensor:
"""
Apply position-wise feed-forward transformation.
TODO: Implement feed-forward forward pass.
STEP-BY-STEP IMPLEMENTATION:
1. Apply first linear transformation: x @ W1 + b1
2. Apply ReLU activation: max(0, linear1)
3. Apply second linear transformation: relu @ W2 + b2
4. Return result with same shape as input
MATHEMATICAL FORMULATION:
hidden = ReLU(x @ W1 + b1)
output = hidden @ W2 + b2
Args:
x: Input tensor with shape (batch_size, seq_len, embed_dim)
Returns:
Output tensor with shape (batch_size, seq_len, embed_dim)
"""
### BEGIN SOLUTION
# Reshape input for matrix multiplication if needed
original_shape = x.shape
if len(x.shape) == 3:
batch_size, seq_len, embed_dim = x.shape
# Reshape to (batch_size * seq_len, embed_dim) for efficient computation
x_reshaped = x.data.reshape(-1, embed_dim)
else:
x_reshaped = x.data
# First linear transformation: x @ W1 + b1
hidden = np.matmul(x_reshaped, self.w1.data) + self.b1.data
# ReLU activation
hidden_relu = np.maximum(0, hidden)
# Second linear transformation: hidden @ W2 + b2
output = np.matmul(hidden_relu, self.w2.data) + self.b2.data
# Reshape back to original shape
if len(original_shape) == 3:
output = output.reshape(original_shape)
return Tensor(output)
### END SOLUTION
def __call__(self, x: Tensor) -> Tensor:
"""Make the class callable."""
return self.forward(x)
def get_memory_usage(self) -> Dict[str, float]:
"""
Calculate memory usage of feed-forward parameters.
This function is PROVIDED to show memory analysis.
"""
# Parameter memory
param_memory_mb = sum(param.data.nbytes for param in self.parameters) / (1024 * 1024)
# Calculate parameter counts
w1_params = self.embed_dim * self.hidden_dim
w2_params = self.hidden_dim * self.embed_dim
bias_params = self.hidden_dim + self.embed_dim
total_params = w1_params + w2_params + bias_params
return {
'parameter_memory_mb': param_memory_mb,
'total_parameters': total_params,
'w1_parameters': w1_params,
'w2_parameters': w2_params,
'bias_parameters': bias_params,
'embed_dim': self.embed_dim,
'hidden_dim': self.hidden_dim
}In [ ]:
def test_unit_feed_forward():
"""Unit test for position-wise feed-forward network."""
print("🔬 Unit Test: Position-wise Feed-Forward Network...")
# Test configuration
embed_dim = 256
hidden_dim = 1024 # Typical 4x expansion
ffn = PositionwiseFeedForward(embed_dim=embed_dim, hidden_dim=hidden_dim)
# Verify initialization
assert ffn.embed_dim == embed_dim, "Should store embedding dimension"
assert ffn.hidden_dim == hidden_dim, "Should store hidden dimension"
assert len(ffn.parameters) == 4, "Should have W1, b1, W2, b2 parameters"
# Verify parameter shapes
assert ffn.w1.shape == (embed_dim, hidden_dim), f"W1 should be ({embed_dim}, {hidden_dim})"
assert ffn.b1.shape == (hidden_dim,), f"b1 should be ({hidden_dim},)"
assert ffn.w2.shape == (hidden_dim, embed_dim), f"W2 should be ({hidden_dim}, {embed_dim})"
assert ffn.b2.shape == (embed_dim,), f"b2 should be ({embed_dim},)"
# Test forward pass with 3D input (typical transformer use)
batch_size = 8
seq_len = 32
x_3d = Tensor(np.random.randn(batch_size, seq_len, embed_dim))
output_3d = ffn.forward(x_3d)
expected_shape = (batch_size, seq_len, embed_dim)
assert output_3d.shape == expected_shape, f"Expected shape {expected_shape}, got {output_3d.shape}"
# Test forward pass with 2D input
x_2d = Tensor(np.random.randn(batch_size, embed_dim))
output_2d = ffn.forward(x_2d)
expected_2d_shape = (batch_size, embed_dim)
assert output_2d.shape == expected_2d_shape, f"Expected 2D shape {expected_2d_shape}, got {output_2d.shape}"
# Test that FFN is applied position-wise (same transformation at each position)
# Extract two positions from the sequence
pos_1_input = Tensor(x_3d.data[:, 0, :]) # First position
pos_2_input = Tensor(x_3d.data[:, 1, :]) # Second position
pos_1_output = ffn.forward(pos_1_input)
pos_2_output = ffn.forward(pos_2_input)
# Compare with full sequence output
assert np.allclose(pos_1_output.data, output_3d.data[:, 0, :]), "Position 0 should match individual processing"
assert np.allclose(pos_2_output.data, output_3d.data[:, 1, :]), "Position 1 should match individual processing"
# Test ReLU activation (some outputs should be zero for negative intermediate values)
# Create input that will definitely produce some negative values after first linear layer
negative_input = Tensor(-np.ones((4, embed_dim)) * 10) # Very negative input
negative_output = ffn.forward(negative_input)
# Not all outputs should be negative (ReLU should clip some values)
assert not np.all(negative_output.data < 0), "ReLU should prevent all outputs from being negative"
# Test callable interface
output_callable = ffn(x_3d)
assert np.allclose(output_callable.data, output_3d.data), "Callable interface should work"
# Test different hidden dimensions
for test_hidden_dim in [512, 2048]:
test_ffn = PositionwiseFeedForward(embed_dim=embed_dim, hidden_dim=test_hidden_dim)
test_output = test_ffn.forward(x_3d)
assert test_output.shape == expected_shape, f"Should work with hidden_dim={test_hidden_dim}"
# Test memory usage calculation
memory_stats = ffn.get_memory_usage()
assert 'parameter_memory_mb' in memory_stats, "Should provide memory statistics"
# Verify parameter counts
expected_w1_params = embed_dim * hidden_dim
expected_w2_params = hidden_dim * embed_dim
expected_total = expected_w1_params + expected_w2_params + hidden_dim + embed_dim
assert memory_stats['w1_parameters'] == expected_w1_params, "Should count W1 parameters correctly"
assert memory_stats['w2_parameters'] == expected_w2_params, "Should count W2 parameters correctly"
assert memory_stats['total_parameters'] == expected_total, "Should count total parameters correctly"
print("✅ Position-wise feed-forward tests passed!")
print(f"✅ Handles 2D and 3D inputs correctly")
print(f"✅ Position-wise processing verified")
print(f"✅ ReLU activation working properly")
print(f"✅ Total parameters: {memory_stats['total_parameters']:,}")
print(f"✅ Parameter memory: {memory_stats['parameter_memory_mb']:.2f}MB")
# Test function defined (called in main block)In [ ]:
#| export
class TransformerBlock:
"""
Complete transformer block with self-attention and feed-forward layers.
Combines multi-head self-attention, layer normalization, residual connections,
and position-wise feed-forward networks into the standard transformer architecture.
"""
def __init__(self, embed_dim: int, num_heads: int, hidden_dim: int,
dropout: float = 0.0, pre_norm: bool = True):
"""
Initialize transformer block with all components.
TODO: Implement transformer block initialization.
STEP-BY-STEP IMPLEMENTATION:
1. Store block configuration
2. Create multi-head attention layer
3. Create two layer normalization layers (for attention and FFN)
4. Create position-wise feed-forward network
5. Set up parameter tracking from all sub-components
ARCHITECTURE CHOICE: Pre-norm vs Post-norm
- Pre-norm: LayerNorm → Attention → Residual (more stable)
- Post-norm: Attention → LayerNorm → Residual (original paper)
Args:
embed_dim: Embedding dimension
num_heads: Number of attention heads
hidden_dim: Feed-forward hidden dimension (typically 4 * embed_dim)
dropout: Dropout rate for regularization
pre_norm: Whether to use pre-normalization (recommended)
"""
### BEGIN SOLUTION
self.embed_dim = embed_dim
self.num_heads = num_heads
self.hidden_dim = hidden_dim
self.dropout = dropout
self.pre_norm = pre_norm
# Multi-head self-attention
self.attention = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout)
# Layer normalization layers
self.norm1 = LayerNorm(embed_dim) # For attention
self.norm2 = LayerNorm(embed_dim) # For feed-forward
# Position-wise feed-forward network
self.ffn = PositionwiseFeedForward(embed_dim=embed_dim, hidden_dim=hidden_dim, dropout=dropout)
# Collect all parameters from sub-components
self.parameters = []
if hasattr(self.attention, 'parameters'):
self.parameters.extend(self.attention.parameters)
self.parameters.extend(self.norm1.parameters)
self.parameters.extend(self.norm2.parameters)
self.parameters.extend(self.ffn.parameters)
### END SOLUTION
def forward(self, x: Tensor, mask: Optional[Tensor] = None,
return_attention_weights: bool = False) -> Union[Tensor, Tuple[Tensor, Tensor]]:
"""
Process input through complete transformer block.
TODO: Implement transformer block forward pass.
STEP-BY-STEP IMPLEMENTATION (Pre-norm):
1. Self-attention with residual: x + attention(norm1(x))
2. Feed-forward with residual: attn_out + ffn(norm2(attn_out))
3. Return final output (and optionally attention weights)
RESIDUAL CONNECTIONS:
Essential for training deep networks - allow gradients to flow directly
Args:
x: Input tensor with shape (batch_size, seq_len, embed_dim)
mask: Optional attention mask
return_attention_weights: Whether to return attention weights
Returns:
Transformer block output with same shape as input
Optionally also attention weights
"""
### BEGIN SOLUTION
if self.pre_norm:
# Pre-normalization: LayerNorm before attention/FFN
# Self-attention with residual connection
norm1_x = self.norm1(x)
if return_attention_weights:
attn_output, attn_weights = self.attention.forward(
norm1_x, norm1_x, norm1_x, mask=mask, return_attention_weights=True
)
else:
attn_output = self.attention.forward(norm1_x, norm1_x, norm1_x, mask=mask)
# Residual connection
x = Tensor(x.data + attn_output.data)
# Feed-forward with residual connection
norm2_x = self.norm2(x)
ffn_output = self.ffn.forward(norm2_x)
# Residual connection
output = Tensor(x.data + ffn_output.data)
else:
# Post-normalization: LayerNorm after attention/FFN (original transformer)
# Self-attention with residual connection
if return_attention_weights:
attn_output, attn_weights = self.attention.forward(
x, x, x, mask=mask, return_attention_weights=True
)
else:
attn_output = self.attention.forward(x, x, x, mask=mask)
# Residual + LayerNorm
attn_residual = Tensor(x.data + attn_output.data)
norm1_output = self.norm1(attn_residual)
# Feed-forward with residual connection
ffn_output = self.ffn.forward(norm1_output)
# Residual + LayerNorm
ffn_residual = Tensor(norm1_output.data + ffn_output.data)
output = self.norm2(ffn_residual)
if return_attention_weights:
return output, attn_weights
else:
return output
### END SOLUTION
def __call__(self, x: Tensor, mask: Optional[Tensor] = None,
return_attention_weights: bool = False) -> Union[Tensor, Tuple[Tensor, Tensor]]:
"""Make the class callable."""
return self.forward(x, mask, return_attention_weights)
def get_memory_usage(self) -> Dict[str, float]:
"""
Calculate memory usage of transformer block components.
This function is PROVIDED to show memory analysis.
"""
# Get memory usage from components
if hasattr(self.attention, 'get_memory_usage'):
attention_memory = self.attention.get_memory_usage()['total_parameter_memory_mb']
else:
attention_memory = 0.0
norm1_memory = self.norm1.get_memory_usage()['parameter_memory_mb']
norm2_memory = self.norm2.get_memory_usage()['parameter_memory_mb']
ffn_memory = self.ffn.get_memory_usage()['parameter_memory_mb']
total_memory = attention_memory + norm1_memory + norm2_memory + ffn_memory
total_params = len(self.parameters) if hasattr(self, 'parameters') else 0
return {
'total_memory_mb': total_memory,
'attention_memory_mb': attention_memory,
'norm_memory_mb': norm1_memory + norm2_memory,
'ffn_memory_mb': ffn_memory,
'total_parameters': sum(p.data.size for p in self.parameters) if hasattr(self, 'parameters') else 0,
'embed_dim': self.embed_dim,
'num_heads': self.num_heads,
'hidden_dim': self.hidden_dim,
'pre_norm': self.pre_norm
}In [ ]:
def test_unit_transformer_block():
"""Unit test for transformer block."""
print("🔬 Unit Test: Transformer Block...")
# Test configuration
embed_dim = 256
num_heads = 8
hidden_dim = 1024
transformer_block = TransformerBlock(
embed_dim=embed_dim,
num_heads=num_heads,
hidden_dim=hidden_dim,
pre_norm=True
)
# Verify initialization
assert transformer_block.embed_dim == embed_dim, "Should store embedding dimension"
assert transformer_block.num_heads == num_heads, "Should store number of heads"
assert transformer_block.hidden_dim == hidden_dim, "Should store hidden dimension"
assert transformer_block.pre_norm == True, "Should store normalization type"
# Verify components exist
assert hasattr(transformer_block, 'attention'), "Should have attention layer"
assert hasattr(transformer_block, 'norm1'), "Should have first norm layer"
assert hasattr(transformer_block, 'norm2'), "Should have second norm layer"
assert hasattr(transformer_block, 'ffn'), "Should have feed-forward network"
# Test forward pass
batch_size = 4
seq_len = 16
x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))
output = transformer_block.forward(x)
expected_shape = (batch_size, seq_len, embed_dim)
assert output.shape == expected_shape, f"Expected shape {expected_shape}, got {output.shape}"
# Test with attention weights return
output_with_attn, attn_weights = transformer_block.forward(x, return_attention_weights=True)
assert output_with_attn.shape == expected_shape, "Output with attention should have correct shape"
expected_attn_shape = (batch_size, num_heads, seq_len, seq_len)
assert attn_weights.shape == expected_attn_shape, f"Expected attention shape {expected_attn_shape}, got {attn_weights.shape}"
# Test with causal mask
causal_mask = np.triu(np.ones((seq_len, seq_len)), k=1)
causal_mask = 1 - causal_mask # Convert to attention mask
masked_output, masked_attn = transformer_block.forward(
x, mask=Tensor(causal_mask), return_attention_weights=True
)
assert masked_output.shape == expected_shape, "Masked output should have correct shape"
# Verify causal masking works
for head in range(num_heads):
for i in range(seq_len):
for j in range(i+1, seq_len):
assert np.all(masked_attn.data[:, head, i, j] < 1e-5), \
f"Position ({i},{j}) should be masked in head {head}"
# Test residual connections by checking that output is different from pure attention
# If we zero out the input, residual connections should preserve some information
zero_input = Tensor(np.zeros((batch_size, seq_len, embed_dim)))
zero_output = transformer_block.forward(zero_input)
# Output should not be exactly zero due to biases and layer norm parameters
assert not np.allclose(zero_output.data, 0), "Residual connections should prevent zero output"
# Test post-normalization variant
post_norm_block = TransformerBlock(
embed_dim=embed_dim,
num_heads=num_heads,
hidden_dim=hidden_dim,
pre_norm=False
)
post_norm_output = post_norm_block.forward(x)
assert post_norm_output.shape == expected_shape, "Post-norm should produce correct shape"
# Pre-norm and post-norm should produce different outputs
pre_norm_output = transformer_block.forward(x)
assert not np.allclose(pre_norm_output.data, post_norm_output.data), \
"Pre-norm and post-norm should produce different outputs"
# Test callable interface
output_callable = transformer_block(x)
assert np.allclose(output_callable.data, output.data), "Callable interface should work"
# Test different configurations
for test_heads in [4, 16]:
if embed_dim % test_heads == 0:
test_block = TransformerBlock(embed_dim=embed_dim, num_heads=test_heads, hidden_dim=hidden_dim)
test_output = test_block.forward(x)
assert test_output.shape == expected_shape, f"Should work with {test_heads} heads"
# Test memory usage calculation
memory_stats = transformer_block.get_memory_usage()
assert 'total_memory_mb' in memory_stats, "Should provide memory statistics"
assert memory_stats['total_memory_mb'] > 0, "Should have positive memory usage"
assert memory_stats['total_parameters'] > 0, "Should count parameters"
print("✅ Transformer block tests passed!")
print(f"✅ Pre-norm and post-norm architectures work correctly")
print(f"✅ Residual connections preserve information flow")
print(f"✅ Causal masking works across all attention heads")
print(f"✅ Total parameters: {memory_stats['total_parameters']:,}")
print(f"✅ Total memory: {memory_stats['total_memory_mb']:.2f}MB")
# Test function defined (called in main block)In [ ]:
#| export
class Transformer:
"""
Complete transformer model for language processing.
Stacks multiple transformer blocks with token embeddings and positional
encoding to create a complete language model architecture.
"""
def __init__(self, vocab_size: int, embed_dim: int, num_heads: int,
num_layers: int, hidden_dim: int, max_seq_length: int = 1024,
dropout: float = 0.0, pre_norm: bool = True):
"""
Initialize complete transformer model.
TODO: Implement transformer model initialization.
STEP-BY-STEP IMPLEMENTATION:
1. Store model configuration
2. Create token embedding layer
3. Create positional encoding
4. Create stack of transformer blocks
5. Create output projection layer (for language modeling)
6. Set up parameter tracking from all components
LANGUAGE MODELING HEAD:
Final linear layer that projects hidden states to vocabulary logits
Args:
vocab_size: Size of vocabulary
embed_dim: Embedding dimension
num_heads: Number of attention heads per layer
num_layers: Number of transformer blocks
hidden_dim: Feed-forward hidden dimension
max_seq_length: Maximum sequence length for positional encoding
dropout: Dropout rate
pre_norm: Whether to use pre-normalization
"""
### BEGIN SOLUTION
self.vocab_size = vocab_size
self.embed_dim = embed_dim
self.num_heads = num_heads
self.num_layers = num_layers
self.hidden_dim = hidden_dim
self.max_seq_length = max_seq_length
self.dropout = dropout
self.pre_norm = pre_norm
# Token embedding layer
self.token_embedding = Embedding(vocab_size=vocab_size, embedding_dim=embed_dim)
# Positional encoding
self.pos_encoding = PositionalEncoding(embedding_dim=embed_dim, max_seq_length=max_seq_length)
# Stack of transformer blocks
self.transformer_blocks = []
for _ in range(num_layers):
block = TransformerBlock(
embed_dim=embed_dim,
num_heads=num_heads,
hidden_dim=hidden_dim,
dropout=dropout,
pre_norm=pre_norm
)
self.transformer_blocks.append(block)
# Final layer normalization (for pre-norm architecture)
if pre_norm:
self.final_norm = LayerNorm(embed_dim)
else:
self.final_norm = None
# Language modeling head (projects to vocabulary)
xavier_bound = math.sqrt(6.0 / (embed_dim + vocab_size))
self.lm_head = Tensor(np.random.uniform(-xavier_bound, xavier_bound, (embed_dim, vocab_size)))
# Collect all parameters
self.parameters = []
if hasattr(self.token_embedding, 'parameters'):
self.parameters.extend(self.token_embedding.parameters)
for block in self.transformer_blocks:
if hasattr(block, 'parameters'):
self.parameters.extend(block.parameters)
if self.final_norm:
self.parameters.extend(self.final_norm.parameters)
self.parameters.append(self.lm_head)
### END SOLUTION
def forward(self, input_ids: Tensor, mask: Optional[Tensor] = None,
return_attention_weights: bool = False) -> Union[Tensor, Tuple[Tensor, List[Tensor]]]:
"""
Process input through complete transformer model.
TODO: Implement transformer model forward pass.
STEP-BY-STEP IMPLEMENTATION:
1. Convert token IDs to embeddings
2. Add positional encoding
3. Process through all transformer blocks
4. Apply final normalization (if pre-norm)
5. Apply language modeling head
6. Return logits (and optionally attention weights)
Args:
input_ids: Token indices with shape (batch_size, seq_len)
mask: Optional attention mask
return_attention_weights: Whether to return all attention weights
Returns:
Logits with shape (batch_size, seq_len, vocab_size)
Optionally also list of attention weights from each layer
"""
### BEGIN SOLUTION
# Token embeddings
embeddings = self.token_embedding.forward(input_ids)
# Add positional encoding
x = self.pos_encoding.forward(embeddings)
# Process through transformer blocks
all_attention_weights = []
for block in self.transformer_blocks:
if return_attention_weights:
x, attn_weights = block.forward(x, mask=mask, return_attention_weights=True)
all_attention_weights.append(attn_weights)
else:
x = block.forward(x, mask=mask)
# Final layer normalization (for pre-norm)
if self.final_norm:
x = self.final_norm.forward(x)
# Language modeling head
# x: (batch_size, seq_len, embed_dim)
# lm_head: (embed_dim, vocab_size)
# output: (batch_size, seq_len, vocab_size)
batch_size, seq_len, embed_dim = x.shape
x_reshaped = x.data.reshape(-1, embed_dim) # (batch_size * seq_len, embed_dim)
logits_reshaped = np.matmul(x_reshaped, self.lm_head.data) # (batch_size * seq_len, vocab_size)
logits = logits_reshaped.reshape(batch_size, seq_len, self.vocab_size)
if return_attention_weights:
return Tensor(logits), all_attention_weights
else:
return Tensor(logits)
### END SOLUTION
def __call__(self, input_ids: Tensor, mask: Optional[Tensor] = None,
return_attention_weights: bool = False) -> Union[Tensor, Tuple[Tensor, List[Tensor]]]:
"""Make the class callable."""
return self.forward(input_ids, mask, return_attention_weights)
def generate(self, input_ids: Tensor, max_new_tokens: int = 50,
temperature: float = 1.0) -> Tensor:
"""
Generate text autoregressively.
This function is PROVIDED to show text generation capability.
"""
batch_size, current_seq_len = input_ids.shape
if current_seq_len >= self.max_seq_length:
raise ValueError(f"Input sequence length {current_seq_len} exceeds max {self.max_seq_length}")
generated_ids = input_ids.data.copy()
for _ in range(max_new_tokens):
# Create causal mask
seq_len = generated_ids.shape[1]
causal_mask = np.triu(np.ones((seq_len, seq_len)), k=1)
causal_mask = 1 - causal_mask
# Forward pass
logits = self.forward(Tensor(generated_ids), mask=Tensor(causal_mask))
# Get logits for last position
last_logits = logits.data[:, -1, :] # (batch_size, vocab_size)
# Apply temperature
last_logits = last_logits / temperature
# Sample next token (using simple sampling)
# Convert to probabilities
exp_logits = np.exp(last_logits - np.max(last_logits, axis=-1, keepdims=True))
probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
# Sample from distribution
next_tokens = []
for i in range(batch_size):
next_token = np.random.choice(self.vocab_size, p=probs[i])
next_tokens.append(next_token)
next_tokens = np.array(next_tokens).reshape(batch_size, 1)
# Append to sequence
generated_ids = np.concatenate([generated_ids, next_tokens], axis=1)
# Stop if we reach max sequence length
if generated_ids.shape[1] >= self.max_seq_length:
break
return Tensor(generated_ids)
def get_memory_usage(self) -> Dict[str, float]:
"""
Calculate memory usage of complete transformer model.
This function is PROVIDED to show memory analysis.
"""
# Token embedding memory
if hasattr(self.token_embedding, 'get_memory_usage'):
embedding_memory = self.token_embedding.get_memory_usage()['total_memory_mb']
else:
embedding_memory = self.vocab_size * self.embed_dim * 4 / (1024 * 1024)
# Transformer blocks memory
block_memory = 0
if self.transformer_blocks:
single_block_memory = self.transformer_blocks[0].get_memory_usage()['total_memory_mb']
block_memory = single_block_memory * self.num_layers
# Final norm memory
final_norm_memory = 0
if self.final_norm:
final_norm_memory = self.final_norm.get_memory_usage()['parameter_memory_mb']
# Language modeling head memory
lm_head_memory = self.lm_head.data.nbytes / (1024 * 1024)
total_memory = embedding_memory + block_memory + final_norm_memory + lm_head_memory
total_params = sum(p.data.size for p in self.parameters) if hasattr(self, 'parameters') else 0
return {
'total_memory_mb': total_memory,
'embedding_memory_mb': embedding_memory,
'transformer_blocks_memory_mb': block_memory,
'lm_head_memory_mb': lm_head_memory,
'total_parameters': total_params,
'vocab_size': self.vocab_size,
'embed_dim': self.embed_dim,
'num_layers': self.num_layers,
'num_heads': self.num_heads,
'hidden_dim': self.hidden_dim
}In [ ]:
def test_unit_transformer_model():
"""Unit test for complete transformer model."""
print("🔬 Unit Test: Complete Transformer Model...")
# Test configuration
vocab_size = 1000
embed_dim = 256
num_heads = 8
num_layers = 4
hidden_dim = 512
max_seq_length = 128
transformer = Transformer(
vocab_size=vocab_size,
embed_dim=embed_dim,
num_heads=num_heads,
num_layers=num_layers,
hidden_dim=hidden_dim,
max_seq_length=max_seq_length,
pre_norm=True
)
# Verify initialization
assert transformer.vocab_size == vocab_size, "Should store vocabulary size"
assert transformer.embed_dim == embed_dim, "Should store embedding dimension"
assert transformer.num_layers == num_layers, "Should store number of layers"
assert len(transformer.transformer_blocks) == num_layers, "Should create correct number of blocks"
# Verify components exist
assert hasattr(transformer, 'token_embedding'), "Should have token embedding"
assert hasattr(transformer, 'pos_encoding'), "Should have positional encoding"
assert hasattr(transformer, 'lm_head'), "Should have language modeling head"
# Test forward pass with token IDs
batch_size = 4
seq_len = 32
input_ids = np.random.randint(0, vocab_size, (batch_size, seq_len))
input_tensor = Tensor(input_ids)
logits = transformer.forward(input_tensor)
expected_shape = (batch_size, seq_len, vocab_size)
assert logits.shape == expected_shape, f"Expected shape {expected_shape}, got {logits.shape}"
# Test with attention weights return
logits_with_attn, all_attention_weights = transformer.forward(input_tensor, return_attention_weights=True)
assert logits_with_attn.shape == expected_shape, "Logits with attention should have correct shape"
assert len(all_attention_weights) == num_layers, f"Should return attention weights from {num_layers} layers"
for i, attn_weights in enumerate(all_attention_weights):
expected_attn_shape = (batch_size, num_heads, seq_len, seq_len)
assert attn_weights.shape == expected_attn_shape, \
f"Layer {i} attention should have shape {expected_attn_shape}, got {attn_weights.shape}"
# Test with causal mask
causal_mask = np.triu(np.ones((seq_len, seq_len)), k=1)
causal_mask = 1 - causal_mask # Convert to attention mask
masked_logits, masked_attention = transformer.forward(
input_tensor, mask=Tensor(causal_mask), return_attention_weights=True
)
assert masked_logits.shape == expected_shape, "Masked logits should have correct shape"
# Verify causal masking propagates through all layers
for layer_idx, attn_weights in enumerate(masked_attention):
for head in range(num_heads):
for i in range(seq_len):
for j in range(i+1, seq_len):
assert np.all(attn_weights.data[:, head, i, j] < 1e-5), \
f"Layer {layer_idx}, head {head}: position ({i},{j}) should be masked"
# Test callable interface
logits_callable = transformer(input_tensor)
assert np.allclose(logits_callable.data, logits.data), "Callable interface should work"
# Test text generation capability
print(" Testing text generation...")
start_tokens = Tensor(np.random.randint(0, vocab_size, (2, 8))) # 2 sequences, 8 tokens each
generated = transformer.generate(start_tokens, max_new_tokens=10, temperature=1.0)
expected_gen_shape = (2, 18) # 8 original + 10 new tokens
assert generated.shape == expected_gen_shape, f"Generated shape should be {expected_gen_shape}, got {generated.shape}"
# Verify original tokens are preserved
assert np.array_equal(generated.data[:, :8], start_tokens.data), "Original tokens should be preserved"
# Test different model configurations
small_transformer = Transformer(
vocab_size=500, embed_dim=128, num_heads=4, num_layers=2, hidden_dim=256
)
small_input = Tensor(np.random.randint(0, 500, (2, 16)))
small_logits = small_transformer.forward(small_input)
expected_small_shape = (2, 16, 500)
assert small_logits.shape == expected_small_shape, "Small transformer should work"
# Test pre-norm vs post-norm
post_norm_transformer = Transformer(
vocab_size=vocab_size, embed_dim=embed_dim, num_heads=num_heads,
num_layers=2, hidden_dim=hidden_dim, pre_norm=False
)
post_norm_logits = post_norm_transformer.forward(input_tensor)
pre_norm_logits = Transformer(
vocab_size=vocab_size, embed_dim=embed_dim, num_heads=num_heads,
num_layers=2, hidden_dim=hidden_dim, pre_norm=True
).forward(input_tensor)
assert not np.allclose(post_norm_logits.data, pre_norm_logits.data), \
"Pre-norm and post-norm should produce different outputs"
# Test memory usage calculation
memory_stats = transformer.get_memory_usage()
assert 'total_memory_mb' in memory_stats, "Should provide memory statistics"
assert memory_stats['total_memory_mb'] > 0, "Should have positive memory usage"
assert memory_stats['total_parameters'] > 0, "Should count parameters"
# Verify memory breakdown
assert memory_stats['embedding_memory_mb'] > 0, "Should have embedding memory"
assert memory_stats['transformer_blocks_memory_mb'] > 0, "Should have transformer block memory"
assert memory_stats['lm_head_memory_mb'] > 0, "Should have language modeling head memory"
print("✅ Complete transformer model tests passed!")
print(f"✅ Forward pass produces correct logit shapes")
print(f"✅ Causal masking works across all {num_layers} layers")
print(f"✅ Text generation capability verified")
print(f"✅ Total parameters: {memory_stats['total_parameters']:,}")
print(f"✅ Total memory: {memory_stats['total_memory_mb']:.2f}MB")
print(f"✅ Pre-norm and post-norm architectures work correctly")
# Test function defined (called in main block)In [ ]:
#| export
import time
class TransformerProfiler:
"""
Performance profiling toolkit for transformer architectures.
Helps ML engineers understand computational costs, memory scaling,
and architectural trade-offs in transformer-based models.
"""
def __init__(self):
self.results = {}
def measure_scaling_with_depth(self, base_config: Dict, layer_counts: List[int]) -> Dict:
"""
Measure how transformer performance scales with number of layers.
TODO: Implement transformer depth scaling measurement.
STEP-BY-STEP IMPLEMENTATION:
1. Create transformers with different layer counts
2. Measure memory usage and computation time for each
3. Calculate scaling patterns (should be linear with depth)
4. Analyze parameter growth and memory requirements
5. Return comprehensive scaling analysis
EXPECTED SCALING:
- Parameters: Linear with depth
- Memory: Linear with depth
- Computation: Linear with depth
- Quality: Generally improves with depth (to a point)
Args:
base_config: Base transformer configuration
layer_counts: List of layer counts to test
Returns:
Dictionary with scaling analysis results
"""
### BEGIN SOLUTION
scaling_results = {}
# Test input
batch_size = 4
seq_len = 32
vocab_size = base_config['vocab_size']
test_input = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))
for num_layers in layer_counts:
# Create transformer with this depth
transformer = Transformer(
vocab_size=base_config['vocab_size'],
embed_dim=base_config['embed_dim'],
num_heads=base_config['num_heads'],
num_layers=num_layers,
hidden_dim=base_config['hidden_dim'],
max_seq_length=base_config.get('max_seq_length', 128)
)
# Measure memory usage
memory_stats = transformer.get_memory_usage()
# Measure computation time
start_time = time.time()
logits = transformer.forward(test_input)
end_time = time.time()
computation_time_ms = (end_time - start_time) * 1000
# Calculate throughput
total_tokens = batch_size * seq_len
tokens_per_second = total_tokens / (end_time - start_time) if end_time > start_time else 0
scaling_results[num_layers] = {
'num_layers': num_layers,
'total_parameters': memory_stats['total_parameters'],
'total_memory_mb': memory_stats['total_memory_mb'],
'computation_time_ms': computation_time_ms,
'tokens_per_second': tokens_per_second,
'memory_per_layer_mb': memory_stats['transformer_blocks_memory_mb'] / num_layers if num_layers > 0 else 0,
'parameters_per_layer': (memory_stats['total_parameters'] -
base_config['vocab_size'] * base_config['embed_dim'] * 2) // num_layers if num_layers > 0 else 0
}
return scaling_results
### END SOLUTION
def analyze_width_vs_depth_tradeoffs(self, base_params: int, configurations: List[Dict]) -> Dict:
"""
Compare different ways to allocate a fixed parameter budget.
This function is PROVIDED to show parameter allocation analysis.
"""
print(f"📊 WIDTH vs DEPTH TRADE-OFF ANALYSIS")
print(f"Target parameter budget: ~{base_params:,} parameters")
print("=" * 70)
results = {}
# Test input
batch_size = 4
seq_len = 32
test_input = Tensor(np.random.randint(0, 1000, (batch_size, seq_len)))
print(f"{'Config':<15} {'Layers':<7} {'Embed':<6} {'Heads':<6} {'Hidden':<7} {'Params':<12} {'Time (ms)':<10} {'Memory'}")
print("-" * 80)
for i, config in enumerate(configurations):
try:
# Create transformer
transformer = Transformer(
vocab_size=1000, # Fixed vocab size
embed_dim=config['embed_dim'],
num_heads=config['num_heads'],
num_layers=config['num_layers'],
hidden_dim=config['hidden_dim'],
max_seq_length=128
)
# Get actual parameter count
memory_stats = transformer.get_memory_usage()
actual_params = memory_stats['total_parameters']
# Measure performance
start_time = time.time()
logits = transformer.forward(test_input)
computation_time = (time.time() - start_time) * 1000
config_name = f"Config_{i+1}"
results[config_name] = {
'config': config,
'actual_parameters': actual_params,
'computation_time_ms': computation_time,
'memory_mb': memory_stats['total_memory_mb'],
'parameter_efficiency': abs(actual_params - base_params) / base_params
}
print(f"{config_name:<15} {config['num_layers']:<7} {config['embed_dim']:<6} "
f"{config['num_heads']:<6} {config['hidden_dim']:<7} {actual_params:<12,} "
f"{computation_time:<10.2f} {memory_stats['total_memory_mb']:.1f}MB")
except Exception as e:
print(f"{config_name:<15} ERROR: {str(e)[:50]}")
# Analysis
print(f"\n💡 TRADE-OFF INSIGHTS:")
print(f" - Deeper models: Better at learning complex patterns, more sequential")
print(f" - Wider models: More parallelizable, can capture diverse features")
print(f" - More heads: Richer attention patterns, more computation")
print(f" - Hidden dimension: Affects FFN capacity, major parameter contributor")
return results
def simulate_production_scaling(self, model_sizes: List[str]) -> Dict:
"""
Simulate memory and computation requirements for production model sizes.
This function is PROVIDED to show production scaling analysis.
"""
print(f"\n🏭 PRODUCTION MODEL SCALING SIMULATION")
print("=" * 60)
# Production model configurations (simplified)
size_configs = {
'Small': {'vocab_size': 50000, 'embed_dim': 512, 'num_heads': 8, 'num_layers': 6, 'hidden_dim': 2048},
'Medium': {'vocab_size': 50000, 'embed_dim': 768, 'num_heads': 12, 'num_layers': 12, 'hidden_dim': 3072},
'Large': {'vocab_size': 50000, 'embed_dim': 1024, 'num_heads': 16, 'num_layers': 24, 'hidden_dim': 4096},
'XL': {'vocab_size': 50000, 'embed_dim': 1280, 'num_heads': 20, 'num_layers': 36, 'hidden_dim': 5120}
}
results = {}
print(f"{'Model Size':<12} {'Parameters':<12} {'Memory (GB)':<12} {'Training GPU':<12} {'Inference'}")
print("-" * 70)
for size in model_sizes:
if size not in size_configs:
continue
config = size_configs[size]
# Estimate parameters
# Embedding: vocab_size * embed_dim * 2 (input + output)
embedding_params = config['vocab_size'] * config['embed_dim'] * 2
# Per layer:
# - Attention: 4 * embed_dim^2 (Q, K, V, O projections)
# - FFN: 2 * embed_dim * hidden_dim + embed_dim + hidden_dim (weights + biases)
# - LayerNorm: 2 * embed_dim * 2 (two norms per layer)
attention_params_per_layer = 4 * config['embed_dim'] ** 2
ffn_params_per_layer = 2 * config['embed_dim'] * config['hidden_dim'] + config['embed_dim'] + config['hidden_dim']
norm_params_per_layer = 4 * config['embed_dim']
layer_params = attention_params_per_layer + ffn_params_per_layer + norm_params_per_layer
total_params = embedding_params + layer_params * config['num_layers']
# Estimate memory (parameters + activations + gradients for training)
param_memory_gb = total_params * 4 / (1024**3) # 4 bytes per float32
# Training memory: parameters + gradients + optimizer states + activations
training_memory_gb = param_memory_gb * 4 # Rough estimate (param + grad + 2x optimizer states)
# Inference memory: just parameters + activations
inference_memory_gb = param_memory_gb * 1.5 # Parameters + activation memory
# GPU requirements (very rough estimates)
if training_memory_gb < 24:
training_gpu = "Single RTX 4090"
elif training_memory_gb < 80:
training_gpu = "Single A100"
else:
training_gpu = "Multi-GPU"
if inference_memory_gb < 12:
inference_req = "RTX 4060 Ti"
elif inference_memory_gb < 24:
inference_req = "RTX 4090"
else:
inference_req = "A100+"
results[size] = {
'config': config,
'total_parameters': total_params,
'training_memory_gb': training_memory_gb,
'inference_memory_gb': inference_memory_gb,
'training_gpu_req': training_gpu,
'inference_gpu_req': inference_req
}
print(f"{size:<12} {total_params/1e6:.1f}M {training_memory_gb:.1f} {training_gpu:<12} {inference_req}")
print(f"\n📈 SCALING OBSERVATIONS:")
print(f" - Model size grows super-linearly with dimension increases")
print(f" - Memory requirements dominate deployment decisions")
print(f" - Training requires 3-4x more memory than inference")
print(f" - Multi-GPU becomes necessary for large models")
return results
def analyze_transformer_system_design():
"""
Comprehensive analysis of transformer system design choices and trade-offs.
This function is PROVIDED to show systems-level design thinking.
"""
print("🏗️ TRANSFORMER SYSTEM DESIGN ANALYSIS")
print("=" * 60)
# Architecture decision analysis
design_choices = {
'Layer Normalization': {
'Pre-norm': {'stability': 'High', 'training': 'Easier', 'performance': 'Good'},
'Post-norm': {'stability': 'Lower', 'training': 'Harder', 'performance': 'Potentially better'}
},
'Attention Patterns': {
'Full attention': {'complexity': 'O(N²)', 'quality': 'Best', 'scalability': 'Limited'},
'Sparse attention': {'complexity': 'O(N√N)', 'quality': 'Good', 'scalability': 'Better'},
'Linear attention': {'complexity': 'O(N)', 'quality': 'Reduced', 'scalability': 'Excellent'}
},
'Feed-Forward Size': {
'2x embed_dim': {'parameters': 'Low', 'capacity': 'Limited', 'speed': 'Fast'},
'4x embed_dim': {'parameters': 'Standard', 'capacity': 'Good', 'speed': 'Medium'},
'8x embed_dim': {'parameters': 'High', 'capacity': 'High', 'speed': 'Slow'}
}
}
print("🎯 ARCHITECTURAL DESIGN CHOICES:")
for category, choices in design_choices.items():
print(f"\n{category}:")
for choice, properties in choices.items():
prop_str = ", ".join([f"{k}: {v}" for k, v in properties.items()])
print(f" - {choice}: {prop_str}")
# Memory scaling analysis
print(f"\n📊 MEMORY SCALING PATTERNS:")
print(f"Component breakdown for typical transformer:")
print(f" - Token embeddings: vocab_size × embed_dim parameters")
print(f" - Position encodings: 0 parameters (sinusoidal) or seq_len × embed_dim (learned)")
print(f" - Attention layers: 4 × embed_dim² parameters per layer")
print(f" - Feed-forward: 2 × embed_dim × hidden_dim parameters per layer")
print(f" - Layer normalization: 2 × embed_dim parameters per layer")
print(f" - Output projection: embed_dim × vocab_size parameters")
print(f"\n🔧 OPTIMIZATION STRATEGIES:")
optimization_techniques = [
"Gradient checkpointing: Trade computation for memory",
"Mixed precision training: Use FP16 for 2x memory reduction",
"Parameter sharing: Share weights across layers",
"Sparse attention: Reduce quadratic scaling",
"Model parallelism: Distribute layers across GPUs",
"Pipeline parallelism: Process different batch elements on different GPUs",
"Activation checkpointing: Recompute activations instead of storing"
]
for technique in optimization_techniques:
print(f" - {technique}")
print(f"\n🎯 PRODUCTION DEPLOYMENT CONSIDERATIONS:")
deployment_factors = [
"Batch size: Larger batches improve GPU utilization but increase memory",
"Sequence length: Quadratic impact on attention memory",
"Model depth: Linear impact on memory and computation",
"Model width: Quadratic impact on attention parameters",
"Precision: FP32 vs FP16 vs INT8 trade-offs",
"Hardware: GPU memory and compute capabilities",
"Latency requirements: Real-time vs batch processing",
"Throughput requirements: Tokens per second targets"
]
for factor in deployment_factors:
print(f" - {factor}")In [ ]:
def test_transformer_profiler():
"""Test transformer profiler with various scenarios."""
print("🔬 Unit Test: Transformer Performance Profiler...")
profiler = TransformerProfiler()
# Test depth scaling measurement
base_config = {
'vocab_size': 500,
'embed_dim': 128,
'num_heads': 4,
'hidden_dim': 256
}
layer_counts = [1, 2, 4]
depth_results = profiler.measure_scaling_with_depth(base_config, layer_counts)
# Verify depth scaling results
assert len(depth_results) == len(layer_counts), f"Should test {len(layer_counts)} layer counts"
for num_layers in layer_counts:
assert num_layers in depth_results, f"Should include results for {num_layers} layers"
result = depth_results[num_layers]
# Verify required metrics
required_keys = ['num_layers', 'total_parameters', 'total_memory_mb',
'computation_time_ms', 'tokens_per_second']
for key in required_keys:
assert key in result, f"Missing metric: {key} for {num_layers} layers"
assert isinstance(result[key], (int, float)), f"Invalid type for {key}"
# Verify reasonable values
assert result['num_layers'] == num_layers, "Should store correct layer count"
assert result['total_parameters'] > 0, "Should have positive parameter count"
assert result['total_memory_mb'] > 0, "Should have positive memory usage"
# Test that parameters and memory scale roughly linearly with depth
if len(layer_counts) >= 2:
shallow = depth_results[layer_counts[0]]
deep = depth_results[layer_counts[-1]]
layer_ratio = deep['num_layers'] / shallow['num_layers']
param_ratio = deep['total_parameters'] / shallow['total_parameters']
memory_ratio = deep['total_memory_mb'] / shallow['total_memory_mb']
# Allow some deviation due to fixed costs (embeddings, etc.)
assert 1.0 < param_ratio < layer_ratio * 2, f"Parameters should scale sub-linearly, got {param_ratio:.2f}"
assert 1.0 < memory_ratio < layer_ratio * 2, f"Memory should scale sub-linearly, got {memory_ratio:.2f}"
print("✅ Depth scaling measurement test passed")
# Test width vs depth analysis
configurations = [
{'embed_dim': 128, 'num_heads': 4, 'num_layers': 4, 'hidden_dim': 256},
{'embed_dim': 256, 'num_heads': 8, 'num_layers': 2, 'hidden_dim': 512},
]
width_depth_results = profiler.analyze_width_vs_depth_tradeoffs(100000, configurations)
# Verify width vs depth results
assert len(width_depth_results) > 0, "Should analyze at least one configuration"
for config_name, result in width_depth_results.items():
assert 'config' in result, "Should include configuration"
assert 'actual_parameters' in result, "Should count actual parameters"
assert 'computation_time_ms' in result, "Should measure computation time"
assert result['actual_parameters'] > 0, "Should have positive parameter count"
print("✅ Width vs depth analysis test passed")
# Test production scaling simulation
production_results = profiler.simulate_production_scaling(['Small', 'Medium'])
# Verify production scaling results
for size, result in production_results.items():
assert 'config' in result, "Should include model configuration"
assert 'total_parameters' in result, "Should estimate total parameters"
assert 'training_memory_gb' in result, "Should estimate training memory"
assert 'inference_memory_gb' in result, "Should estimate inference memory"
# Verify reasonable scaling
assert result['total_parameters'] > 1e6, "Should have millions of parameters"
assert result['training_memory_gb'] > result['inference_memory_gb'], "Training should require more memory"
print("✅ Production scaling simulation test passed")
print("🎯 Transformer Profiler: All tests passed!")
# Test function defined (called in main block)In [ ]:
def test_complete_language_model_pipeline():
"""Test complete language model pipeline integration."""
print("🧪 Integration Test: Complete Language Model Pipeline...")
# Create a small but complete language model
vocab_size = 1000
embed_dim = 256
num_heads = 8
num_layers = 4
hidden_dim = 512
max_seq_length = 64
print(f" Creating transformer with {num_layers} layers, {embed_dim} dimensions...")
transformer = Transformer(
vocab_size=vocab_size,
embed_dim=embed_dim,
num_heads=num_heads,
num_layers=num_layers,
hidden_dim=hidden_dim,
max_seq_length=max_seq_length
)
# Test 1: Basic text processing pipeline
print(" Testing basic text processing pipeline...")
batch_size = 4
seq_len = 32
# Simulate tokenized input
input_ids = np.random.randint(0, vocab_size, (batch_size, seq_len))
input_tensor = Tensor(input_ids)
# Forward pass
logits = transformer.forward(input_tensor)
expected_shape = (batch_size, seq_len, vocab_size)
assert logits.shape == expected_shape, f"Expected {expected_shape}, got {logits.shape}"
# Test that logits are reasonable (not all zeros/inf/nan)
assert not np.all(logits.data == 0), "Logits should not all be zero"
assert not np.any(np.isinf(logits.data)), "Logits should not contain inf"
assert not np.any(np.isnan(logits.data)), "Logits should not contain nan"
print(f" Forward pass successful: {logits.shape}")
# Test 2: Language modeling with causal mask
print(" Testing language modeling with causal attention...")
causal_mask = np.triu(np.ones((seq_len, seq_len)), k=1)
causal_mask = 1 - causal_mask # Convert to attention mask
masked_logits, all_attention = transformer.forward(
input_tensor, mask=Tensor(causal_mask), return_attention_weights=True
)
assert len(all_attention) == num_layers, f"Should return attention from {num_layers} layers"
# Verify causal masking works across all layers
for layer_idx, attn_weights in enumerate(all_attention):
# Check a few positions to ensure masking works
for i in range(min(5, seq_len)):
for j in range(i+1, min(i+5, seq_len)):
future_attention = attn_weights.data[:, :, i, j] # All heads, all batches
assert np.all(future_attention < 1e-5), \
f"Layer {layer_idx}: future attention at ({i},{j}) should be ~0"
print(f" Causal masking verified across all layers")
# Test 3: Text generation
print(" Testing autoregressive text generation...")
# Start with a shorter sequence for generation
gen_start = Tensor(np.random.randint(0, vocab_size, (2, 8)))
generated = transformer.generate(gen_start, max_new_tokens=8, temperature=1.0)
expected_gen_shape = (2, 16) # 8 start + 8 generated
assert generated.shape == expected_gen_shape, f"Expected {expected_gen_shape}, got {generated.shape}"
# Verify original tokens preserved
assert np.array_equal(generated.data[:, :8], gen_start.data), "Should preserve original tokens"
# Verify new tokens are valid
new_tokens = generated.data[:, 8:]
assert np.all(new_tokens >= 0), "Generated tokens should be >= 0"
assert np.all(new_tokens < vocab_size), f"Generated tokens should be < {vocab_size}"
print(f" Generated {new_tokens.shape[1]} new tokens successfully")
# Test 4: Different sequence lengths
print(" Testing variable sequence lengths...")
for test_seq_len in [16, 32, 48]:
if test_seq_len > max_seq_length:
continue
test_input = Tensor(np.random.randint(0, vocab_size, (2, test_seq_len)))
test_logits = transformer.forward(test_input)
expected_test_shape = (2, test_seq_len, vocab_size)
assert test_logits.shape == expected_test_shape, f"Failed for seq_len {test_seq_len}"
print(f" Variable sequence lengths work correctly")
# Test 5: Memory usage analysis
print(" Analyzing memory usage...")
memory_stats = transformer.get_memory_usage()
print(f" Model parameters: {memory_stats['total_parameters']:,}")
print(f" Model memory: {memory_stats['total_memory_mb']:.1f}MB")
print(f" Embedding memory: {memory_stats['embedding_memory_mb']:.1f}MB")
print(f" Transformer blocks: {memory_stats['transformer_blocks_memory_mb']:.1f}MB")
print(f" LM head: {memory_stats['lm_head_memory_mb']:.1f}MB")
# Verify memory breakdown makes sense
component_memory = (memory_stats['embedding_memory_mb'] +
memory_stats['transformer_blocks_memory_mb'] +
memory_stats['lm_head_memory_mb'])
# Allow small difference due to final norm layer
memory_diff = abs(memory_stats['total_memory_mb'] - component_memory)
assert memory_diff < 1.0, f"Memory breakdown doesn't add up: {memory_diff:.2f}MB difference"
# Test 6: Performance characteristics
print(" Testing performance characteristics...")
# Time multiple forward passes
num_iterations = 5
start_time = time.time()
for _ in range(num_iterations):
_ = transformer.forward(input_tensor)
total_time = time.time() - start_time
avg_time_per_forward = total_time / num_iterations
tokens_per_second = (batch_size * seq_len) / avg_time_per_forward
print(f" Average forward pass: {avg_time_per_forward*1000:.2f}ms")
print(f" Processing speed: {tokens_per_second:.0f} tokens/second")
# Verify reasonable performance
assert avg_time_per_forward < 1.0, "Forward pass should be < 1 second"
assert tokens_per_second > 50, "Should process > 50 tokens/second"
# Test 7: Gradient flow (simulated)
print(" Testing gradient flow through layers...")
# Create slightly different inputs to test sensitivity
input_1 = Tensor(input_ids.copy())
input_2 = Tensor(input_ids.copy())
input_2.data[0, 0] = (input_2.data[0, 0] + 1) % vocab_size # Change one token
logits_1 = transformer.forward(input_1)
logits_2 = transformer.forward(input_2)
# Outputs should be different (model is sensitive to input changes)
output_diff = np.mean(np.abs(logits_1.data - logits_2.data))
assert output_diff > 1e-6, f"Model should be sensitive to input changes, diff: {output_diff}"
# But not too different (model should be stable)
assert output_diff < 100, f"Model should be stable, large diff: {output_diff}"
print(f" Model shows appropriate sensitivity to input changes")
print("✅ Complete language model pipeline integration test passed!")
print(f"✅ Forward pass, masking, generation, and performance verified")
print(f"✅ Model processes {tokens_per_second:.0f} tokens/second")
print(f"✅ Memory footprint: {memory_stats['total_memory_mb']:.1f}MB")
# Test function defined (called in main block)In [ ]:
if __name__ == "__main__":
# Run all unit tests
test_unit_layer_norm()
test_unit_feed_forward()
test_unit_transformer_block()
test_unit_transformer_model()
test_transformer_profiler()
test_complete_language_model_pipeline()
print("\n" + "="*60)
print("🔍 TRANSFORMER SYSTEMS ANALYSIS")
print("="*60)
# Performance analysis
profiler = TransformerProfiler()
# Test transformer scaling with different depths
print("📈 TRANSFORMER DEPTH SCALING ANALYSIS:")
base_config = {
'vocab_size': 1000,
'embed_dim': 256,
'num_heads': 8,
'hidden_dim': 1024
}
layer_counts = [2, 4, 8, 12]
depth_results = profiler.measure_scaling_with_depth(base_config, layer_counts)
# Analyze scaling patterns
print(f"\n{'Layers':<7} {'Parameters':<12} {'Memory (MB)':<12} {'Time (ms)':<10} {'Tokens/sec':<10}")
print("-" * 60)
for num_layers in layer_counts:
result = depth_results[num_layers]
print(f"{num_layers:<7} {result['total_parameters']:<12,} {result['total_memory_mb']:<12.1f} "
f"{result['computation_time_ms']:<10.2f} {result['tokens_per_second']:<10.0f}")
# Width vs depth trade-off analysis
print("\n" + "="*60)
configurations = [
{'embed_dim': 256, 'num_heads': 8, 'num_layers': 8, 'hidden_dim': 1024}, # Deep & narrow
{'embed_dim': 512, 'num_heads': 16, 'num_layers': 4, 'hidden_dim': 2048}, # Wide & shallow
{'embed_dim': 384, 'num_heads': 12, 'num_layers': 6, 'hidden_dim': 1536}, # Balanced
]
width_depth_results = profiler.analyze_width_vs_depth_tradeoffs(2000000, configurations)
# Production scaling simulation
print("\n" + "="*60)
production_results = profiler.simulate_production_scaling(['Small', 'Medium', 'Large'])
# Systems design analysis
print("\n" + "="*60)
analyze_transformer_system_design()
# Demonstrate realistic language model setup
print("\n" + "="*60)
print("🏗️ REALISTIC LANGUAGE MODEL DEMONSTRATION")
print("="*60)
# Create a realistic small language model
vocab_size = 5000
embed_dim = 512
num_heads = 8
num_layers = 6
hidden_dim = 2048
max_seq_length = 256
print(f"Language model configuration:")
print(f" Vocabulary: {vocab_size:,} tokens")
print(f" Embedding dimension: {embed_dim}")
print(f" Attention heads: {num_heads}")
print(f" Transformer layers: {num_layers}")
print(f" Feed-forward dimension: {hidden_dim}")
print(f" Max sequence length: {max_seq_length}")
# Create the model
language_model = Transformer(
vocab_size=vocab_size,
embed_dim=embed_dim,
num_heads=num_heads,
num_layers=num_layers,
hidden_dim=hidden_dim,
max_seq_length=max_seq_length,
pre_norm=True
)
# Analyze model characteristics
memory_stats = language_model.get_memory_usage()
print(f"\nModel characteristics:")
print(f" Total parameters: {memory_stats['total_parameters']:,}")
print(f" Model size: {memory_stats['total_memory_mb']:.1f}MB")
print(f" Embedding table: {memory_stats['embedding_memory_mb']:.1f}MB ({memory_stats['embedding_memory_mb']/memory_stats['total_memory_mb']*100:.1f}%)")
print(f" Transformer layers: {memory_stats['transformer_blocks_memory_mb']:.1f}MB ({memory_stats['transformer_blocks_memory_mb']/memory_stats['total_memory_mb']*100:.1f}%)")
print(f" Output projection: {memory_stats['lm_head_memory_mb']:.1f}MB ({memory_stats['lm_head_memory_mb']/memory_stats['total_memory_mb']*100:.1f}%)")
# Performance simulation
batch_size = 8
seq_len = 128
test_input = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))
start_time = time.time()
logits = language_model.forward(test_input)
forward_time = time.time() - start_time
tokens_per_second = (batch_size * seq_len) / forward_time
print(f"\nPerformance simulation:")
print(f" Batch size: {batch_size}, Sequence length: {seq_len}")
print(f" Forward pass time: {forward_time*1000:.2f}ms")
print(f" Throughput: {tokens_per_second:.0f} tokens/second")
print(f" Memory for batch: {logits.data.nbytes/(1024*1024):.1f}MB")
# Text generation example
print(f"\nText generation example:")
start_sequence = Tensor(np.random.randint(0, vocab_size, (1, 10)))
generated = language_model.generate(start_sequence, max_new_tokens=20, temperature=0.8)
print(f" Input sequence: {start_sequence.data[0].tolist()}")
print(f" Generated tokens: {generated.data[0, 10:].tolist()}")
print(f" Generation completed successfully")
# Scaling predictions
print(f"\nScaling analysis:")
current_params = memory_stats['total_parameters']
# Estimate for different scales
scaling_factors = [2, 5, 10]
for factor in scaling_factors:
scaled_params = current_params * factor
scaled_memory_gb = memory_stats['total_memory_mb'] * factor / 1024
print(f" {factor}x scale: {scaled_params/1e6:.0f}M params, ~{scaled_memory_gb:.1f}GB memory")
print("\n" + "="*60)
print("🎯 TRANSFORMERS MODULE COMPLETE!")
print("="*60)
print("All transformer tests passed!")
print("Complete language model architecture implemented!")
print("Ready for production deployment and optimization!")In [ ]:
"""
YOUR REFLECTION ON TRANSFORMER ARCHITECTURE OPTIMIZATION:
TODO: Replace this text with your thoughtful response about transformer architecture optimization for diverse deployment scenarios.
Consider addressing:
- How would you allocate parameter budgets across depth, width, and attention heads for different scenarios?
- What architecture search strategies would you use to optimize within hardware constraints?
- How would you implement adaptive model scaling that adjusts to available resources?
- What approaches would you use to maintain model quality across different deployment environments?
- How would you balance latency, throughput, and resource constraints in architectural decisions?
Write a strategic analysis connecting your transformer implementations to real architecture optimization challenges.
GRADING RUBRIC (Instructor Use):
- Demonstrates understanding of transformer architecture trade-offs and optimization (3 points)
- Designs practical approaches to parameter allocation and architecture search (3 points)
- Addresses adaptive scaling and hardware-aware optimization (2 points)
- Shows systems thinking about production deployment optimization (2 points)
- Clear strategic reasoning with architecture optimization insights (bonus points for innovative approaches)
"""
### BEGIN SOLUTION
# Student response area - instructor will replace this section during grading setup
# This is a manually graded question requiring strategic analysis of transformer architecture optimization
# Students should demonstrate understanding of architecture design and production deployment challenges
### END SOLUTIONIn [ ]:
"""
YOUR REFLECTION ON TRANSFORMER TRAINING AND INFERENCE SYSTEM DESIGN:
TODO: Replace this text with your thoughtful response about transformer training and inference system architecture.
Consider addressing:
- How would you design distributed training for billion-parameter transformers with memory constraints?
- What strategies would you use for efficient inference serving with millisecond latency requirements?
- How would you manage model deployment across heterogeneous hardware environments?
- What approaches would you use to maintain numerical stability during distributed training?
- How would you ensure consistent inference performance across different deployment targets?
Write a system design analysis connecting your transformer implementation to large-scale training and serving challenges.
GRADING RUBRIC (Instructor Use):
- Shows understanding of distributed training and inference serving challenges (3 points)
- Designs practical approaches to memory management and latency optimization (3 points)
- Addresses heterogeneous deployment and numerical stability considerations (2 points)
- Demonstrates systems thinking about training-inference system coordination (2 points)
- Clear system design reasoning with scalability insights (bonus points for comprehensive system architecture)
"""
### BEGIN SOLUTION
# Student response area - instructor will replace this section during grading setup
# This is a manually graded question requiring system design for transformer training and inference
# Students should demonstrate knowledge of distributed systems and production deployment architecture
### END SOLUTIONIn [ ]:
"""
YOUR REFLECTION ON TRANSFORMER OPTIMIZATION AND PRODUCTION DEPLOYMENT:
TODO: Replace this text with your thoughtful response about transformer production deployment system design.
Consider addressing:
- How would you implement end-to-end optimization spanning tokenization through generation?
- What strategies would you use for efficient model serving with dynamic batching and request routing?
- How would you enable seamless model updates without service interruption?
- What approaches would you use to maintain pipeline modularity while optimizing holistically?
- How would you support diverse model variants and fine-tuned versions in production?
Write a deployment analysis connecting your transformer implementation to complete production system optimization.
GRADING RUBRIC (Instructor Use):
- Understands end-to-end optimization and production deployment challenges (3 points)
- Designs practical approaches to model serving and continuous deployment (3 points)
- Addresses modularity and system integration considerations (2 points)
- Shows systems thinking about holistic pipeline optimization (2 points)
- Clear deployment reasoning with production optimization insights (bonus points for innovative system design)
"""
### BEGIN SOLUTION
# Student response area - instructor will replace this section during grading setup
# This is a manually graded question requiring understanding of production transformer deployment optimization
# Students should demonstrate knowledge of end-to-end system design and continuous deployment strategies
### END SOLUTION