Transformer Internals for Fine-Tuning

A transformer decoder block contains two submodules: multi-head self-attention and a feed-forward network (FFN). Each submodule is preceded by a layer norm (RMSNorm in Llama-family models). Fine-tuning updates the weight matrices inside these submodules — understanding which matrices exist, how many parameters they contain, and what each one does determines which fine-tuning strategy to use.

Attention weight matrices

Each attention head computes: Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

The input hidden state x of dimension d_model is projected into queries, keys, and values through learned linear transformations:

Q = x @ W_q where W_q is (d_model, d_model) — 4096 x 4096 = 16.7M params in Llama 3 8B
K = x @ W_k where W_k is (d_model, d_kv) — 4096 x 1024 = 4.2M params (GQA: 8 KV heads)
V = x @ W_v where W_v is (d_model, d_kv) — 4096 x 1024 = 4.2M params
O = attn_output @ W_o where W_o is (d_model, d_model) — 4096 x 4096 = 16.7M params

Llama 3 8B uses Grouped Query Attention (GQA) with 32 query heads and 8 key-value heads. This reduces KV parameters by 4x compared to standard multi-head attention while maintaining quality (Ainslie et al. 2023).

Per attention block: 16.7M + 4.2M + 4.2M + 16.7M = 41.8M parameters.

Feed-forward network

Llama 3 uses a gated FFN (SwiGLU activation, Shazeer 2020):

FFN(x) = down_proj(silu(gate_proj(x)) * up_proj(x))

Three weight matrices per FFN block:

gate_proj: (d_model, d_ff) — 4096 x 14336 = 58.7M params
up_proj: (d_model, d_ff) — 4096 x 14336 = 58.7M params
down_proj: (d_ff, d_model) — 14336 x 4096 = 58.7M params

Per FFN block: 176.2M parameters. The FFN contains 4.2x more parameters than the attention block.

Layer norms and embeddings

RMSNorm layers contain a single learnable scale vector of dimension d_model = 4096 parameters each. With 2 norms per layer (pre-attention, pre-FFN) across 32 layers plus a final norm: 65 x 4096 = 266K parameters. Negligible.

The token embedding and language model head share a weight matrix of shape (vocab_size, d_model) = 128256 x 4096 = 525M parameters.

Total parameter accounting

import torch
from transformers import AutoModelForCausalLM, AutoConfig

config = AutoConfig.from_pretrained("meta-llama/Meta-Llama-3-8B")

d_model = config.hidden_size           # 4096
d_ff = config.intermediate_size        # 14336
n_layers = config.num_hidden_layers    # 32
n_heads = config.num_attention_heads   # 32
n_kv_heads = config.num_key_value_heads  # 8
vocab_size = config.vocab_size         # 128256
d_kv = d_model // n_heads * n_kv_heads # 1024

# Per layer
attn_params = d_model * d_model + d_model * d_kv * 2 + d_model * d_model  # Q, K, V, O
ffn_params = d_model * d_ff * 3  # gate, up, down
norm_params = d_model * 2  # 2 RMSNorms per layer
layer_params = attn_params + ffn_params + norm_params

# Total
embed_params = vocab_size * d_model
total_params = n_layers * layer_params + embed_params + d_model  # +final norm

print(f"Per attention block: {attn_params:,}")     # 41,943,040
print(f"Per FFN block: {ffn_params:,}")            # 176,160,768
print(f"Per layer: {layer_params:,}")              # 218,112,000
print(f"Embeddings: {embed_params:,}")             # 525,336,576
print(f"Total: {total_params:,}")                  # 8,030,261,248

VRAM budget for full fine-tuning

Training requires storing four things in memory simultaneously:

Model parameters (fp16): 8B params x 2 bytes = 16 GB
Gradients (fp16): 8B params x 2 bytes = 16 GB
Optimizer states (AdamW, fp32): 8B params x 8 bytes = 64 GB (two momentum terms + master weights)
Activations (depends on batch size and sequence length): 2-8 GB with gradient checkpointing

Total for full fine-tuning: ~60-64 GB VRAM minimum with gradient checkpointing and batch size 1. Without gradient checkpointing, activations for sequence length 2048 add another 20-40 GB.

def estimate_vram_full_finetune(
    num_params: int,
    batch_size: int = 1,
    seq_len: int = 2048,
    d_model: int = 4096,
    n_layers: int = 32,
    gradient_checkpointing: bool = True,
):
    # Parameters in fp16
    param_memory = num_params * 2  # bytes

    # Gradients in fp16
    gradient_memory = num_params * 2

    # AdamW optimizer: fp32 copy + 2 momentum buffers
    optimizer_memory = num_params * (4 + 4 + 4)  # 12 bytes per param

    # Activations (rough estimate)
    if gradient_checkpointing:
        # Only store activations at checkpoint boundaries (sqrt(n_layers))
        activation_memory = batch_size * seq_len * d_model * 2 * (n_layers ** 0.5)
    else:
        activation_memory = batch_size * seq_len * d_model * 2 * n_layers * 4

    total_bytes = param_memory + gradient_memory + optimizer_memory + activation_memory
    total_gb = total_bytes / (1024 ** 3)

    print(f"Parameters (fp16):    {param_memory / 1e9:.1f} GB")
    print(f"Gradients (fp16):     {gradient_memory / 1e9:.1f} GB")
    print(f"Optimizer (fp32):     {optimizer_memory / 1e9:.1f} GB")
    print(f"Activations:          {activation_memory / 1e9:.1f} GB")
    print(f"Total:                {total_gb:.1f} GB")
    return total_gb

estimate_vram_full_finetune(8_030_000_000)
# Parameters (fp16):    16.1 GB
# Gradients (fp16):     16.1 GB
# Optimizer (fp32):     96.4 GB   <-- this dominates
# Activations:          3.0 GB
# Total:                ~60 GB (with mixed precision optimizations reducing optimizer to ~32 GB)

In practice, frameworks like DeepSpeed ZeRO-3 shard optimizer states across GPUs. A single-GPU solution requires reducing the number of trainable parameters — which is exactly what LoRA does.

Why weight updates are low-rank

Aghajanyan et al. (2020) demonstrated that pre-trained language models have a low intrinsic dimensionality. Fine-tuning on downstream tasks produces weight changes delta_W = W_finetuned - W_pretrained that concentrate in a small subspace. Measuring the singular values of delta_W across fine-tuning runs shows rapid decay — the top 1% of singular values capture 90%+ of the variance in the weight change.

This observation has a geometric interpretation: the pre-trained model already occupies a good region of parameter space. Task adaptation only needs to move along a few directions (the task-relevant subspace), not all d x d directions available in the full weight matrix.

import torch
import numpy as np

# Simulate: measure rank of weight change after fine-tuning
# (In practice, compute delta_W from checkpoints)
d = 4096
pretrained_W = torch.randn(d, d) * 0.02  # simulated pretrained weights
finetuned_W = pretrained_W + torch.randn(d, d) * 0.001  # small perturbation

delta_W = finetuned_W - pretrained_W
U, S, V = torch.svd(delta_W)

# Cumulative explained variance
cumvar = torch.cumsum(S ** 2, dim=0) / (S ** 2).sum()
rank_90 = (cumvar < 0.90).sum().item() + 1
rank_99 = (cumvar < 0.99).sum().item() + 1

print(f"Rank to explain 90% of delta_W: {rank_90}")  # typically 16-64
print(f"Rank to explain 99% of delta_W: {rank_99}")  # typically 128-256
print(f"Full rank: {d}")                              # 4096
# The weight change is ~100x compressible

Which layers matter most

Not all layers contribute equally to task adaptation. Empirical findings (Hu et al. 2021):

Attention projection matrices (especially Q and V) capture most of the task-specific signal
FFN layers matter for knowledge injection but less for style/format tasks
Early layers (0-4) capture syntax and low-level features — rarely need updating
Middle layers (8-24) capture semantic representations — most impactful for fine-tuning
Final layers (28-31) specialize for the output distribution

LoRA's default in PEFT targets q_proj and v_proj across all layers. Extending to k_proj, o_proj, and FFN layers (gate_proj, up_proj, down_proj) increases capacity at the cost of more trainable parameters. For instruction tuning, targeting all linear layers typically outperforms attention-only by 2-5% on benchmarks (Dettmers et al. 2023).

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.float16,
)

# Inspect all named parameters and their shapes
for name, param in model.named_parameters():
    if "layers.0." in name:  # first layer as example
        print(f"{name}: {param.shape} ({param.numel():,} params)")

# Output:
# model.layers.0.self_attn.q_proj.weight: torch.Size([4096, 4096]) (16,777,216 params)
# model.layers.0.self_attn.k_proj.weight: torch.Size([1024, 4096]) (4,194,304 params)
# model.layers.0.self_attn.v_proj.weight: torch.Size([1024, 4096]) (4,194,304 params)
# model.layers.0.self_attn.o_proj.weight: torch.Size([4096, 4096]) (16,777,216 params)
# model.layers.0.mlp.gate_proj.weight: torch.Size([14336, 4096]) (58,720,256 params)
# model.layers.0.mlp.up_proj.weight: torch.Size([14336, 4096]) (58,720,256 params)
# model.layers.0.mlp.down_proj.weight: torch.Size([4096, 14336]) (58,720,256 params)
# model.layers.0.input_layernorm.weight: torch.Size([4096]) (4,096 params)
# model.layers.0.post_attention_layernorm.weight: torch.Size([4096]) (4,096 params)

This per-layer anatomy — 41.8M attention params + 176.2M FFN params = 218M per layer x 32 layers — is what makes parameter-efficient methods essential. Updating all 8B parameters requires 4x A100 80GB GPUs. Updating 0.5% of them via LoRA fits on a single consumer GPU.

← Previous