Grouped Query Attention › Converting MHA to GQA

Michael BrenndoerferUpdated August 6, 20254 min read
Language AI HandbookAll chapters
Part 29 of 67Modern Decoder Models
Chapter 209Grouped Query Attention › Converting MHA to GQA
Reading section 10 of 14 · Converting MHA to GQA · 4 minSection 10 of 14 · 4 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 209 of 513

Section 10 of 14 · 4 min read

Grouped Query Attention › Converting MHA to GQAOne section · Sections in this chapter10 / 14 ↓
Sections in this chapter14
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

What if you have an existing MHA model and want the benefits of GQA? The LLaMA 2 paper describes a conversion process called "uptraining" that adapts pretrained MHA models to use GQA.

The conversion process works as follows:

  1. Mean pooling: Average the key and value projection weights across the heads that will share a KV representation
  2. Uptraining: Continue training the converted model on a small fraction (typically 5%) of the original training data
  3. Validation: Verify quality recovery on held-out benchmarks
In[24]:
Code
def convert_mha_to_gqa(
    mha_k_weight: torch.Tensor,  # (d_model, num_heads * head_dim)
    mha_v_weight: torch.Tensor,  # (d_model, num_heads * head_dim)
    num_heads: int,
    num_kv_groups: int,
) -> tuple:
    """Convert MHA K,V weights to GQA by mean pooling within groups."""
    d_model = mha_k_weight.shape[0]
    head_dim = mha_k_weight.shape[1] // num_heads
    heads_per_group = num_heads // num_kv_groups

    # Reshape to (d_model, num_heads, head_dim)
    k_reshaped = mha_k_weight.view(d_model, num_heads, head_dim)
    v_reshaped = mha_v_weight.view(d_model, num_heads, head_dim)

    # Reshape to (d_model, num_kv_groups, heads_per_group, head_dim)
    k_grouped = k_reshaped.view(
        d_model, num_kv_groups, heads_per_group, head_dim
    )
    v_grouped = v_reshaped.view(
        d_model, num_kv_groups, heads_per_group, head_dim
    )

    # Mean pool within groups: (d_model, num_kv_groups, head_dim)
    gqa_k_weight = k_grouped.mean(dim=2)
    gqa_v_weight = v_grouped.mean(dim=2)

    # Reshape to (d_model, num_kv_groups * head_dim)
    gqa_k_weight = gqa_k_weight.view(d_model, num_kv_groups * head_dim)
    gqa_v_weight = gqa_v_weight.view(d_model, num_kv_groups * head_dim)

    return gqa_k_weight, gqa_v_weight


# Example conversion
d_model, num_heads, head_dim = 512, 8, 64
num_kv_groups = 2

# Simulate MHA weights
mha_k = torch.randn(d_model, num_heads * head_dim)
mha_v = torch.randn(d_model, num_heads * head_dim)

# Convert to GQA
gqa_k, gqa_v = convert_mha_to_gqa(mha_k, mha_v, num_heads, num_kv_groups)
Out[25]:
Console
MHA to GQA Weight Conversion
==================================================

Original MHA weights:
  K shape: torch.Size([512, 512])
  V shape: torch.Size([512, 512])

Converted GQA weights:
  K shape: torch.Size([512, 128])
  V shape: torch.Size([512, 128])

Parameter reduction: 4x for K,V projections
Out[26]:
Visualization
Heatmap of MHA key weight matrix slice showing 8 head columns with thin boundary lines between heads.
Original MHA key weights with 8 heads.
Heatmap of the same MHA key weights with thick green lines marking the two group boundaries for mean pooling.
Weights grouped for mean pooling (2 groups of 4).
Heatmap of GQA key weight matrix slice after mean pooling, showing 2 head columns.
GQA key weights after mean pooling (2 heads).

Mean pooling provides a reasonable initialization, but the converted model typically shows 2-5% quality degradation before uptraining. The uptraining phase allows the model to adapt its query projections to work effectively with the shared key-value representations. Research suggests that 5% of the original training compute is sufficient to recover most of the lost quality.