Thread (29 posts)

Looking at the gemma_pytorch repository structure, here's how the Gemma 2 source code is organized:

Architecture Integration

Gemma 2 is integrated into a unified codebase that supports multiple Gemma architectures (1, 2, and 3). The code uses an Architecture enum to differentiate between versions:

class Architecture(enum.Enum):
    GEMMA_1 = 1
    GEMMA_2 = 2
    GEMMA_3 = 3

Configuration System

Gemma 2 has dedicated configuration functions in gemma/config.py:

• get_config_for_2b_v2() - 2B parameter variant
• get_config_for_9b() - 9B parameter variant
• get_config_for_27b() - 27B parameter variant

Each configuration specifies Gemma 2's unique features:

• Architecture: Architecture.GEMMA_2
• Pre/Post FFW Normalization: use_pre_ffw_norm=True, use_post_ffw_norm=True
• Logit Softcapping: final_logit_softcapping=30.0, attn_logit_softcapping=50.0
• Attention Types: Alternating pattern of LOCAL_SLIDING and GLOBAL attention
• Sliding Window: sliding_window_size=4096

Model Implementation

The model architecture in gemma/model.py uses conditional logic to handle different architectures:

Decoder Layer

Gemma 2 has its own decoder layer class Gemma2DecoderLayer with specific features:

• Additional normalization layers (pre/post feedforward)
• Support for both local sliding and global attention types
• Attention logit softcapping

Model Construction

for i in range(config.num_hidden_layers):
    if config.architecture == gemma_config.Architecture.GEMMA_1:
        self.layers.append(GemmaDecoderLayer(config))
    elif config.architecture in (Architecture.GEMMA_2, Architecture.GEMMA_3):
        attn_type = config.attn_types[i % len(config.attn_types)]
        self.layers.append(Gemma2DecoderLayer(config, attn_type))

Key Gemma 2 Features

1. Mixed Attention Pattern: Alternates between local sliding window attention and global attention across layers
2. Additional Normalizations: Extra RMSNorm layers before and after the MLP block
3. Softcapping: Applied to both attention logits and final output logits
4. Query Pre-attention Scaling: Custom scaling factor for certain variants (e.g., 27B model)

Variant Selection

The get_model_config() function maps variant strings to configurations:

• '2b-v2' → Gemma 2 2B model
• '9b' → Gemma 2 9B model
• '27b' → Gemma 2 27B model

This organization allows the codebase to support multiple Gemma architectures while sharing common components and maintaining architecture-specific features through configuration and conditional logic.

Let me trace the data flow through the Gemma 2b-v2 model from input tokens to output token selection:

1. Input Tokenization and Embedding

Function: GemmaForCausalLM.forward()

# Input: input_token_ids [batch_size, input_len]
hidden_states = self.embedder(input_token_ids)
# Output: [batch_size, input_len, hidden_size=2304]

# Normalize embeddings by sqrt(hidden_size)
normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
hidden_states = hidden_states * normalizer

2. Positional Encoding Setup

Function: GemmaForCausalLM.forward()

# Create frequency encodings for rotary embeddings
freqs_cis[AttentionType.LOCAL_SLIDING] = self.freqs_cis.index_select(0, input_positions)
freqs_cis[AttentionType.GLOBAL] = self.freqs_cis.index_select(0, input_positions)

3. Model Layers (26 layers for 2b-v2)

Function: GemmaModel.forward() → Gemma2DecoderLayer.forward()

Each layer alternates between LOCAL_SLIDING and GLOBAL attention:

3a. Pre-Attention Processing

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)  # RMSNorm

3b. Self-Attention

Function: GemmaAttention.forward()

# Project to Q, K, V
qkv = self.qkv_proj(hidden_states)  # Linear projection
xq, xk, xv = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

# Reshape for multi-head attention
# 2b-v2: 8 attention heads, 4 KV heads, head_dim=256
xq = xq.view(batch_size, -1, 8, 256)
xk = xk.view(batch_size, -1, 4, 256)
xv = xv.view(batch_size, -1, 4, 256)

# Apply rotary embeddings
xq = apply_rotary_emb(xq, freqs_cis=freqs_cis)
xk = apply_rotary_emb(xk, freqs_cis=freqs_cis)

# Update KV cache
k_cache.index_copy_(1, kv_write_indices, xk)
v_cache.index_copy_(1, kv_write_indices, xv)

# Compute attention scores
q.mul_(self.scaling)  # Scale by 1/sqrt(head_dim)
scores = torch.matmul(q, k.transpose(2, 3))

# Apply attention type-specific mask
if self.attn_type == AttentionType.LOCAL_SLIDING:
    # Use sliding window of 4096 tokens
    scores = scores + local_mask
else:
    # Use global causal mask
    scores = scores + mask

# Apply attention logit softcapping (50.0 for 2b-v2)
scores = scores / 50.0
scores = torch.tanh(scores)
scores = scores * 50.0

# Softmax and apply to values
scores = F.softmax(scores.float(), dim=-1)
output = torch.matmul(scores, v)

# Project back
output = self.o_proj(output)  # Linear projection

3c. Post-Attention Processing

hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = residual + hidden_states  # Residual connection

3d. MLP Block

Function: GemmaMLP.forward()

residual = hidden_states
hidden_states = self.pre_feedforward_layernorm(hidden_states)  # Pre-norm

# MLP with GeGLU activation
gate = self.gate_proj(hidden_states)  # Project to intermediate_size=9216
gate = F.gelu(gate, approximate="tanh")
up = self.up_proj(hidden_states)      # Project to intermediate_size=9216
fuse = gate * up                       # Element-wise multiply
hidden_states = self.down_proj(fuse)  # Project back to hidden_size=2304

hidden_states = self.post_feedforward_layernorm(hidden_states)  # Post-norm
hidden_states = residual + hidden_states  # Residual connection

4. Final Normalization

Function: GemmaModel.forward()

hidden_states = self.norm(hidden_states)  # Final RMSNorm
# Output: [batch_size, input_len, hidden_size=2304]

5. Token Selection

Function: Sampler.forward()

# Select output position
hidden_states = hidden_states.index_select(1, output_positions).squeeze(dim=1)
# Shape: [batch_size, hidden_size=2304]

# Compute logits
logits = torch.matmul(hidden_states, embedding.t())
# Shape: [batch_size, vocab_size=256000]

# Apply final logit softcapping (30.0 for 2b-v2)
logits = logits / 30.0
logits = torch.tanh(logits)
logits = logits * 30.0

# Temperature scaling (if temperature provided)
logits.div_(temperatures.unsqueeze(dim=1))

# Convert to probabilities
probs = torch.softmax(logits, dim=-1, dtype=torch.float)

# Apply top-p and top-k filtering
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
# ... filtering logic ...

# Sample next token
next_token_ids = torch.multinomial(probs, num_samples=1)

Summary of Data Flow

1. Embedding: Tokens → 2304-dim embeddings (normalized)
2. 26 Transformer Layers: Each with:
   • RMSNorm → Attention (local/global) → RMSNorm → Residual
   • RMSNorm → MLP (GeGLU) → RMSNorm → Residual
3. Final Norm: RMSNorm on hidden states
4. Logit Generation: Project to vocabulary space
5. Sampling: Softmax → Top-p/Top-k filtering → Multinomial sampling

The 2b-v2 model's unique features include:

• Alternating local (4096 window) and global attention
• Double normalization (pre/post) around MLP blocks
• Softcapping on both attention (50.0) and final logits (30.0)
• 8 attention heads with 4 KV heads (2:1 ratio)

For the Gemma 2b-v2 model, here are the QKV dimensions at each stage:

Initial Configuration (2b-v2)

• hidden_size: 2304
• num_attention_heads: 8
• num_key_value_heads: 4
• head_dim: 256

1. QKV Projection Input/Output

# Input to qkv_proj
hidden_states: [batch_size, input_len, 2304]

# After projection
qkv = self.qkv_proj(hidden_states)
# qkv: [batch_size, input_len, 4096]
# Where 4096 = q_size + kv_size + kv_size = 2048 + 1024 + 1024

2. After Splitting

xq, xk, xv = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
# xq: [batch_size, input_len, 2048]  # 8 heads × 256 dim
# xk: [batch_size, input_len, 1024]  # 4 heads × 256 dim  
# xv: [batch_size, input_len, 1024]  # 4 heads × 256 dim

3. After Reshaping for Multi-Head

xq = xq.view(batch_size, -1, 8, 256)   # [batch_size, input_len, 8, 256]
xk = xk.view(batch_size, -1, 4, 256)   # [batch_size, input_len, 4, 256]
xv = xv.view(batch_size, -1, 4, 256)   # [batch_size, input_len, 4, 256]

4. KV Cache Storage

k_cache: [batch_size, max_seq_len, 4, 256]
v_cache: [batch_size, max_seq_len, 4, 256]

5. For Attention Computation

# After transposing
q: [batch_size, 8, input_len, 256]

# K and V are expanded to match Q's head count
key = torch.repeat_interleave(key, 2, dim=2)    # 4→8 heads
value = torch.repeat_interleave(value, 2, dim=2) # 4→8 heads

# After expansion and transpose
k: [batch_size, 8, max_seq_len, 256]
v: [batch_size, 8, max_seq_len, 256]

6. Attention Scores

scores = torch.matmul(q, k.transpose(2, 3))
# scores: [batch_size, 8, input_len, max_seq_len]

Key Points:

• Q uses 8 attention heads (full)
• K and V use only 4 heads (half) for memory efficiency
• Each KV head is shared by 2 Q heads (Multi-Query Attention variant)
• Total QKV parameters: 2304 × 4096 = 9,437,184 parameters
• Memory saved: Using 4 KV heads instead of 8 saves ~25% of KV cache memory

This is a form of Grouped Query Attention (GQA) where multiple query heads share the same key-value pairs, reducing memory requirements while maintaining model quality.

KV Cache: General Concept

The KV cache is an optimization technique for autoregressive generation in transformers that avoids redundant computation.

The Problem

During generation, each new token needs to attend to all previous tokens. Without caching:

• Token 1: Compute K,V for position 0
• Token 2: Recompute K,V for positions 0,1
• Token 3: Recompute K,V for positions 0,1,2
• ...
• Token N: Recompute K,V for positions 0,1,2,...,N-1

This leads to O(n²) complexity for generating n tokens.

The Solution

Cache the K and V projections after computing them once:

• Token 1: Compute and cache K,V for position 0
• Token 2: Use cached K,V for position 0, compute and cache for position 1
• Token 3: Use cached K,V for positions 0,1, compute and cache for position 2

This reduces complexity to O(n).

Gemma Implementation

1. KV Cache Initialization

# In generate() method
kv_caches = []
for _ in range(self.config.num_hidden_layers):  # 26 layers for 2b-v2
    size = (batch_size, max_seq_len, self.config.num_key_value_heads, self.config.head_dim)
    # For 2b-v2: (batch_size, max_seq_len, 4, 256)
    k_cache = torch.zeros(size=size, dtype=dtype, device=device)
    v_cache = torch.zeros(size=size, dtype=dtype, device=device)
    kv_caches.append((k_cache, v_cache))

Each layer gets its own K and V cache tensors, pre-allocated for the maximum sequence length.

2. Writing to Cache

In GemmaAttention.forward():

# Compute new K,V for current positions
xk = xk.view(batch_size, -1, self.num_kv_heads, self.head_dim)
xv = xv.view(batch_size, -1, self.num_kv_heads, self.head_dim)

# Apply rotary embeddings (positional encoding)
xk = apply_rotary_emb(xk, freqs_cis=freqs_cis)

# Write to cache at specific positions
k_cache, v_cache = kv_cache
k_cache.index_copy_(1, kv_write_indices, xk)  # kv_write_indices = input_positions
v_cache.index_copy_(1, kv_write_indices, xv)

3. Reading from Cache

# Use the entire cache for attention
key = k_cache     # [batch_size, max_seq_len, num_kv_heads, head_dim]
value = v_cache   # [batch_size, max_seq_len, num_kv_heads, head_dim]

# Expand KV heads to match Q heads (4 → 8 for 2b-v2)
if self.num_kv_heads != self.num_heads:
    key = torch.repeat_interleave(key, self.num_queries_per_kv, dim=2)
    value = torch.repeat_interleave(value, self.num_queries_per_kv, dim=2)

4. Generation Flow

Prefill Phase (First Forward Pass):

# Process all prompt tokens at once
input_token_ids: [batch_size, prompt_len]
input_positions: [0, 1, 2, ..., prompt_len-1]

# Compute and cache K,V for all prompt positions
# k_cache[:, :prompt_len, :, :] now contains K values
# v_cache[:, :prompt_len, :, :] now contains V values

Generation Phase (Subsequent Passes):

# Process one new token at a time
input_token_ids: [batch_size, 1]  # Just the last generated token
input_positions: [current_position]  # e.g., [prompt_len], [prompt_len+1], ...

# Only compute K,V for the new position
# Cache is updated incrementally

5. Attention Computation with Cache

# Q is only computed for current tokens
q: [batch_size, num_heads, current_len, head_dim]

# K,V contain all positions up to max_seq_len
k: [batch_size, num_heads, max_seq_len, head_dim]
v: [batch_size, num_heads, max_seq_len, head_dim]

# Attention scores use full K cache
scores = torch.matmul(q, k.transpose(2, 3))
# scores: [batch_size, num_heads, current_len, max_seq_len]

# Mask ensures we only attend to valid positions
scores = scores + mask  # Causal mask prevents attending to future

Key Implementation Details

1. Pre-allocation: Caches are allocated for max_seq_len to avoid dynamic memory allocation
2. Sparse Updates: Only positions being processed get updated via index_copy_
3. Full Attention: During attention, the full cache is used but masked appropriately
4. Memory Efficiency: Using 4 KV heads instead of 8 reduces cache size by 50%

Example: Generating 3 Tokens

Starting with prompt "The cat":

1. Prefill:

   • Input: ["The", "cat"]
   • Cache positions 0,1 with K,V for "The", "cat"

2. Generate "sat":

   • Input: ["sat"] (from previous generation)
   • Cache position 2 with K,V for "sat"
   • Attention uses cached K,V from positions 0,1,2

3. Generate "on":

   • Input: ["on"]
   • Cache position 3 with K,V for "on"
   • Attention uses cached K,V from positions 0,1,2,3

Without caching, step 3 would recompute K,V for all previous tokens. With caching, we only compute K,V for the new token.