xFormers developer on retainer: query tensor pre-scaled by 1/√head_dim before memory_efficient_attention() — xFormers also applies 1/√head_dim internally — effective scale 1/head_dim — attention weights uniform over all positions — GPT-NeoX-20B loss plateaus at perplexity 500 — no convergence — 3 weeks — xFormers developer on monthly retainer
xformers 0.0.26 • Python 3.11 • NVIDIA A100 40 GB • GPT-NeoX-20B fine-tuning • PyTorch 2.3.x • q = q / math.sqrt(head_dim) preserved from manual SDPA • memory_efficient_attention() applies scale=1/sqrt(96) internally • effective scale 1/96 • attention logits 9.8× too small • softmax output near-uniform • entropy 7.6 nats per head • loss plateaus at perplexity 500 for 10,000 steps • developer checks learning rate, weight decay, gradient clipping — 3 weeks • fix: remove pre-scaling line • loss 6.2→2.4 in 500 steps • perplexity 500→11 • 1.5h retainer engagement
Northfield xFormers Efficiency Ltd. was fine-tuning GPT-NeoX-20B on a 50-million-token domain corpus. Their existing training pipeline used a custom PyTorch attention implementation with manual scaled dot-product attention. The developer replaced the manual SDPA with xformers.ops.memory_efficient_attention() to reduce peak GPU memory and enable longer sequence lengths on the A100 40 GB. The migration appeared straightforward: swap the torch.matmul-based attention block for a single xformers call. The first training run after migration produced the correct output shapes and no CUDA errors. The developer began a full fine-tuning run on the 50M-token corpus.
After 200 training steps, the loss stopped decreasing. The loss had dropped from 6.2 to 5.8 in the first 100 steps, then held at 5.8–5.9 for the next 10,000 steps regardless of learning rate adjustments. Gradient norms were healthy (1.2–1.8), indicating the optimizer was taking steps. No NaN or Inf values appeared anywhere in the computation. The developer verified the data pipeline, confirmed the training corpus was not degenerate, and ran a sanity-check with the original PyTorch SDPA implementation: loss decreased normally. The developer attributed the plateau to a hyperparameter mismatch and spent three weeks adjusting learning rate from 1e-5 to 3e-4, weight decay from 0.0 to 0.1, gradient clipping from 0.5 to 1.0, and warmup steps from 100 to 1,000. None resolved the plateau. After three weeks the developer engaged an xFormers retainer developer.
The attention implementation at the time of the engagement:
# Custom attention — Northfield xFormers Efficiency Ltd.
# xformers 0.0.26, PyTorch 2.3.x, Python 3.11
# GPT-NeoX-20B: num_layers=44, num_heads=64, head_dim=96, hidden_size=6144
import math
import xformers.ops as xops
from xformers.ops import LowerTriangularMask
class GPTNeoXAttention(nn.Module):
def __init__(self, config):
super().__init__()
self.num_heads = config.num_attention_heads # 64
self.head_dim = config.hidden_size // config.num_attention_heads # 96
self.qkv_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size)
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
def forward(self, hidden_states, attention_mask=None):
batch, seq_len, hidden_size = hidden_states.shape
# QKV projection
qkv = self.qkv_proj(hidden_states) # (batch, seq_len, 3*hidden_size)
q, k, v = qkv.chunk(3, dim=-1) # each: (batch, seq_len, hidden_size)
# Reshape to multi-head
q = q.view(batch, seq_len, self.num_heads, self.head_dim) # (B, S, H, D)
k = k.view(batch, seq_len, self.num_heads, self.head_dim)
v = v.view(batch, seq_len, self.num_heads, self.head_dim)
# BUG: pre-scaling preserved from the original manual SDPA implementation
# Original code (before xformers migration):
# q = q / math.sqrt(self.head_dim)
# attn = torch.matmul(q, k.transpose(-2,-1)) # already scaled here
# attn = torch.softmax(attn, dim=-1)
# out = torch.matmul(attn, v)
#
# After xformers migration — developer kept the pre-scaling:
q = q / math.sqrt(self.head_dim) # ← q divided by sqrt(96) = 9.798
# math.sqrt(96) = 9.798
# q values now ≈ 9.8× smaller than v and k
# xformers.ops.memory_efficient_attention() default: scale=None
# When scale=None, xformers internally computes scale = 1 / math.sqrt(q.shape[-1])
# q.shape[-1] = head_dim = 96
# xformers internal scale = 1 / sqrt(96) = 1 / 9.798 = 0.1021
#
# Effective attention scale = pre-scale × xformers-internal-scale
# = (1/9.798) × (1/9.798) = 1/96 = 0.01042
# Correct scale: 1/9.798 = 0.1021
# Actual scale: 0.01042 (9.8× too small)
#
# Attention logits Q_scaled · K^T:
# correct: Q · K^T / 9.798 → logits in range [-3, 3] for typical Q, K
# actual: (Q/9.798) · K^T / 9.798 = Q · K^T / 96 → logits in [-0.3, 0.3]
#
# softmax of logits in [-0.3, 0.3]: near-uniform distribution over all positions
# entropy ≈ log(2048) = 7.6 nats per head (maximum = completely uniform attention)
out = xops.memory_efficient_attention(
q, k, v,
attn_bias=LowerTriangularMask(), # causal mask
# scale=None ← default: xformers applies 1/sqrt(q.shape[-1]) internally
)
# out shape: (batch, seq_len, num_heads, head_dim)
# Reshape back
out = out.reshape(batch, seq_len, hidden_size)
return self.out_proj(out)Why double scaling makes attention uniform
xFormers memory_efficient_attention() computes softmax(Q·K^T × scale) × V where scale = 1/sqrt(q.shape[-1]) when scale=None (the default). For head_dim=96, scale = 1/sqrt(96) ≈ 0.102. The developer’s pre-scaling divided q by sqrt(96) before passing it to xformers, so the actual dot product computed is (Q/9.798) · K^T × 0.102 = Q · K^T / 96. The correct scale factor is Q · K^T / 9.798.
For a typical attention head, Q · K^T entries are in the range [-30, 30]. With the correct scale of 1/9.798, logits span [-3, 3]; softmax([-3, 3]) produces peaked distributions concentrating weight on a few tokens. With the double scale of 1/96, logits span [-0.3, 0.3]; softmax([-0.3, 0.3]) over 2,048 positions produces weights that differ by less than 0.02% across all positions—effectively uniform. Each attention head produces a weighted average over all 2,048 context positions with nearly equal weight, destroying all positional selectivity. The model is equivalent to a mean-pooling layer followed by a linear projection.
# Diagnostic: check attention weight entropy during training
# Healthy attention should have low entropy (peaked on a few tokens)
# Double-scaled attention has high entropy (uniform over all positions)
import torch
import math
def attention_entropy_check(model, sample_batch):
"""
Temporarily hook attention weights to measure entropy.
Healthy training: entropy should decrease from log(seq_len) toward 1-3 nats.
Double scaling: entropy stays at log(seq_len) for all training steps.
"""
entropies = []
def hook_fn(module, input, output):
# Compute attention entropy from q, k if accessible
# Alternatively: measure indirectly via output diversity
pass
# Quick diagnostic: check q tensor magnitude before xformers call
q_before_xformers = ... # capture via hook or print in forward()
expected_q_magnitude = 1.0 # unscaled q has magnitude ~1.0 (depends on init)
prescaled_q_magnitude = 1.0 / math.sqrt(96) # ≈ 0.102 — too small
print(f"q.abs().mean() = {q_before_xformers.abs().mean():.4f}")
# If this prints ≈0.010-0.020, q was pre-scaled AND xformers will scale again
# If this prints ≈0.100-0.200, q is pre-scaled correctly (use scale=1.0 in xformers)
# If this prints ≈0.900-1.100, q is unscaled (let xformers use default scale=None)
# Explicit entropy measurement (insert into forward during debug run):
# q: (batch, seq_len, num_heads, head_dim) — xformers layout
# k: (batch, seq_len, num_heads, head_dim)
with torch.no_grad():
# Compute raw attention logits for one head to check scale
q_head = q[:, :, 0, :] # (batch, seq_len, head_dim)
k_head = k[:, :, 0, :] # (batch, seq_len, head_dim)
logits = torch.matmul(q_head, k_head.transpose(-2, -1)) # (batch, seq_len, seq_len)
print(f"Attention logit range: [{logits.min():.3f}, {logits.max():.3f}]")
# Healthy: [-3.0, 3.0] (xformers will scale correctly)
# Double-scaled: [-0.3, 0.3] (logits already divided by head_dim)
weights = torch.softmax(logits, dim=-1) # (batch, seq_len, seq_len)
entropy = -(weights * (weights + 1e-9).log()).sum(-1) # (batch, seq_len)
print(f"Mean attention entropy: {entropy.mean():.2f} nats")
print(f"log(seq_len={q.shape[1]}): {math.log(q.shape[1]):.2f} nats")
# Double scaling: 7.60 ≈ log(2048) — uniform
# Healthy (step 0): 6.5-7.0 — slightly peaked
# Healthy (trained): 1.5-3.0 — clearly peaked
return entropies
# Developer's configuration (buggy — q pre-scaled, xformers scales again):
# q.abs().mean() ≈ 0.012 (q/9.798 → values ≈ 9.8× too small)
# Attention logit range: [-0.31, 0.29]
# Mean attention entropy: 7.59 nats ≈ log(2048) = uniform attention ✗
# Training loss after 10,000 steps: 5.8 (plateau)The fix: remove the pre-scaling line
The fix is to remove the q = q / math.sqrt(self.head_dim) line from the forward pass and let memory_efficient_attention() apply the correct scale internally via its scale=None default. No other changes are needed.
# Fixed attention — Northfield xFormers Efficiency Ltd.
# xformers 0.0.26, PyTorch 2.3.x
class GPTNeoXAttention(nn.Module):
def forward(self, hidden_states, attention_mask=None):
batch, seq_len, hidden_size = hidden_states.shape
qkv = self.qkv_proj(hidden_states)
q, k, v = qkv.chunk(3, dim=-1)
q = q.view(batch, seq_len, self.num_heads, self.head_dim)
k = k.view(batch, seq_len, self.num_heads, self.head_dim)
v = v.view(batch, seq_len, self.num_heads, self.head_dim)
# FIXED: pre-scaling line removed
# q = q / math.sqrt(self.head_dim) ← DELETED
# xformers default scale=None applies 1/sqrt(q.shape[-1]) = 1/sqrt(96) = 0.102
# q is now unscaled: values ≈ 0.9-1.1
# Effective scale: 0.102 (correct)
# Attention logits Q·K^T / 9.798: range [-3, 3] for typical Q, K
# Softmax of [-3, 3] over 2048 positions: peaked attention ✓
out = xops.memory_efficient_attention(
q, k, v,
attn_bias=LowerTriangularMask(),
# scale=None ← default; xformers applies 1/sqrt(head_dim) correctly
)
out = out.reshape(batch, seq_len, hidden_size)
return self.out_proj(out)
# Alternative: keep pre-scaling, suppress xformers internal scaling
# (use this if you need pre-scaling for other reasons, e.g., RoPE applied to scaled q)
class GPTNeoXAttentionAlt(nn.Module):
def forward(self, hidden_states, attention_mask=None):
batch, seq_len, hidden_size = hidden_states.shape
qkv = self.qkv_proj(hidden_states)
q, k, v = qkv.chunk(3, dim=-1)
q = q.view(batch, seq_len, self.num_heads, self.head_dim)
k = k.view(batch, seq_len, self.num_heads, self.head_dim)
v = v.view(batch, seq_len, self.num_heads, self.head_dim)
q = q / math.sqrt(self.head_dim) # manual pre-scaling
# Set scale=1.0 to suppress xformers internal scaling:
out = xops.memory_efficient_attention(
q, k, v,
attn_bias=LowerTriangularMask(),
scale=1.0, # ← suppress internal scaling since q is already pre-scaled
)
# Effective scale: pre-scale (1/9.798) × xformers-scale (1.0) = 1/9.798 ✓
out = out.reshape(batch, seq_len, hidden_size)
return self.out_proj(out)
# Numerical correctness test (run before training to verify fix):
import torch
torch.manual_seed(42)
q_test = torch.randn(2, 64, 4, 96, device='cuda', dtype=torch.float16)
k_test = torch.randn(2, 64, 4, 96, device='cuda', dtype=torch.float16)
v_test = torch.randn(2, 64, 4, 96, device='cuda', dtype=torch.float16)
# xformers result:
out_xformers = xops.memory_efficient_attention(q_test, k_test, v_test)
# Reference PyTorch SDPA (batch, heads, seq, dim) — transpose for PyTorch:
q_pt = q_test.transpose(1, 2) # (2, 4, 64, 96)
k_pt = k_test.transpose(1, 2)
v_pt = v_test.transpose(1, 2)
out_pt = torch.nn.functional.scaled_dot_product_attention(q_pt, k_pt, v_pt)
out_pt = out_pt.transpose(1, 2) # back to (2, 64, 4, 96)
max_diff = (out_xformers - out_pt).abs().max().item()
print(f"Max diff xformers vs PyTorch SDPA: {max_diff:.4f}")
# Should be < 0.005 (fp16 rounding); if > 0.1, scaling mismatch still presentValidation
# Validation log — November 2026
# xformers 0.0.26, GPT-NeoX-20B, NVIDIA A100 40GB
# Fine-tuning on 50M-token domain corpus, batch_size=4, seq_len=2048
# ---- Numerical correctness test ----
# Max diff xformers vs PyTorch SDPA: 0.0021 ← within fp16 rounding tolerance ✓
# ---- Before fix (q pre-scaled, xformers scales again) ----
# Step 0: loss 6.19 | perplexity 490 | attn_entropy 7.60 nats (= log(2048))
# Step 100: loss 5.84 | perplexity 343 | attn_entropy 7.57 nats
# Step 500: loss 5.81 | perplexity 334 | attn_entropy 7.55 nats
# Step 1000: loss 5.80 | perplexity 330 | attn_entropy 7.53 nats (plateau)
# Step 5000: loss 5.79 | perplexity 327 | attn_entropy 7.52 nats (plateau)
# → No convergence. Loss plateau. Attention entropy = maximum (uniform weights).
# ---- After fix (pre-scaling removed, xformers handles scale) ----
# Step 0: loss 6.21 | perplexity 498 | attn_entropy 7.60 nats
# Step 100: loss 4.72 | perplexity 112 | attn_entropy 5.14 nats ← decreasing ✓
# Step 500: loss 2.42 | perplexity 11 | attn_entropy 2.78 nats ← converging ✓
# Step 1000: loss 1.98 | perplexity 7 | attn_entropy 2.11 nats ✓
# Step 5000: loss 1.41 | perplexity 4 | attn_entropy 1.63 nats ✓
# → Loss converges normally. Attention entropy drops as model learns to attend selectively.
# ---- Numerical correctness after fix ----
# Max diff xformers vs PyTorch SDPA: 0.0019 ✓
# ---- Memory comparison ----
# Manual PyTorch SDPA (before xformers): 38.4 GB VRAM, seq_len=2048, batch=4
# xformers (after fix): 22.1 GB VRAM, seq_len=2048, batch=4
# Memory saving: 42% ✓ (original motivation for xformers migration)
# ---- Hour breakdown ----
# Triage: 0.5h — loss plateau at step 200; gradient norms 1.2-1.8 (healthy);
# no NaN/Inf; printed q.abs().mean()=0.012 before xformers call
# (expected ≈0.1 for unscaled q with head_dim=96 initialization);
# attention entropy 7.6 nats = log(2048) = uniform; q pre-scaling line
# found at line 47 of attention.py from original SDPA migration
# Root cause: 0.25h — q/sqrt(96) × xformers(1/sqrt(96)) = 1/96; effective scale wrong
# Fix: 0.25h — removed pre-scaling line; numerical test passes
# Validation: 0.5h — 500 training steps; loss 6.2→2.4; entropy 7.6→2.8; converging
# Total: 1.5hRelated retainer posts
- FlashInfer developer on retainer — paged KV cache allocated for single-sequence capacity — batch sequences 1–7 read out-of-bounds pages — garbage attention output — no CUDA error — fix: allocate max_pages_per_seq×batch_size — 3 weeks
- TensorRT-LLM developer on retainer — engine built with
--max_beam_width 1— runtimenum_beams=4— KV cache overflow between batch sequences — quality degraded positions 3–8 — no error — 3 weeks - vLLM developer on retainer — PagedAttention KV cache configuration — tensor parallelism deadlocks — AWQ quantization compatibility — throughput issues — 3 weeks
Track retainer hours without a client portal
HourTab turns your time-tracker CSV into a public URL your client bookmarks. No login. No portal. CSV in, URL out.
See pricing →