modeling_rne_tiny_gpt.py

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPooling

from .configuration_rne_tiny_gpt import RNETinyGPTConfig


class CausalSelfAttention(nn.Module):
    def __init__(self, config: RNETinyGPTConfig):
        super().__init__()

        if config.n_embd % config.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")

        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.attention_backend = getattr(config, "attention_backend", "sage")
        self.torch_fallback = bool(getattr(config, "torch_fallback", False))

        if self.attention_backend not in ("sage", "torch"):
            raise ValueError("attention_backend must be 'sage' or 'torch'")

        if self.attention_backend == "sage" and self.head_dim not in (64, 96, 128):
            raise ValueError(f"SageAttention requires head_dim in [64, 96, 128], got {self.head_dim}.")

        if self.attention_backend == "sage" and config.dropout != 0.0:
            raise ValueError("SageAttention strict mode requires dropout=0.0")

        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(config.ctx_len, config.ctx_len, dtype=torch.bool))
        self.register_buffer("mask", mask.view(1, 1, config.ctx_len, config.ctx_len), persistent=False)

        self.sageattn = None
        if self.attention_backend == "sage":
            try:
                from sageattention import sageattn
                self.sageattn = sageattn
            except Exception as exc:
                if self.torch_fallback:
                    self.attention_backend = "torch"
                    self.sageattn = None
                else:
                    raise RuntimeError(
                        "Ce modèle a été entraîné avec SageAttention. "
                        "Installe sageattention: pip install sageattention"
                    ) from exc

    def _torch_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, t: int) -> torch.Tensor:
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.mask[:, :, :t, :t] == 0, float("-inf"))
        att = F.softmax(scores.float(), dim=-1).to(q.dtype)
        att = self.dropout(att)
        y = att @ v
        return y

    def _sage_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        if self.sageattn is None:
            raise RuntimeError("SageAttention demandé mais sageattn est None")
        if not q.is_cuda:
            if self.torch_fallback:
                return None
            raise RuntimeError(
                "SageAttention exige CUDA. Passe le modèle sur CUDA avec model.cuda(), "
                "ou active torch_fallback dans config.json."
            )
        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
        return self.sageattn(q, k, v, tensor_layout="HND", is_causal=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, t, c = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()

        if self.attention_backend == "sage":
            y = self._sage_attention(q, k, v)
            if y is None:
                y = self._torch_attention(q, k, v, t)
        else:
            y = self._torch_attention(q, k, v, t)

        y = y.transpose(1, 2).contiguous().view(b, t, c)
        y = self.proj(y)
        return y


class MLP(nn.Module):
    def __init__(self, config: RNETinyGPTConfig):
        super().__init__()
        self.fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
        self.proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc(x)
        x = F.gelu(x)
        x = self.proj(x)
        x = self.dropout(x)
        return x


class Block(nn.Module):
    def __init__(self, config: RNETinyGPTConfig):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class RNETinyGPTPreTrainedModel(PreTrainedModel):
    config_class = RNETinyGPTConfig
    base_model_prefix = "rne_tiny_gpt"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        if isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)


class RNETinyGPTModel(RNETinyGPTPreTrainedModel):
    def __init__(self, config: RNETinyGPTConfig):
        super().__init__(config)
        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.pos_emb = nn.Embedding(config.ctx_len, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.post_init()

    def _mean_pool(self, hidden: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor) -> torch.Tensor:
        if attention_mask is None:
            mask = input_ids.ne(self.config.pad_token_id)
        else:
            mask = attention_mask.bool()
        mask = mask.unsqueeze(-1).to(hidden.dtype)
        summed = (hidden * mask).sum(dim=1)
        denom = mask.sum(dim=1).clamp(min=1.0)
        return summed / denom

    def _last_pool(self, hidden: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor) -> torch.Tensor:
        if attention_mask is None:
            mask = input_ids.ne(self.config.pad_token_id)
        else:
            mask = attention_mask.bool()
        lengths = mask.sum(dim=1).clamp(min=1)
        last_pos = lengths - 1
        batch_idx = torch.arange(input_ids.size(0), device=input_ids.device)
        return hidden[batch_idx, last_pos, :]

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = True,
        **kwargs,
    ):
        b, t = input_ids.shape
        if t > self.config.ctx_len:
            raise ValueError(f"Input length {t} > ctx_len {self.config.ctx_len}. Truncate before calling the model.")

        pos = torch.arange(0, t, dtype=torch.long, device=input_ids.device).unsqueeze(0)
        x = self.tok_emb(input_ids) + self.pos_emb(pos)
        x = self.drop(x)
        for block in self.blocks:
            x = block(x)
        hidden = self.ln_f(x)

        if self.config.pooling == "last":
            pooled = self._last_pool(hidden, attention_mask, input_ids)
        else:
            pooled = self._mean_pool(hidden, attention_mask, input_ids)

        pooled = pooled.float()
        if self.config.normalize_embeddings:
            pooled = F.normalize(pooled, p=2, dim=-1)

        if not return_dict:
            return (hidden, pooled)

        return BaseModelOutputWithPooling(
            last_hidden_state=hidden,
            pooler_output=pooled,
            hidden_states=None,
            attentions=None,
        )

    @torch.no_grad()
    def encode(self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        out = self.forward(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
        return out.pooler_output