modeling_swa.py

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 <s> doc </s> 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 <s>); every other caller gets ids from the <s> 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)