modeling_gpt2.py

import math
from typing import Optional, Tuple, Dict, Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin, AutoConfig, AutoModelForCausalLM
from transformers import PretrainedConfig

class GPT2CustomConfig(PretrainedConfig):
    model_type = "gpt2_custom"
    auto_map = {
        "AutoConfig": "configuration_gpt2.GPT2CustomConfig",
        "AutoModelForCausalLM": "modeling_gpt2.GPT2CustomLMHeadModel"
    }
    
    def __init__(
        self,
        vocab_size=16384,
        n_positions=512,
        n_embd=512,
        n_layer=12,
        n_head=8,
        n_inner=1360,
        activation_function="gelu_new",
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
        layer_norm_epsilon=1e-5,
        initializer_range=0.02,
        bos_token_id=2,
        eos_token_id=3,
        **kwargs
    ):
        self.vocab_size = vocab_size
        self.n_positions = n_positions
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_inner = n_inner
        self.activation_function = activation_function
        self.resid_pdrop = resid_pdrop
        self.embd_pdrop = embd_pdrop
        self.attn_pdrop = attn_pdrop
        self.layer_norm_epsilon = layer_norm_epsilon
        self.initializer_range = initializer_range
        self.hidden_size = n_embd
        self.num_attention_heads = n_head
        self.num_hidden_layers = n_layer
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

# ─────────────────────────────────────────────────────────────
# ROPE HELPERS
# ─────────────────────────────────────────────────────────────

def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos(), emb.sin()

def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    L = x.size(2)
    cos = cos[:L, :].unsqueeze(0).unsqueeze(1)
    sin = sin[:L, :].unsqueeze(0).unsqueeze(1)

    half_dim = x.size(-1) // 2
    x1 = x[..., :half_dim]
    x2 = x[..., half_dim:]
    rotated_x = torch.cat((-x2, x1), dim=-1)

    return (x * cos) + (rotated_x * sin)

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.num_heads = config.n_head
        self.head_dim = config.n_embd // config.n_head
        
        self.q_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.k_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.v_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.o_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)

    def forward(self, x, past_kv=None, use_cache: bool = False):
        B, L, D = x.size()
        
        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)

        if past_kv is not None:
            past_k, past_v = past_kv
            past_len = past_k.size(2)
            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        else:
            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
            q = _apply_rope(q, cos, sin)
            k = _apply_rope(k, cos, sin)

        L_kv = k.size(2)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # Standard full context causal attention mask (no sliding window)
        past_len = L_kv - L
        pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
        pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
        causal_mask = (pos_i - pos_j) >= 0

        scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))

        attn = torch.softmax(scores, dim=-1)
        attn = self.attn_dropout(attn)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, L, D)
        out = self.resid_dropout(self.o_proj(out))

        present_kv = (k, v) if use_cache else None
        return out, present_kv

class SwiGLU(nn.Module):
    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
        self.fc2 = nn.Linear(dim, hidden_dim, bias=False)
        self.fc3 = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x):
        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))

class GPT2Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.mlp = SwiGLU(config.n_embd, config.n_inner)

    def forward(self, x, past_kv=None, use_cache: bool = False):
        attn_out, present_kv = self.attn(self.ln_1(x), past_kv=past_kv, use_cache=use_cache)
        x = x + attn_out
        x = x + self.mlp(self.ln_2(x))
        return x, present_kv

class GPT2CustomLMHeadModel(PreTrainedModel, GenerationMixin):
    config_class = GPT2CustomConfig
    base_model_prefix = "transformer"
    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
    
    def __init__(self, config):
        super().__init__(config)
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            drop = nn.Dropout(config.embd_pdrop),
            h = nn.ModuleList([GPT2Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.transformer.wte

    def set_input_embeddings(self, new_embeddings):
        self.transformer.wte = new_embeddings

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

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

    def forward(self, input_ids, labels=None, past_key_values=None, use_cache: bool = False, **kwargs):
        device = input_ids.device
        
        x = self.transformer.wte(input_ids)
        x = self.transformer.drop(x)
        
        new_kvs = []
        for i, block in enumerate(self.transformer.h):
            pkv = past_key_values[i] if past_key_values is not None else None
            x, nkv = block(x, past_kv=pkv, use_cache=use_cache)
            if use_cache:
                new_kvs.append(nkv)
            
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        
        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)
            
        present_kvs = tuple(new_kvs) if use_cache else None
        from transformers.modeling_outputs import CausalLMOutputWithPast
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=present_kvs
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        return {"input_ids": input_ids, "past_key_values": past_key_values}

# Register configuration and model for auto mapping
AutoConfig.register("gpt2_custom", GPT2CustomConfig)
AutoModelForCausalLM.register(GPT2CustomConfig, GPT2CustomLMHeadModel)