Upload 2 files

configuration_fwkv.py

@@ -0,0 +1,44 @@
+from transformers import PretrainedConfig
+
+
+class FWKVConfig(PretrainedConfig):
+    model_type = "fwkv"
+
+    def __init__(
+        self,
+        vocab_size=50257,
+        d_model=768,
+        d_emb=384,
+        n_layers=13,
+        ffn_mult=4,
+        head_size=64,
+        wkv_floor=0.1,
+        wkv_ceil=0.999,
+        tau_min=2.0,
+        tau_max=1000.0,
+        seq_len=1024,
+        ce_chunk=512,
+        arch_version=2,
+        tie_word_embeddings=True,
+        **kwargs,
+    ):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_emb = d_emb
+        self.n_layers = n_layers
+        self.ffn_mult = ffn_mult
+        self.head_size = head_size
+        self.wkv_floor = wkv_floor
+        self.wkv_ceil = wkv_ceil
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+        self.seq_len = seq_len
+        self.ce_chunk = ce_chunk
+        self.arch_version = arch_version
+
+        # ─── Standard HF aliases ───
+        self.num_hidden_layers = n_layers
+        self.hidden_size = d_model
+        self.num_attention_heads = 1       # dummy, not used by FWKV
+        self.num_key_value_heads = 1       # dummy

modeling_fwkv.py

@@ -0,0 +1,345 @@
+# modeling_fwkv.py
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, AutoModel, AutoConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.generation import GenerationMixin
+from transformers.cache_utils import Cache  # needed to detect Cache objects
+from .configuration_fwkv import FWKVConfig
+
+
+# ════════════════════════════════════════════════════════════════════════════
+# utility functions
+# ════════════════════════════════════════════════════════════════════════════
+def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
+    """REFERENCE implementation (Hillis-Steele). y_t = a_t + W * y_{t-1}.
+    Kept for parity/testing with decay_scan; not used at inference time."""
+    W = W.to(dtype=a.dtype)
+    val = a
+    T = a.shape[1]
+    d = 1
+    while d < T:
+        shifted = F.pad(val[:, :-d, :], (0, 0, d, 0))
+        val = val + (W ** d) * shifted
+        d *= 2
+    return val
+
+
+def decay_scan(a: torch.Tensor, W: torch.Tensor, state: torch.Tensor | None = None) -> torch.Tensor:
+    """y_t = a_t + W * y_{t-1}, y_{-1} = state, as a causal FFT convolution in fp32.
+
+    a: (B, T, d), W: (d,) with 0 < W < 1, state: (B, d) or None.
+    Matches the training-time implementation exactly (same math as
+    parallel_scan_decay, but computed as an FFT convolution for speed).
+    """
+    with torch.autocast(device_type=a.device.type, enabled=False):
+        a = a.float()
+        W = W.float()
+        B, T, d = a.shape
+        if state is not None:
+            first = a[:, :1] + W.view(1, 1, d) * state.float().unsqueeze(1)
+            a = torch.cat([first, a[:, 1:]], dim=1)
+        if T == 1:
+            return a
+        n = 1 << (2 * T - 1).bit_length()          # >= 2T-1, avoids circular wrap
+        t = torch.arange(T, device=a.device, dtype=a.dtype)
+        K = torch.exp(t.unsqueeze(1) * torch.log(W).unsqueeze(0))   # (T, d) = W^t
+        Fa = torch.fft.rfft(a, n=n, dim=1)
+        Fk = torch.fft.rfft(K, n=n, dim=0)
+        y = torch.fft.irfft(Fa * Fk.unsqueeze(0), n=n, dim=1)[:, :T]
+        return y.contiguous()
+
+
+def chunked_cross_entropy(x, weight, labels, chunk_size=512, ignore_index=-100):
+    x_2d = x.reshape(-1, x.size(-1))
+    labels_1d = labels.reshape(-1)
+    N = x_2d.size(0)
+    total_loss = x_2d.new_zeros((), dtype=torch.float32)
+    total_count = labels_1d.new_zeros(())
+    for start in range(0, N, chunk_size):
+        end = min(start + chunk_size, N)
+        label_chunk = labels_1d[start:end]
+        count = (label_chunk != ignore_index).sum()
+        if count == 0:
+            continue
+        logits_chunk = F.linear(x_2d[start:end], weight)
+        loss_chunk = F.cross_entropy(logits_chunk.float(), label_chunk, ignore_index=ignore_index, reduction="sum")
+        total_loss = total_loss + loss_chunk
+        total_count = total_count + count
+    return total_loss / total_count.clamp(min=1)
+
+
+# ════════════════════════════════════════════════════════════════════════════
+# modules — mirror the trained Myosotis v2 / FWKVBlock exactly
+# ════════════════════════════════════════════════════════════════════════════
+class FactorizedTiedHead(nn.Module):
+    def __init__(self, vocab_size, d_model, d_emb):
+        super().__init__()
+        self.d_model = d_model
+        self.d_emb = d_emb
+        self.weight = nn.Parameter(torch.empty(vocab_size, d_emb))
+        self.proj = nn.Linear(d_emb, d_model, bias=False)
+        nn.init.normal_(self.weight, std=0.02)
+        nn.init.normal_(self.proj.weight, std=0.02)
+
+    def embed(self, input_ids):
+        return self.proj(F.embedding(input_ids, self.weight))
+
+    # kept as an alias so either calling convention works
+    def forward(self, input_ids):
+        return self.embed(input_ids)
+
+    def to_emb_space(self, x):
+        return F.linear(x, self.proj.weight.t())
+
+    def logits(self, x_emb):
+        return F.linear(x_emb, self.weight)
+
+    def resize_vocab(self, new_vocab_size):
+        old = self.weight.shape[0]
+        if new_vocab_size == old:
+            return
+        new_w = torch.empty(new_vocab_size, self.d_emb, device=self.weight.device, dtype=self.weight.dtype)
+        nn.init.normal_(new_w, std=0.02)
+        new_w[:old] = self.weight.data
+        self.weight = nn.Parameter(new_w)
+
+
+class FWKVBlock(nn.Module):
+    """Pre-norm block: token-shifted WKV mixer + token-shifted FFN.
+
+    Per-layer recurrent state = (wkv_state, shift_state_timemix, shift_state_chanmix),
+    each (B, d). Must match the training-time block bit-for-bit (module names and
+    math), since checkpoints are loaded by state_dict key.
+    """
+
+    def __init__(self, d, layer_id, n_layers, ffn_mult=4, floor=0.1, ceil=0.999,
+                 head_size=64, tau_min=2.0, tau_max=1000.0):
+        super().__init__()
+        assert d % head_size == 0, "d_model must be divisible by head_size"
+        self.d = d
+        self.floor = floor
+        self.ceil = ceil
+
+        self.ln1 = nn.LayerNorm(d)
+        self.ln2 = nn.LayerNorm(d)
+        self.proj_k = nn.Linear(d, d, bias=False)
+        self.proj_v = nn.Linear(d, d, bias=False)
+        self.proj_r = nn.Linear(d, d, bias=False)
+        self.proj_out = nn.Linear(d, d, bias=False)
+        self.proj_out._is_out_proj = True
+        self.ln_x = nn.GroupNorm(d // head_size, d)  # per-head norm on mixer output
+        self.ffn = nn.Sequential(
+            nn.Linear(d, ffn_mult * d, bias=False),
+            nn.GELU(),
+            nn.Linear(ffn_mult * d, d, bias=False),
+        )
+        self.ffn[2]._is_out_proj = True
+
+        # Decay: channels span timescales tau in [tau_min, tau_max] (W = 1 - 1/tau).
+        tau = torch.exp(torch.linspace(math.log(tau_min), math.log(tau_max), d))
+        W0 = (1.0 - 1.0 / tau).clamp(floor + 1e-3, ceil - 1e-4)
+        p = (W0 - floor) / (ceil - floor)
+        self.w = nn.Parameter(torch.log(p / (1.0 - p)))
+
+        # Token-shift mix coefficients (RWKV-4 style init).
+        ratio_1_to_almost0 = 1.0 - (layer_id / n_layers)
+        ratio_0_to_1 = layer_id / max(1, n_layers - 1)
+        ramp = torch.arange(d, dtype=torch.float32) / d
+        self.mu_k = nn.Parameter(ramp ** ratio_1_to_almost0)
+        self.mu_v = nn.Parameter(ramp ** ratio_1_to_almost0 + 0.3 * ratio_0_to_1)
+        self.mu_r = nn.Parameter(ramp ** (0.5 * ratio_1_to_almost0))
+        self.mu_c = nn.Parameter(ramp ** ratio_1_to_almost0)
+
+    @property
+    def W(self):
+        return self.floor + (self.ceil - self.floor) * torch.sigmoid(self.w)
+
+    def forward(self, x, state=None):
+        B, T, d = x.shape
+        if state is None:
+            wkv_s = None
+            sh_tm = x.new_zeros(B, d)
+            sh_cm = x.new_zeros(B, d)
+        else:
+            wkv_s, sh_tm, sh_cm = state
+
+        # ── time mix (WKV) ────────────────────────────────────────────────
+        h = self.ln1(x)
+        hp = torch.cat([sh_tm.unsqueeze(1).to(h.dtype), h[:, :-1]], dim=1)   # x_{t-1}
+        xk = h * self.mu_k + hp * (1 - self.mu_k)
+        xv = h * self.mu_v + hp * (1 - self.mu_v)
+        xr = h * self.mu_r + hp * (1 - self.mu_r)
+        k = self.proj_k(xk)
+        v = self.proj_v(xv)
+        r = torch.sigmoid(self.proj_r(xr))
+
+        W = self.W
+        gain = torch.sqrt(1.0 - W * W)                 # unit stationary variance per channel
+        a = k.float() * v.float() * gain
+        y = decay_scan(a, W, wkv_s)                    # (B, T, d) fp32
+        new_wkv = y[:, -1].detach()
+        yn = self.ln_x(y.reshape(B * T, d)).reshape(B, T, d)
+        x = x + self.proj_out(r * yn.to(r.dtype))
+
+        # ── channel mix (FFN) ────────────────────────────────────────────
+        h2 = self.ln2(x)
+        h2p = torch.cat([sh_cm.unsqueeze(1).to(h2.dtype), h2[:, :-1]], dim=1)
+        xc = h2 * self.mu_c + h2p * (1 - self.mu_c)
+        x = x + self.ffn(xc)
+
+        new_state = (new_wkv, h[:, -1].detach(), h2[:, -1].detach())
+        return x, new_state
+
+
+# ════════════════════════════════════════════════════════════════════════════
+# HF wrapper
+# ════════════════════════════════════════════════════════════════════════════
+class FWKVPreTrainedModel(PreTrainedModel):
+    config_class = FWKVConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            std = 0.02
+            if getattr(module, "_is_out_proj", False):
+                std = 0.02 / math.sqrt(2 * self.config.n_layers)
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, FactorizedTiedHead):
+            pass
+
+
+class FWKVLanguageModel(FWKVPreTrainedModel, GenerationMixin):
+    # This model does NOT use the standard KV-Cache machinery (DynamicCache/
+    # StaticCache). Without this, newer transformers versions (4.45+/5.x)
+    # auto-wrap past_key_values in a Cache object during generate(), which this
+    # custom recurrent-state architecture can't consume.
+    _supports_cache_class = False
+    _supports_static_cache = False
+    _supports_quantized_cache = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.amp_dtype = None    # optionally set by a training/inference driver
+        self.shared = FactorizedTiedHead(config.vocab_size, config.d_model, config.d_emb)
+        self.ln0 = nn.LayerNorm(config.d_model)
+        self.blocks = nn.ModuleList([
+            FWKVBlock(config.d_model, i, config.n_layers, config.ffn_mult,
+                      config.wkv_floor, config.wkv_ceil, config.head_size,
+                      config.tau_min, config.tau_max)
+            for i in range(config.n_layers)
+        ])
+        self.norm = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, value):
+        self.shared = value
+
+    def resize_vocab(self, new_vocab_size):
+        self.shared.resize_vocab(new_vocab_size)
+        self.config.vocab_size = new_vocab_size
+
+    def count_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        labels=None,
+        use_cache=True,
+        **kwargs
+    ):
+        if input_ids is None:
+            raise ValueError("input_ids must be provided")
+        if not isinstance(input_ids, torch.Tensor):
+            input_ids = torch.tensor(input_ids)
+        if input_ids.dtype != torch.long:
+            input_ids = input_ids.long()
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)
+
+        x = self.ln0(self.shared.embed(input_ids))
+
+        # If transformers handed us a Cache object (DynamicCache, etc.) instead
+        # of our custom list/tuple of per-layer state tensors, treat it as
+        # "no state yet" rather than trying to iterate/index into it.
+        if past_key_values is None or isinstance(past_key_values, Cache):
+            states_in = [None] * len(self.blocks)
+        else:
+            states_in = past_key_values
+
+        states_out = []
+        for block, state in zip(self.blocks, states_in):
+            x, new_state = block(x, state)
+            states_out.append(new_state)
+
+        x = self.norm(x)
+        x_emb = self.shared.to_emb_space(x)
+
+        loss = None
+        logits = None
+        if labels is not None:
+            if not isinstance(labels, torch.Tensor):
+                labels = torch.tensor(labels)
+            if labels.dtype != torch.long:
+                labels = labels.long()
+            loss = chunked_cross_entropy(
+                x_emb, self.shared.weight, labels,
+                chunk_size=self.config.ce_chunk, ignore_index=-100
+            )
+        else:
+            logits = self.shared.logits(x_emb)
+
+        past_key_values_out = tuple(states_out) if use_cache else None
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values_out,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:] if input_ids.dim() > 1 else input_ids[-1:]
+
+        # Don't coerce a Cache object into a tuple — pass it through as-is,
+        # forward() knows how to treat it as "empty" state.
+        if past_key_values is not None and not isinstance(past_key_values, (tuple, Cache)):
+            past_key_values = tuple(past_key_values)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        if past_key_values is None or isinstance(past_key_values, Cache):
+            return past_key_values
+        reordered = []
+        for state in past_key_values:
+            if state is None:
+                reordered.append(None)
+                continue
+            # each per-layer state is itself a 3-tuple (wkv_state, sh_tm, sh_cm)
+            reordered.append(tuple(s.index_select(0, beam_idx) for s in state))
+        return tuple(reordered)
+
+
+# Register the model with AutoModel
+AutoConfig.register("fwkv", FWKVConfig)
+AutoModel.register(FWKVConfig, FWKVLanguageModel)