modeling_statehead.py

from __future__ import annotations

import os
import shutil
import warnings
from collections.abc import Sequence
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from .configuration_statehead import StateHeadConfig


def _sqrt_with_clipped_grad(x: torch.Tensor, max_grad: float) -> torch.Tensor:
    eps = 1.0 / (4.0 * max_grad * max_grad)
    return torch.sqrt(torch.clamp(x, min=0.0) + eps)


def _retention_write_scale_from_logits(a_logits: torch.Tensor, max_grad: float) -> torch.Tensor:
    a = torch.sigmoid(a_logits)
    scale_square = torch.sigmoid(-a_logits) * (1.0 + a)
    return _sqrt_with_clipped_grad(scale_square, max_grad)


def _statehead_scan_torch(
    gates: torch.Tensor,
    initial_state: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    write_scale_by_retention: bool = False,
    write_scale_by_retention_multiplier: float = 1.0,
    write_scale_sqrt_grad_clip: float = 1000.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    batch, seq_len, four, n_head, head_dim = gates.shape
    if four != 4:
        raise ValueError("gates must have shape [batch, seq, 4, n_head, head_dim]")
    if initial_state.shape != (batch, n_head, head_dim):
        raise ValueError("initial_state must have shape [batch, n_head, head_dim]")

    if attention_mask is not None:
        if attention_mask.shape != (batch, seq_len):
            raise ValueError("attention_mask must have shape [batch, seq]")
        attention_mask = attention_mask.to(dtype=gates.dtype, device=gates.device)

    state = initial_state
    outputs: list[torch.Tensor] = []
    for t in range(seq_len):
        a_logits, b_logits, c_logits, o_logits = gates[:, t].unbind(dim=1)
        a = torch.sigmoid(a_logits)
        b = torch.sigmoid(b_logits)
        c = torch.tanh(c_logits)
        o = torch.sigmoid(o_logits)
        write = b * c
        if write_scale_by_retention:
            write = (
                write_scale_by_retention_multiplier
                * _retention_write_scale_from_logits(a_logits, write_scale_sqrt_grad_clip)
                * write
            )
        next_state = a * state + write
        y = o * next_state
        if attention_mask is not None:
            mask = attention_mask[:, t].view(batch, 1, 1)
            state = mask * next_state + (1.0 - mask) * state
            y = mask * y
        else:
            state = next_state
        outputs.append(y)
    return torch.stack(outputs, dim=1), state


def _statehead_scan_parallel(
    gates: torch.Tensor,
    initial_state: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    chunk_size: int = 64,
    write_scale_by_retention: bool = False,
    write_scale_by_retention_multiplier: float = 1.0,
    write_scale_sqrt_grad_clip: float = 1000.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    batch, seq_len, four, n_head, head_dim = gates.shape
    if four != 4:
        raise ValueError("gates must have shape [batch, seq, 4, n_head, head_dim]")
    if initial_state.shape != (batch, n_head, head_dim):
        raise ValueError("initial_state must have shape [batch, n_head, head_dim]")
    if attention_mask is not None and attention_mask.shape != (batch, seq_len):
        raise ValueError("attention_mask must have shape [batch, seq]")
    if chunk_size < 1:
        raise ValueError("chunk_size must be positive")

    original_dtype = gates.dtype
    scan_dtype = torch.float32 if gates.dtype in {torch.float16, torch.bfloat16} else gates.dtype
    gates = gates.to(scan_dtype)
    state = initial_state.to(scan_dtype)
    if attention_mask is not None:
        attention_mask = attention_mask.to(dtype=scan_dtype, device=gates.device)

    a_logits, b_logits, c_logits, o_logits = gates.unbind(dim=2)
    a = torch.sigmoid(a_logits)
    u = torch.sigmoid(b_logits) * torch.tanh(c_logits)
    if write_scale_by_retention:
        u = (
            write_scale_by_retention_multiplier
            * _retention_write_scale_from_logits(a_logits, write_scale_sqrt_grad_clip)
            * u
        )
    o = torch.sigmoid(o_logits)

    mask = None
    if attention_mask is not None:
        mask = attention_mask.view(batch, seq_len, 1, 1)
        a = mask * a + (1.0 - mask)
        u = mask * u

    chunk_size = min(chunk_size, seq_len)
    n_chunks = (seq_len + chunk_size - 1) // chunk_size
    padded_seq_len = n_chunks * chunk_size
    pad_len = padded_seq_len - seq_len
    if pad_len:
        pad_shape = (batch, pad_len, n_head, head_dim)
        a = torch.cat([a, torch.ones(pad_shape, dtype=scan_dtype, device=gates.device)], dim=1)
        u = torch.cat([u, torch.zeros(pad_shape, dtype=scan_dtype, device=gates.device)], dim=1)

    a = a.view(batch, n_chunks, chunk_size, n_head, head_dim)
    u = u.view(batch, n_chunks, chunk_size, n_head, head_dim)
    prefix_a, prefix_u = _statehead_associative_prefix(a, u, dim=2)

    chunk_a = prefix_a[:, :, -1]
    chunk_u = prefix_u[:, :, -1]
    chunk_prefix_a, chunk_prefix_u = _statehead_associative_prefix(chunk_a, chunk_u, dim=1)
    states_after_chunks = chunk_prefix_a * state.unsqueeze(1) + chunk_prefix_u
    chunk_initials = torch.cat([state.unsqueeze(1), states_after_chunks[:, :-1]], dim=1)

    states = prefix_a * chunk_initials.unsqueeze(2) + prefix_u
    states = states.reshape(batch, padded_seq_len, n_head, head_dim)[:, :seq_len]
    y = o * states
    if mask is not None:
        y = mask * y
    return y.to(original_dtype), states_after_chunks[:, -1].to(original_dtype)


def _statehead_associative_prefix(
    a: torch.Tensor,
    u: torch.Tensor,
    *,
    dim: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Parallel prefix for elementwise recurrences s[t] = a[t] * s[t-1] + u[t]."""
    size = a.shape[dim]
    offset = 1
    while offset < size:
        head = [slice(None)] * a.ndim
        tail = [slice(None)] * a.ndim
        prev = [slice(None)] * a.ndim
        head[dim] = slice(0, offset)
        tail[dim] = slice(offset, None)
        prev[dim] = slice(0, -offset)

        head_idx = tuple(head)
        tail_idx = tuple(tail)
        prev_idx = tuple(prev)
        a_tail = a[tail_idx]
        u_tail = u[tail_idx]
        a = torch.cat([a[head_idx], a_tail * a[prev_idx]], dim=dim)
        u = torch.cat([u[head_idx], a_tail * u[prev_idx] + u_tail], dim=dim)
        offset <<= 1
    return a, u


def _statehead_scan_chunk_loop(
    gates: torch.Tensor,
    initial_state: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    chunk_size: int = 64,
    write_scale_by_retention: bool = False,
    write_scale_by_retention_multiplier: float = 1.0,
    write_scale_sqrt_grad_clip: float = 1000.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    batch, seq_len, four, n_head, head_dim = gates.shape
    if four != 4:
        raise ValueError("gates must have shape [batch, seq, 4, n_head, head_dim]")
    if initial_state.shape != (batch, n_head, head_dim):
        raise ValueError("initial_state must have shape [batch, n_head, head_dim]")
    if attention_mask is not None and attention_mask.shape != (batch, seq_len):
        raise ValueError("attention_mask must have shape [batch, seq]")
    if chunk_size < 1:
        raise ValueError("chunk_size must be positive")

    original_dtype = gates.dtype
    scan_dtype = torch.float32 if gates.dtype in {torch.float16, torch.bfloat16} else gates.dtype
    gates = gates.to(scan_dtype)
    state = initial_state.to(scan_dtype)
    if attention_mask is not None:
        attention_mask = attention_mask.to(dtype=scan_dtype, device=gates.device)

    chunks: list[torch.Tensor] = []
    for start in range(0, seq_len, chunk_size):
        stop = min(start + chunk_size, seq_len)
        chunk_outputs: list[torch.Tensor] = []
        for t in range(start, stop):
            a_logits, b_logits, c_logits, o_logits = gates[:, t].unbind(dim=1)
            a = torch.sigmoid(a_logits)
            b = torch.sigmoid(b_logits)
            c = torch.tanh(c_logits)
            o = torch.sigmoid(o_logits)
            write = b * c
            if write_scale_by_retention:
                write = (
                    write_scale_by_retention_multiplier
                    * _retention_write_scale_from_logits(a_logits, write_scale_sqrt_grad_clip)
                    * write
                )
            next_state = a * state + write
            y = o * next_state
            if attention_mask is not None:
                mask = attention_mask[:, t].view(batch, 1, 1)
                state = mask * next_state + (1.0 - mask) * state
                y = mask * y
            else:
                state = next_state
            chunk_outputs.append(y)
        chunks.append(torch.stack(chunk_outputs, dim=1))

    return torch.cat(chunks, dim=1).to(original_dtype), state.to(original_dtype)


_compiled_parallel_scan = None


def _parallel_scan_fn():
    """Opt-in compiled scan: set STATEHEAD_COMPILE=1. Inference-only (compiled
    backward measured slower), dynamic shapes, bit-exact vs eager; falls back to
    eager permanently if compilation or the first compiled call fails."""
    global _compiled_parallel_scan
    if _compiled_parallel_scan is None:
        try:
            _compiled_parallel_scan = torch.compile(_statehead_scan_parallel, dynamic=True)
        except Exception:
            _compiled_parallel_scan = _statehead_scan_parallel
    return _compiled_parallel_scan


def statehead_scan(
    gates: torch.Tensor,
    initial_state: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    backend: str = "torch",
    chunk_size: int = 64,
    write_scale_by_retention: bool = False,
    write_scale_by_retention_multiplier: float = 1.0,
    write_scale_sqrt_grad_clip: float = 1000.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    if backend == "auto":
        backend = "parallel"
    if backend == "torch":
        backend = "parallel"
    if backend == "parallel":
        args = (
            gates,
            initial_state,
            attention_mask,
            chunk_size,
            write_scale_by_retention,
            write_scale_by_retention_multiplier,
            write_scale_sqrt_grad_clip,
        )
        if os.environ.get("STATEHEAD_COMPILE", "0") == "1" and not torch.is_grad_enabled():
            global _compiled_parallel_scan
            try:
                return _parallel_scan_fn()(*args)
            except Exception:
                _compiled_parallel_scan = _statehead_scan_parallel
        return _statehead_scan_parallel(*args)
    if backend == "chunk_loop":
        return _statehead_scan_chunk_loop(
            gates,
            initial_state,
            attention_mask,
            chunk_size,
            write_scale_by_retention,
            write_scale_by_retention_multiplier,
            write_scale_sqrt_grad_clip,
        )
    if backend == "torch_loop":
        return _statehead_scan_torch(
            gates,
            initial_state,
            attention_mask,
            write_scale_by_retention,
            write_scale_by_retention_multiplier,
            write_scale_sqrt_grad_clip,
        )
    if backend == "triton":
        warnings.warn(
            "The Triton StateHead scan is not wired yet; falling back to the PyTorch scan.",
            RuntimeWarning,
            stacklevel=2,
        )
        return _statehead_scan_parallel(
            gates,
            initial_state,
            attention_mask,
            chunk_size,
            write_scale_by_retention,
            write_scale_by_retention_multiplier,
            write_scale_sqrt_grad_clip,
        )
    raise ValueError(f"Unknown StateHead scan backend: {backend}")


def count_unique_parameters(model: nn.Module) -> int:
    seen: set[int] = set()
    total = 0
    for parameter in model.parameters():
        pointer = id(parameter)
        if pointer in seen:
            continue
        seen.add(pointer)
        total += parameter.numel()
    return total


class StateHeadMLP(nn.Module):
    def __init__(self, config: StateHeadConfig) -> None:
        super().__init__()
        hidden_dim = config.mlp_ratio * config.n_embd
        self.up = nn.Linear(config.n_embd, hidden_dim)
        self.down = nn.Linear(hidden_dim, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.down(F.gelu(self.up(x))))


class StateHeadBank(nn.Module):
    def __init__(self, config: StateHeadConfig) -> None:
        super().__init__()
        self.n_head = config.n_head
        self.head_dim = config.head_dim
        self.n_embd = config.n_embd
        self.scan_backend = config.scan_backend
        self.scan_chunk_size = config.scan_chunk_size
        self.write_scale_by_retention = config.state_write_scale_by_retention
        self.write_scale_by_retention_multiplier = config.state_write_scale_by_retention_multiplier
        self.write_scale_sqrt_grad_clip = config.state_write_scale_sqrt_grad_clip
        self.gate = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.proj = nn.Linear(config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)
        if config.learned_initial_state:
            initial_state = torch.zeros(config.n_head, config.head_dim)
            if config.initial_state_std > 0:
                initial_state.normal_(mean=0.0, std=config.initial_state_std)
            self.initial_state = nn.Parameter(initial_state)
        else:
            self.register_parameter("initial_state", None)

    def reset_statehead_bias(self, retention_bias: float, write_gate_bias: float, output_gate_bias: float) -> None:
        with torch.no_grad():
            self.gate.bias.zero_()
            self.gate.bias[: self.n_embd].fill_(retention_bias)
            self.gate.bias[self.n_embd : 2 * self.n_embd].fill_(write_gate_bias)
            self.gate.bias[3 * self.n_embd :].fill_(output_gate_bias)

    def zero_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        return torch.zeros(batch_size, self.n_head, self.head_dim, device=device, dtype=dtype)

    def fresh_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        if self.initial_state is None:
            return self.zero_state(batch_size, device=device, dtype=dtype)
        return self.initial_state.to(device=device, dtype=dtype).unsqueeze(0).expand(batch_size, -1, -1)

    def forward(
        self,
        x: torch.Tensor,
        state: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, seq_len, _ = x.shape
        if state is None:
            state = self.fresh_state(batch, device=x.device, dtype=x.dtype)
        gates = self.gate(x).view(batch, seq_len, 4, self.n_head, self.head_dim)
        y, next_state = statehead_scan(
            gates,
            state,
            attention_mask,
            self.scan_backend,
            self.scan_chunk_size,
            self.write_scale_by_retention,
            self.write_scale_by_retention_multiplier,
            self.write_scale_sqrt_grad_clip,
        )
        y = y.reshape(batch, seq_len, self.n_embd)
        return self.dropout(self.proj(y)), next_state


class StateHeadBlock(nn.Module):
    def __init__(self, config: StateHeadConfig) -> None:
        super().__init__()
        self.residual_scale = config.residual_scale
        self.norm_state = nn.RMSNorm(config.n_embd, eps=config.layer_norm_eps)
        self.state_bank = StateHeadBank(config)
        self.norm_mlp = nn.RMSNorm(config.n_embd, eps=config.layer_norm_eps)
        self.mlp = StateHeadMLP(config)

    def forward(
        self,
        x: torch.Tensor,
        state: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        y, next_state = self.state_bank(self.norm_state(x), state, attention_mask)
        x = x + self.residual_scale * y
        x = x + self.residual_scale * self.mlp(self.norm_mlp(x))
        return x, next_state


class StateHeadBackbone(nn.Module):
    def __init__(self, config: StateHeadConfig) -> None:
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = (
            nn.Embedding(config.block_size, config.n_embd) if config.use_position_embeddings else None
        )
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList(StateHeadBlock(config) for _ in range(config.n_layer))
        self.norm = nn.RMSNorm(config.n_embd, eps=config.layer_norm_eps)

    def reset_statehead_biases(self) -> None:
        for block in self.blocks:
            block.state_bank.reset_statehead_bias(
                self.config.state_retention_bias,
                self.config.state_write_gate_bias,
                self.config.state_output_gate_bias,
            )

    def zero_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, ...]:
        return tuple(block.state_bank.fresh_state(batch_size, device=device, dtype=dtype) for block in self.blocks)

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Sequence[torch.Tensor] | None = None,
        use_cache: bool = False,
        output_hidden_states: bool = False,
    ) -> BaseModelOutputWithPast:
        if input_ids is None and inputs_embeds is None:
            raise ValueError("Either input_ids or inputs_embeds must be provided")
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("Only one of input_ids or inputs_embeds may be provided")

        if inputs_embeds is None:
            inputs_embeds = self.token_embedding(input_ids)
        batch_size, seq_len, _ = inputs_embeds.shape

        if attention_mask is not None:
            if attention_mask.shape[-1] != seq_len:
                attention_mask = attention_mask[:, -seq_len:]
            attention_mask = attention_mask.to(device=inputs_embeds.device)

        x = inputs_embeds
        if self.position_embedding is not None:
            if position_ids is None:
                position_ids = torch.arange(seq_len, device=x.device).unsqueeze(0)
            if position_ids.shape[-1] != seq_len:
                position_ids = position_ids[:, -seq_len:]
            position_ids = position_ids.clamp(max=self.config.block_size - 1)
            x = x + self.position_embedding(position_ids)
        x = self.dropout(x)

        if past_key_values is None:
            states: Sequence[torch.Tensor | None] = [None] * len(self.blocks)
        else:
            if len(past_key_values) != len(self.blocks):
                raise ValueError("past_key_values must contain one state per layer")
            states = past_key_values

        all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None
        if all_hidden_states is not None:
            all_hidden_states += (x,)

        next_states: list[torch.Tensor] = []
        for block, state in zip(self.blocks, states, strict=True):
            x, next_state = block(x, state, attention_mask)
            next_states.append(next_state)
            if all_hidden_states is not None:
                all_hidden_states += (x,)

        x = self.norm(x)
        return BaseModelOutputWithPast(
            last_hidden_state=x,
            past_key_values=tuple(next_states) if use_cache else None,
            hidden_states=all_hidden_states,
        )


class StateHeadPreTrainedModel(PreTrainedModel):
    config_class = StateHeadConfig
    base_model_prefix = "statehead"
    supports_gradient_checkpointing = False
    _no_split_modules = ["StateHeadBlock"]

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def save_pretrained(self, save_directory: str | Path, *args: Any, **kwargs: Any):
        result = super().save_pretrained(save_directory, *args, **kwargs)
        save_path = Path(save_directory)
        source_dir = Path(__file__).parent
        for filename in ("configuration_statehead.py", "modeling_statehead.py"):
            shutil.copy2(source_dir / filename, save_path / filename)
        return result


class StateHeadModel(StateHeadPreTrainedModel):
    def __init__(self, config: StateHeadConfig) -> None:
        super().__init__(config)
        self.statehead = StateHeadBackbone(config)
        self.post_init()
        self.statehead.reset_statehead_biases()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.statehead.token_embedding

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.statehead.token_embedding = value

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Sequence[torch.Tensor] | None = None,
        use_cache: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **_: Any,
    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, ...]:
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.statehead(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
        )
        if return_dict:
            return outputs
        result: tuple[torch.Tensor, ...] = (outputs.last_hidden_state,)
        if use_cache:
            result += (outputs.past_key_values,)
        if output_hidden_states:
            result += (outputs.hidden_states,)
        return result


class StateHeadForCausalLM(StateHeadPreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: StateHeadConfig) -> None:
        super().__init__(config)
        self.statehead = StateHeadBackbone(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.post_init()
        self.statehead.reset_statehead_biases()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.statehead.token_embedding

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.statehead.token_embedding = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
        self.lm_head = new_embeddings

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Sequence[torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
            if attention_mask is not None:
                attention_mask = attention_mask[:, -1:]
        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
            "use_cache": kwargs.get("use_cache", True),
        }

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Sequence[torch.Tensor] | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **_: Any,
    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.statehead(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
        )
        logits = self.lm_head(outputs.last_hidden_state)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if return_dict:
            return CausalLMOutputWithPast(
                loss=loss,
                logits=logits,
                past_key_values=outputs.past_key_values,
                hidden_states=outputs.hidden_states,
            )

        result: tuple[torch.Tensor, ...] = (logits,)
        if use_cache:
            result += (outputs.past_key_values,)
        if output_hidden_states:
            result += (outputs.hidden_states,)
        return ((loss,) + result) if loss is not None else result