Fables / modeling_fwkv.py

Upload 2 files

dd08deb verified about 22 hours ago

14.7 kB

# 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)

	# 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)