111 lines
4.4 KiB
Python
111 lines
4.4 KiB
Python
|
"""
|
|||
|
|
Gravity-2 attention for Qwen2 / VibeThinker-3B (transformers 5.x interface).
|
|||
|
|
|
|||
|
|
Replaces softmax(QKᵀ·scaling) with a physically-motivated score:
|
|||
|
|
|
|||
|
|
score(i,j) = M_h² / (||q_i − k_j||² + eps) # then standard softmax over j
|
|||
|
|
|
|||
|
|
• M_h = softplus(gravity_mass_log[h]) — one learnable mass per QUERY head (16/layer)
|
|||
|
|
• ||q_i − k_j||² = ||q||² + ||k||² − 2·q·k # GQA: K repeated 2→16 first
|
|||
|
|
• eps guards the singularity at q==k
|
|||
|
|
|
|||
|
|
Integration uses the transformers-5.x AttentionInterface dispatch (NOT a forward
|
|||
|
|
monkeypatch): we register a "gravity" attention fn + alias its mask to "eager" so
|
|||
|
|
the framework keeps building the additive causal mask, handling RoPE/cache itself.
|
|||
|
|
"""
|
|||
|
|
import math
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from transformers.models.qwen2.modeling_qwen2 import repeat_kv
|
|||
|
|
from transformers.modeling_utils import AttentionInterface
|
|||
|
|
from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
|
|||
|
|
|
|||
|
|
ATTN_NAME = "gravity"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def gravity_attention_forward(module, query, key, value, attention_mask,
|
|||
|
|
scaling=None, dropout=0.0, **kwargs):
|
|||
|
|
"""AttentionInterface contract.
|
|||
|
|
|
|||
|
|
query: (B, Hq, Tq, D) key/value: (B, Hkv, Tk, D)
|
|||
|
|
returns: (attn_output (B, Tq, Hq, D), attn_weights (B, Hq, Tq, Tk))
|
|||
|
|
`scaling` is intentionally ignored — gravity replaces the 1/√d scale.
|
|||
|
|
"""
|
|||
|
|
# GQA: expand 2 KV heads up to 16 so distances live in per-query-head space
|
|||
|
|
key = repeat_kv(key, module.num_key_value_groups)
|
|||
|
|
value = repeat_kv(value, module.num_key_value_groups)
|
|||
|
|
|
|||
|
|
# ||q_i - k_j||^2 in fp32 for numerical stability
|
|||
|
|
q = query.float()
|
|||
|
|
k = key.float()
|
|||
|
|
q_sq = (q * q).sum(-1, keepdim=True) # (B,Hq,Tq,1)
|
|||
|
|
k_sq = (k * k).sum(-1, keepdim=True).transpose(-2, -1) # (B,Hq,1,Tk)
|
|||
|
|
qk = torch.matmul(q, k.transpose(-2, -1)) # (B,Hq,Tq,Tk)
|
|||
|
|
d_sq = (q_sq + k_sq - 2.0 * qk).clamp_min(0.0)
|
|||
|
|
|
|||
|
|
mass = F.softplus(module.gravity_mass_log).float().view(1, -1, 1, 1) # (1,Hq,1,1)
|
|||
|
|
scores = (mass * mass) / (d_sq + module.gravity_eps) # (B,Hq,Tq,Tk), fp32
|
|||
|
|
|
|||
|
|
if attention_mask is not None:
|
|||
|
|
# additive causal mask (eager-style), already correct length
|
|||
|
|
scores = scores + attention_mask[..., : key.shape[-2]].float()
|
|||
|
|
|
|||
|
|
attn = F.softmax(scores, dim=-1, dtype=torch.float32)
|
|||
|
|
|
|||
|
|
# AER: optionally stash mean per-row attention entropy (flag-gated, ~free when off)
|
|||
|
|
if getattr(module, "_capture_entropy", False):
|
|||
|
|
ent = -(attn.clamp_min(1e-12) * attn.clamp_min(1e-12).log()).sum(-1)
|
|||
|
|
module._last_entropy = ent.mean().detach()
|
|||
|
|
|
|||
|
|
attn = F.dropout(attn, p=dropout, training=module.training)
|
|||
|
|
attn = attn.to(value.dtype)
|
|||
|
|
|
|||
|
|
out = torch.matmul(attn, value) # (B,Hq,Tq,D)
|
|||
|
|
out = out.transpose(1, 2).contiguous() # (B,Tq,Hq,D)
|
|||
|
|
return out, attn
|
|||
|
|
|
|||
|
|
|
|||
|
|
_REGISTERED = False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _register():
|
|||
|
|
global _REGISTERED
|
|||
|
|
if _REGISTERED:
|
|||
|
|
return
|
|||
|
|
AttentionInterface.register(ATTN_NAME, gravity_attention_forward)
|
|||
|
|
# reuse the eager additive-mask builder for our custom impl
|
|||
|
|
ALL_MASK_ATTENTION_FUNCTIONS.register(ATTN_NAME, ALL_MASK_ATTENTION_FUNCTIONS["eager"])
|
|||
|
|
_REGISTERED = True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def patch_qwen_with_gravity(model, eps: float = 0.1, init_mass: float = 0.5):
|
|||
|
|
"""Add per-head gravity_mass_log to every Qwen2 self-attn and switch dispatch.
|
|||
|
|
|
|||
|
|
Leaves q/k/v/o_proj weights untouched. gravity_mass_log kept in fp32.
|
|||
|
|
"""
|
|||
|
|
_register()
|
|||
|
|
init_log = math.log(math.exp(init_mass) - 1.0) # softplus^{-1}(init_mass)
|
|||
|
|
H = model.config.num_attention_heads
|
|||
|
|
n = 0
|
|||
|
|
for layer in model.model.layers:
|
|||
|
|
attn = layer.self_attn
|
|||
|
|
dev = attn.q_proj.weight.device
|
|||
|
|
attn.gravity_mass_log = nn.Parameter(
|
|||
|
|
torch.full((H,), init_log, device=dev, dtype=torch.float32)
|
|||
|
|
)
|
|||
|
|
attn.gravity_eps = float(eps)
|
|||
|
|
# config object is shared, but set defensively
|
|||
|
|
attn.config._attn_implementation = ATTN_NAME
|
|||
|
|
n += 1
|
|||
|
|
model.config._attn_implementation = ATTN_NAME
|
|||
|
|
print(f"[gravity] patched {n} Qwen2 layers (heads={H}, eps={eps}, init_mass={init_mass})")
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def gravity_mass_state_dict(model):
|
|||
|
|
"""Extract only the gravity_mass_log params (for saving separately from base)."""
|
|||
|
|
return {f"model.layers.{i}.self_attn.gravity_mass_log":
|
|||
|
|
layer.self_attn.gravity_mass_log.detach().cpu()
|
|||
|
|
for i, layer in enumerate(model.model.layers)}
|