File size: 1,043 Bytes
a20151e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import torch
from transformers import GPT2Model
# Load pretrained GPT-2
model = GPT2Model.from_pretrained("openai-community/gpt2")
ratios = []
# Loop through transformer blocks
for i, block in enumerate(model.h): # model.h = list of transformer layers
# GPT2 uses Conv1D for QKV, packed into one matrix
W_qkv = block.attn.c_attn.weight.detach() # shape [768, 2304]
# Split into Q, K, V
hidden_size = model.config.hidden_size # 768
W_q, W_k, W_v = W_qkv.split(hidden_size, dim=1) # each [768, 768]
# Compute norms
l2_q = torch.norm(W_q, p=2)
l2_k = torch.norm(W_k, p=2)
ratio = l2_q / l2_k
ratios.append(ratio.item())
print(f"Layer {i+1}: Q/K L2 ratio = {ratio.item():.4f}")
# Convert to tensor for stats
ratios_tensor = torch.tensor(ratios)
mean_ratio = torch.mean(ratios_tensor).item()
std_ratio = torch.std(ratios_tensor).item()
print(f"\nMean Q/K ratio over {len(ratios)} layers: {mean_ratio:.4f}")
print(f"Std Q/K ratio over {len(ratios)} layers: {std_ratio:.4f}")
|