import importlib.util, math, os, sys, types, warnings, torch, torch.nn as nn, torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
def _text_only_if_torchvision_broken():
"""A torchvision built for a different torch ("operator torchvision::nms does not exist") crashes transformers' lazy
vision imports, which run when PreTrainedModel is first imported. Only if torchvision fails to import: drop its
half-imported modules and tell transformers it is unavailable (this model is text-only). A working install is untouched."""
try: import torchvision # noqa: F401
except Exception:
for name in [m for m in sys.modules if m == "torchvision" or m.startswith("torchvision.")]: del sys.modules[name]
import transformers.utils as utils
from transformers.utils import import_utils
if hasattr(import_utils, "_torchvision_available"): import_utils._torchvision_available = False
for name in ("is_torchvision_available", "is_torchvision_v2_available"):
old = getattr(import_utils, name, None)
if callable(getattr(old, "cache_clear", None)): old.cache_clear()
setattr(import_utils, name, lambda: False); setattr(utils, name, lambda: False)
_text_only_if_torchvision_broken()
from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
def rope(x, cos, sin): a, b = x.chunk(2, -1); return torch.cat((a * cos - b * sin, a * sin + b * cos), -1)
def causal_conv(x, w, docs): # depthwise causal conv over time, w (K, C): y[t] = sum_j w[j] * x[t - K + 1 + j]; shifted sums fuse
K, T, w = w.shape[0], x.shape[1], w.to(x.dtype); x, d = F.pad(x, (0, 0, K - 1, 0)), F.pad(docs, (K - 1, 0), value=-1)
return sum(x[:, j:j + T] * w[j] * (d[:, j:j + T] == docs).unsqueeze(-1).to(x.dtype) for j in range(K)) # no taps across documents
class RMSNorm(nn.RMSNorm): # weight cast to the input dtype so bf16 activations take the fused kernel
def forward(s, x): return F.rms_norm(x, s.normalized_shape, s.weight.to(x.dtype), s.eps)
def block_masks(c, docs): # per attention kind: causal (+ window) and inside one document; docs (B, T) = document id per token
(B, T), w = docs.shape, c.window
swa = lambda b, h, q, kv: (q >= kv) & (q - kv < w) & (docs[b, q] == docs[b, kv])
full = lambda b, h, q, kv: (q >= kv) & (docs[b, q] == docs[b, kv])
# inside the compiled model the mask build is traced and fused; eager callers (HF loading, eval) compile it on their own
# when the kernels can compile here, and build it uncompiled otherwise (same mask)
def make(f):
if not torch.compiler.is_compiling() and _flex.compiled is not None:
try: return create_block_mask(f, B, None, T, T, device=docs.device, _compile=True)
except Exception as e: _flex.give_up(e)
return create_block_mask(f, B, None, T, T, device=docs.device, _compile=False)
return {"swa": make(swa), "full": make(full)}
def _can_compile(): # compiled flex kernels need triton and a CUDA GPU of compute capability 8.0+ (Ampere or newer)
if os.environ.get("SWA_NO_COMPILE"): return False # set SWA_NO_COMPILE=1 to force the eager path
try: return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 0) and importlib.util.find_spec("triton") is not None
except Exception: return False
class SafeFlex:
"""flex_attention compiled where it can be (fused, memory-light). On CPU, pre-Ampere GPUs (T4, V100), a missing
triton, or the first compile failure, `compiled` is None and the model uses chunked_attention instead (the call
that hit the failure finishes on eager flex_attention)."""
def __init__(s): s.compiled = torch.compile(flex_attention) if _can_compile() else None
def give_up(s, e):
warnings.warn(f"compiled flex_attention unavailable ({type(e).__name__}: {str(e).splitlines()[0][:200] if str(e) else ''}); "
"using eager attention (same results, slower)"); s.compiled = None
def __call__(s, *a, **kw):
if s.compiled is not None:
try: return s.compiled(*a, **kw)
except Exception as e: s.give_up(e)
return flex_attention(*a, **kw)
# Compiled flex_attention caches on the mask functions' code objects, and transformers imports every exported folder as
# a new copy of this file. The first copy in a process registers the masks and the compiled kernel; later copies (export
# checks, eval loads) reuse them, so the kernel compiles once instead of once per load (8 loads -> unfused fallback).
# v2: document masks (block_masks takes document ids); v3: SafeFlex (eager fallback), so a process holding older
# exports never mixes them.
_shared = sys.modules.setdefault("swa_lm_kernels_v3", types.ModuleType("swa_lm_kernels_v3"))
if not hasattr(_shared, "flex"): _shared.flex, _shared.block_masks = SafeFlex(), block_masks
_flex, block_masks = _shared.flex, _shared.block_masks # eager callers (HF loading, eval) still get the fused kernel
def cached_attention(q, k, v, qpos, kpos, qdoc, kdoc, window, sink):
"""Attention of new queries over cached + new keys (decode steps), the same function as the flex path: causal,
same document, optional sliding window, GQA, and a per-head sink slot with zero value.
q (B, T, H, D); k, v (B, L, KV, D); qpos (T,), kpos (L,) absolute positions; qdoc (B, T), kdoc (B, L)."""
B, T, H, D = q.shape; g = H // k.shape[2]
k, v = k.repeat_interleave(g, 2), v.repeat_interleave(g, 2) # query head i uses kv head i // g, as enable_gqa
sc = torch.einsum("bthd,blhd->bhtl", q.float(), k.float()) / math.sqrt(D)
ok = (kpos[None, None, :] <= qpos[None, :, None]) & (kdoc[:, None, :] == qdoc[:, :, None])
if window: ok = ok & (qpos[None, :, None] - kpos[None, None, :] < window)
sc = sc.masked_fill(~ok[:, None], float("-inf"))
sc = torch.cat([sc, sink.float().view(1, H, 1, 1).expand(B, H, T, 1)], -1) # sink logit, value 0
return torch.einsum("bhtl,blhd->bthd", sc.softmax(-1)[..., :-1], v.float()).to(q.dtype)
def chunked_attention(q, k, v, pos, docs, window, sink, chunk=512):
"""The flex path without compiling: queries in chunks, each over only the keys it can see (the last `window`
positions for sliding-window layers, everything before it for full layers), so no T x T score matrix is built.
q (B, T, H, D); k, v (B, T, KV, D); pos (T,); docs (B, T). Same result as flex with the block masks."""
out = []
for i in range(0, q.shape[1], chunk):
j = min(i + chunk, q.shape[1]); a = max(0, i - window + 1) if window else 0
out.append(cached_attention(q[:, i:j], k[:, a:j], v[:, a:j], pos[i:j], pos[a:j], docs[:, i:j], docs[:, a:j], window, sink))
return torch.cat(out, 1)
class SWACache:
"""KV cache for generation. Per layer: keys/values after norms, HoPE and value residual (sliding-window layers keep
only the last `window` tokens, full layers keep all), their positions and document ids, and the last conv - 1 raw
q|k|v rows so the causal conv sees the right neighbours for new tokens. Inference only; training never builds one."""
def __init__(s, layers): s.layers, s.length, s.last_doc, s.ids = [dict() for _ in range(layers)], 0, None, None
def get_seq_length(s, layer_idx=0): return s.length
def __len__(s): return len(s.layers)
class Block(nn.Module):
def __init__(s, c, i):
super().__init__(); lin = lambda i, o: nn.Linear(i, o, bias=False); s.c, s.kind = c, "full" if i in c.full else "swa"
s.wqkv, s.wo = lin(c.d, (c.h + 2 * c.kv) * c.hd), lin(c.h * c.hd, c.d) # fused q|k|v: Muon orthogonalizes it as one matrix
s.w13, s.w2 = lin(c.d, 2 * c.m), lin(c.m, c.d) # fused SwiGLU gate|up, then down
s.n1, s.n2, s.qn, s.kn = RMSNorm(c.d), RMSNorm(c.d), RMSNorm(c.hd), RMSNorm(c.hd)
s.sink = nn.Parameter(torch.zeros(c.h)) # per-head logit of a zero-value "attend to nothing" slot
s.conv = nn.Parameter(torch.zeros(c.conv, (c.h + 2 * c.kv) * c.hd)) # residual short conv on q|k|v (Primer / Canon); 0 = identity at init
if i: s.lam = nn.Parameter(torch.full((2,), .5)) # value residual (ResFormer): v = lam0 * v + lam1 * v of block 0
def forward(s, x, cos, sin, masks, docs, v1=None, cache=None, pos=None): # pre-norm block; returns (x, block-0 values)
c, (B, T, _), h = s.c, x.shape, s.n1(x)
raw = s.wqkv(h)
if cache is not None and "craw" in cache: # decode: conv over cached rows + new rows
full, fdoc = torch.cat([cache["craw"], raw], 1), torch.cat([cache["cdoc"], docs], 1)
qkv = raw + causal_conv(full, s.conv, fdoc)[:, -T:]
else:
full, fdoc = raw, docs; qkv = raw + causal_conv(raw, s.conv, docs) # each q/k/v also sees up to conv - 1 earlier tokens of its document
if cache is not None: n = max(0, full.shape[1] - (c.conv - 1)); cache["craw"], cache["cdoc"] = full[:, n:], fdoc[:, n:]
q, k, v = qkv.split((c.h * c.hd, c.kv * c.hd, c.kv * c.hd), -1)
q, k, v = q.view(B, T, c.h, c.hd), k.view(B, T, c.kv, c.hd), v.view(B, T, c.kv, c.hd)
if v1 is None: v1 = v # block 0 hands its values to every later block
else: lam = s.lam.to(v.dtype); v = lam[0] * v + lam[1] * v1
q, k = rope(s.qn(q), cos, sin).type_as(v), rope(s.kn(k), cos, sin).type_as(v)
if (cache is None or "k" not in cache) and masks is not None and (torch.compiler.is_compiling() or _flex.compiled is not None):
flex = flex_attention if torch.compiler.is_compiling() else _flex # training, eval, and the prompt pass (fused)
a, lse = flex(*(t.transpose(1, 2) for t in (q, k, v)), block_mask=masks[s.kind], enable_gqa=True, return_lse=True)
a = a * torch.sigmoid(lse - s.sink.view(1, -1, 1)).unsqueeze(-1).type_as(a) # softmax with sink: Z / (Z + e^sink)
o = a.transpose(1, 2).reshape(B, T, c.h * c.hd)
K, V, P, D = k, v, pos, docs
elif cache is None or "k" not in cache: # the same, where flex cannot compile
P = pos if pos is not None else torch.arange(T, device=x.device)
o = chunked_attention(q, k, v, P, docs, c.window if s.kind == "swa" else 0, s.sink).reshape(B, T, c.h * c.hd)
K, V, D = k, v, docs
else: # decode: new queries over cached + new keys
K, V = torch.cat([cache["k"], k], 1), torch.cat([cache["v"], v], 1)
P, D = torch.cat([cache["pos"], pos]), torch.cat([cache["doc"], docs], 1)
o = cached_attention(q, K, V, pos, P, docs, D, c.window if s.kind == "swa" else 0, s.sink).reshape(B, T, c.h * c.hd)
if cache is not None:
n = c.window if s.kind == "swa" else K.shape[1] # sliding-window layers never need older keys
cache["k"], cache["v"], cache["pos"], cache["doc"] = K[:, -n:], V[:, -n:], P[-n:], D[:, -n:]
x = x + s.wo(o); g, u = s.w13(s.n2(x)).chunk(2, -1)
return x + s.w2(F.silu(g) * u), v1
class SWA(nn.Module):
def __init__(s, c):
super().__init__(); s.c = c
s.emb, s.blocks, s.norm = nn.Embedding(c.v, c.d), nn.ModuleList(Block(c, i) for i in range(c.l)), RMSNorm(c.d)
for p in s.parameters(): # GPT-2 / Megatron init: N(0, 0.02), residual output projections scaled by 1/sqrt(2 * layers)
if p.ndim == 2: nn.init.normal_(p, std=.02)
for b in s.blocks:
for w in (b.wo.weight, b.w2.weight): nn.init.normal_(w, std=.02 / math.sqrt(2 * c.l))
nn.init.zeros_(b.conv) # 2-D, so the loop above drew it; starts as identity
def forward(s, x, docs=None, cache=None): # packed doc rows: each token attends only within its own document
# docs (B, T): segment id per token. Training passes them (they also break at mixture-source boundaries, where a
# slice starts mid-document without ); every other caller gets ids from the positions.
# cache: an SWACache to read and extend (generation); None everywhere else.
c, T, ids = s.c, x.shape[1], x; off = cache.length if cache is not None else 0
if docs is None:
docs = (x == c.bos_token_id).cumsum(1)
if off: docs = docs + cache.last_doc[:, None] # continue the cached document ids
fused = not off and (torch.compiler.is_compiling() or _flex.compiled is not None and x.is_cuda) # compiled kernels: CUDA inputs only
masks = block_masks(c, docs) if fused else None # decode steps and uncompiled runs use cached_attention
x = s.emb(x)
if torch.is_autocast_enabled(): x = x.to(torch.get_autocast_dtype("cuda"))
inv = 1 / c.theta ** (torch.arange(0, c.hd, 2, device=x.device) / c.hd) # rotary frequency per q/k dimension pair
if c.hope: inv = inv * (inv >= 2 * math.pi / c.seq) # HoPE: pairs that turn less than once per training length stay unrotated
pos = torch.arange(off, off + T, device=x.device)
t = torch.outer(pos.float(), inv)
cos, sin = t.cos()[:, None], t.sin()[:, None]
v1 = None
for i, b in enumerate(s.blocks):
if cache is None: x, v1 = b(x, cos, sin, masks, docs, v1)
else: x, v1 = b(x, cos, sin, masks, docs, v1, cache.layers[i], pos)
if cache is not None:
cache.length += T; cache.last_doc = docs[:, -1]
cache.ids = ids if cache.ids is None else torch.cat([cache.ids, ids], 1)
return s.norm(x)
class SWAConfig(PretrainedConfig):
model_type = "swa_lm"
def __init__(s, d=576, h=9, kv=3, hd=64, m=1792, l=20, v=32000, seq=16384, window=512, full=(6, 14), theta=500_000., conv=4, hope=False, **kw):
s.d, s.h, s.kv, s.hd, s.m, s.l, s.v, s.seq, s.window, s.full, s.theta, s.conv = d, h, kv, hd, m, l, v, seq, window, list(full), theta, conv
s.hope = hope # HoPE: freeze rotary pairs slower than one turn per `seq` tokens
s.vocab_size, s.hidden_size, s.num_hidden_layers, s.max_position_embeddings = v, d, l, seq
for k, val in dict(bos_token_id=1, eos_token_id=2, pad_token_id=2, use_cache=False, tie_word_embeddings=False).items(): kw.setdefault(k, val)
super().__init__(**kw)
class SWAForCausalLM(PreTrainedModel, GenerationMixin):
config_class, base_model_prefix, _no_split_modules = SWAConfig, "model", ["Block"]
def __init__(s, config, *args, **kw): # config overrides passed by loaders (e.g. use_cache) go onto the config
for k, v in kw.items(): setattr(config, k, v)
super().__init__(config); s.model = SWA(config); s.post_init()
def _init_weights(s, m): pass # SWA initializes itself; checkpoints overwrite
def get_input_embeddings(s): return s.model.emb
def set_input_embeddings(s, e): s.model.emb = e
def prepare_inputs_for_generation(s, input_ids, past_key_values=None, attention_mask=None, **kwargs):
"""generate(): keep one SWACache per sequence and feed only the tokens it has not seen. The cache is found in
past_key_values, or (if this transformers version keeps its own cache object there) in the last cache this
model built, reused only when its tokens are exactly the prefix of input_ids. Batch rows must be unpadded."""
cache = past_key_values if isinstance(past_key_values, SWACache) else s.__dict__.get("_swa_last")
L = cache.length if cache is not None else 0
if not (cache is not None and 0 < L < input_ids.shape[1] and cache.ids.shape[0] == input_ids.shape[0]
and torch.equal(cache.ids, input_ids[:, :L])):
cache, L = SWACache(s.config.l), 0 # new sequence: start from the full prompt
return {"input_ids": input_ids[:, L:], "past_key_values": cache, "use_cache": True}
def forward(s, input_ids=None, attention_mask=None, labels=None, past_key_values=None, use_cache=None, **kw): # tied head
if attention_mask is not None and not bool(attention_mask.all()): # padding is not supported: fail loudly, never silently
raise ValueError("SWAForCausalLM does not support padded batches (attention_mask has zeros); "
"run sequences one at a time or batch only equal-length sequences")
cache =past_key_values if isinstance(past_key_values, SWACache) else SWACache(s.config.l) if use_cache else None
logits = F.linear(s.model(input_ids, cache=cache), s.model.emb.weight).float()
if cache is not None: s.__dict__["_swa_last"] = cache # plain attribute, not a registered module
loss = None if labels is None else F.cross_entropy(logits[:, :-1].flatten(0, 1), labels[:, 1:].flatten(), ignore_index=-100)
return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=cache)