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)

Let me trace how the attention output flows back into the hidden states through the residual connections:

Attention Block Flow in Gemma2DecoderLayer

Step 3a: Pre-Attention

# Save the original input for residual connection
residual = hidden_states  # Shape: [batch_size, seq_len, 2304]

# Normalize for attention computation
hidden_states = self.input_layernorm(hidden_states)  # RMSNorm

Step 3b: Attention Computation

# The attention mechanism processes the normalized hidden states
hidden_states = self.self_attn(
    hidden_states=hidden_states,  # Normalized input
    freqs_cis=freqs_cis,
    kv_write_indices=kv_write_indices,
    kv_cache=kv_cache,
    mask=mask,
    local_mask=local_mask,
)
# Output shape: [batch_size, seq_len, 2304]

Step 3c: Post-Attention Integration

# Apply post-attention normalization
hidden_states = self.post_attention_layernorm(hidden_states)

# CRITICAL: Add residual connection
hidden_states = residual + hidden_states

The Residual Connection Pattern

The key insight is the residual connection that adds the original input back to the attention output:

Original Input (residual) ──────────────────┐
     │                                      │
     ↓                                      │
Input LayerNorm                             │
     │                                      │
     ↓                                      │
Self-Attention                              │
     │                                      │
     ↓                                      │
Post-Attention LayerNorm                    │
     │                                      │
     ↓                                      ↓
     └──────────────────> Addition <────────┘
                              │
                              ↓
                    Output (new hidden_states)

Why This Matters

1. Gradient Flow: The residual connection allows gradients to flow directly through the addition operation, bypassing the attention mechanism if needed.

2. Information Preservation: The original input information is preserved and combined with the attention-transformed information.

3. Mathematical Form:

   hidden_states_out = hidden_states_in + LayerNorm(Attention(LayerNorm(hidden_states_in)))
   

The Same Pattern Repeats for MLP

# After attention block completes, save new residual
residual = hidden_states  # This now includes attention output

# Pre-MLP normalization
hidden_states = self.pre_feedforward_layernorm(hidden_states)

# MLP computation
hidden_states = self.mlp(hidden_states)

# Post-MLP normalization
hidden_states = self.post_feedforward_layernorm(hidden_states)

# Another residual connection
hidden_states = residual + hidden_states

Data Flow Summary

For each Gemma2 layer:

1. Input: hidden_states from previous layer
2. Attention Block:
   • Save input as residual
   • Transform through normalized attention
   • Add back residual
3. MLP Block:
   • Save attention output as new residual
   • Transform through normalized MLP
   • Add back residual
4. Output: Enhanced hidden_states to next layer

This double-residual pattern with pre/post normalization is a key architectural feature of Gemma 2 that helps with training stability and allows the model to learn both local refinements (through attention/MLP) and preserve global information flow (through residuals).

Looking at the code, Gemma models always start with at least one token - typically the BOS (Beginning of Sequence) token. But let me explain the initial attention state:

Initial KV Cache State

When the model starts, the KV caches are initialized to zeros:

# From generate() method
kv_caches = []
for _ in range(self.config.num_hidden_layers):
    size = (batch_size, max_seq_len, self.config.num_key_value_heads,
            self.config.head_dim)
    dtype = self.config.get_dtype()
    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))

First Token Processing

When processing the first token (usually BOS):

1. Tokenization

# From Tokenizer.encode()
def encode(self, s: str, bos: bool = True, eos: bool = False) -> List[int]:
    t = self.sp_model.EncodeAsIds(s)
    if bos:
        t = [self.bos_id] + t  # Prepends BOS token
    return t

2. First Attention Computation

For the very first token at position 0:

# Input position is [0]
input_positions_tensor = torch.arange(0, 1)  # Just position 0

# The attention mask allows the first token to only attend to itself
mask_tensor = torch.full((1, 1, max_seq_len, max_seq_len), -2.3819763e38)
mask_tensor = torch.triu(mask_tensor, diagonal=1)
# This creates a causal mask where position 0 can only see position 0

3. Attention Mechanics

When the first token goes through attention:

# Query from first token
q = xq[:, 0:1, :, :]  # Shape: [batch, 1, heads, head_dim]

# Key and Value caches
# Initially all zeros except position 0 which gets updated:
k_cache[:, 0:1, :, :] = xk  # First token's key
v_cache[:, 0:1, :, :] = xv  # First token's value

# Attention scores
scores = torch.matmul(q, k.transpose(2, 3))
# Only position 0 has non-zero keys, so attention focuses entirely on position 0

The "Empty" Attention State

If hypothetically there were truly no tokens:

• The model couldn't function because embedding lookup would fail
• There would be no query vectors to compute attention

The closest conceptual "empty" state is:

• KV caches: All zeros (no previous context stored)
• Attention mask: Prevents looking at future positions
• First attention: Self-attention only (token attends to itself)

Practical Initialization

In practice, generation always starts with at least a BOS token:

# Example from the codebase
prompt_tokens = [self.tokenizer.encode(prompt) for prompt in prompts]
# This automatically includes BOS token due to default bos=True

So the initial attention state is:

1. Position 0: Contains BOS token embeddings
2. KV Cache: Zeros everywhere except position 0 after first forward pass
3. Attention Pattern: First token can only attend to itself
4. Subsequent tokens: Can attend to all previous positions including BOS

This is why language models can generate coherent text from just a BOS token - it provides the initial hidden state that seeds all subsequent generation.