Update modeling code

configuration_pit.py

@@ -12,13 +12,23 @@ class PITConfig(PretrainedConfig):
         n_layer: int = 20,
         n_head: int = 32,
         n_embd: int = 4096,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = True,
         **kwargs,
     ):
+        kwargs.pop("tie_word_embeddings", None)
         super().__init__(**kwargs)
         self.vocab_size = vocab_size
         self.n_layer = n_layer
         self.n_head = n_head
         self.n_embd = n_embd
+        # KV caching during generation; the base PretrainedConfig does not
+        # define this, so the model reads it from here.
+        self.use_cache = use_cache
+        # lm_head shares its weight with the input embedding. transformers v5
+        # treats an *absent* `tie_word_embeddings` as False and would leave the
+        # two untied, so state it explicitly.
+        self.tie_word_embeddings = tie_word_embeddings
         # Standard aliases expected by transformers internals
         self.num_hidden_layers = n_layer
         self.hidden_size = n_embd

modeling_pit.py

@@ -3,12 +3,17 @@ PIT (Point-In-Time) GPT model — self-contained for trust_remote_code=True load
 
 Architecture: decoder-only Transformer with RoPE, RMSNorm on Q/K, squared-ReLU
 MLP, and weight-tied input/output embeddings.
+
+Generation is KV-cached: past keys and values are kept in a `DynamicCache`, so
+each step runs attention with a single query position against the cache instead
+of re-running the whole prefix.
 """
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PreTrainedModel
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.cache_utils import Cache, DynamicCache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from .configuration_pit import PITConfig
@@ -23,24 +28,23 @@ class Rotary(nn.Module):
         super().__init__()
         self.dim = dim
         self.base = base * scaling_factor
-        self.seq_len_cached: int | None = None
-        self.cos_cached: torch.Tensor | None = None
-        self.sin_cached: torch.Tensor | None = None
-
-    def forward(self, x: torch.Tensor):
-        seq_len = x.shape[1]
-        if seq_len != self.seq_len_cached:
-            self.seq_len_cached = seq_len
-            # Compute inv_freq on-the-fly on the correct device — never stored
-            # as a buffer so device_map="auto" / meta-device loading can't break it.
-            inv_freq = 1.0 / (self.base ** (
-                torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32) / self.dim
+        self.inv_freq: torch.Tensor | None = None
+
+    def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Absolute positions [B, T] -> (cos, sin), each of shape [B, T, 1, dim // 2].
+
+        Positions are passed in rather than derived from the sequence length so
+        that a cached decode step rotates the new token by its *absolute*
+        position, not by 0.
+        """
+        if self.inv_freq is None or self.inv_freq.device != position_ids.device:
+            # Computed on-the-fly on the correct device — never stored as a
+            # buffer so device_map="auto" / meta-device loading can't break it.
+            self.inv_freq = 1.0 / (self.base ** (
+                torch.arange(0, self.dim, 2, device=position_ids.device, dtype=torch.float32) / self.dim
             ))
-            t = torch.arange(seq_len, device=x.device, dtype=torch.float32)
-            freqs = torch.outer(t, inv_freq)
-            self.cos_cached = freqs.cos().bfloat16()
-            self.sin_cached = freqs.sin().bfloat16()
-        return self.cos_cached[None, :, None, :], self.sin_cached[None, :, None, :]
+        freqs = position_ids.float()[:, :, None] * self.inv_freq
+        return freqs.cos().bfloat16()[:, :, None, :], freqs.sin().bfloat16()[:, :, None, :]
 
 
 def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
@@ -50,8 +54,9 @@ def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) ->
 
 
 class CausalSelfAttention(nn.Module):
-    def __init__(self, config: PITConfig):
+    def __init__(self, config: PITConfig, layer_idx: int = 0):
         super().__init__()
+        self.layer_idx = layer_idx
         self.n_head = config.n_head
         self.n_embd = config.n_embd
         self.head_dim = config.n_embd // config.n_head
@@ -60,19 +65,30 @@ class CausalSelfAttention(nn.Module):
         self.c_v = nn.Linear(config.n_embd, config.n_embd, bias=False)
         self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
         self.c_proj.weight.data.zero_()
-        self.rotary = Rotary(self.head_dim)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        is_causal: bool = False,
+        past_key_values: Cache | None = None,
+    ) -> torch.Tensor:
         B, T, C = x.size()
         q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
         k = self.c_k(x).view(B, T, self.n_head, self.head_dim)
         v = self.c_v(x).view(B, T, self.n_head, self.head_dim)
-        cos, sin = self.rotary(q)
         q = _apply_rotary_emb(F.rms_norm(q, (q.size(-1),)), cos, sin)
         k = _apply_rotary_emb(F.rms_norm(k, (k.size(-1),)), cos, sin)
-        y = F.scaled_dot_product_attention(
-            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True
-        )
+
+        # [B, T, H, D] -> [B, H, T, D], the layout the cache and SDPA expect.
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+        if past_key_values is not None:
+            # RoPE is already applied, so cached keys stay valid for later steps.
+            k, v = past_key_values.update(k, v, self.layer_idx)
+
+        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal)
         return self.c_proj(y.transpose(1, 2).contiguous().view_as(x))
 
 
@@ -88,13 +104,23 @@ class MLP(nn.Module):
 
 
 class Block(nn.Module):
-    def __init__(self, config: PITConfig):
+    def __init__(self, config: PITConfig, layer_idx: int = 0):
         super().__init__()
-        self.attn = CausalSelfAttention(config)
+        self.attn = CausalSelfAttention(config, layer_idx)
         self.mlp  = MLP(config)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x + self.attn(F.rms_norm(x, (x.size(-1),)))
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        is_causal: bool = False,
+        past_key_values: Cache | None = None,
+    ) -> torch.Tensor:
+        x = x + self.attn(
+            F.rms_norm(x, (x.size(-1),)), cos, sin, attn_mask, is_causal, past_key_values
+        )
         x = x + self.mlp(F.rms_norm(x, (x.size(-1),)))
         return x
 
@@ -103,7 +129,7 @@ class Block(nn.Module):
 # HuggingFace PreTrainedModel wrapper
 # ---------------------------------------------------------------------------
 
-class PITForCausalLM(PreTrainedModel):
+class PITForCausalLM(PreTrainedModel, GenerationMixin):
     """
     Point-In-Time GPT wrapped as a HuggingFace CausalLM.
 
@@ -123,17 +149,25 @@ class PITForCausalLM(PreTrainedModel):
 
     config_class = PITConfig
     _no_split_modules = ["Block"]
-    _supports_cache_class = False
-    # Weight tying: lm_head and transformer.wte share parameters.
-    _tied_weights_keys = ["lm_head.weight", "transformer.wte.weight"]
+    _supports_cache_class = True          # only read by transformers < 4.45
+    # `bool(attention_mask.all())` below forces a host sync, and the KV cache is
+    # dynamically sized — neither is fullgraph-compilable.
+    _can_compile_fullgraph = False
+    # Weight tying: lm_head and transformer.wte share parameters. transformers
+    # v5 expects {tied key: source key}; older versions iterate it as a list of
+    # tied keys, which yields "lm_head.weight" — also correct.
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
 
     def __init__(self, config: PITConfig):
         super().__init__(config)
         self.transformer = nn.ModuleDict({
             "wte": nn.Embedding(config.vocab_size, config.n_embd),
-            "h":   nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            "h":   nn.ModuleList([Block(config, i) for i in range(config.n_layer)]),
         })
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # Parameter-free, so it adds nothing to the state dict. Shared by every
+        # block: cos/sin depend only on position, not on the layer.
+        self.rotary = Rotary(config.n_embd // config.n_head)
         # Tie weights (re-tied after load_state_dict via tie_weights())
         self.transformer["wte"].weight = self.lm_head.weight
         self.post_init()
@@ -152,20 +186,86 @@ class PITForCausalLM(PreTrainedModel):
     def set_output_embeddings(self, value: nn.Linear) -> None:
         self.lm_head = value
 
+    # -- attention masking ---------------------------------------------------
+
+    @staticmethod
+    def _causal_mask(
+        attention_mask: torch.Tensor | None,
+        q_len: int,
+        past_len: int,
+        device: torch.device,
+    ) -> tuple[torch.Tensor | None, bool]:
+        """Return the (attn_mask, is_causal) pair to hand to SDPA.
+
+        The two `None` cases are the fast paths — SDPA can pick a fused kernel
+        only when no explicit mask is materialised:
+          • one query token: it may attend to the entire cache, no mask needed;
+          • uncached prefill: plain `is_causal=True`.
+        Anything else (chunked prefill on top of a cache, or padded batches)
+        needs an explicit mask, because `is_causal=True` aligns to the *top
+        left* of a non-square score matrix and would mask the cache away.
+        """
+        if attention_mask is not None:
+            if attention_mask.dim() == 4:
+                return attention_mask, False           # already prepared by the caller
+            if bool(attention_mask.all()):
+                attention_mask = None                  # all-ones carries no information
+
+        if attention_mask is None:
+            if q_len == 1:
+                return None, False
+            if past_len == 0:
+                return None, True
+
+        kv_len = past_len + q_len
+        q_pos  = torch.arange(past_len, kv_len, device=device)[:, None]
+        kv_pos = torch.arange(kv_len, device=device)[None, :]
+        mask = (kv_pos <= q_pos)[None, None]           # [1, 1, q_len, kv_len]
+        if attention_mask is not None:
+            mask = mask & attention_mask[:, None, None, :].bool()
+            # A left-padded row can end up fully masked, which makes softmax
+            # produce NaN. Let those (discarded) pad rows attend freely.
+            mask = mask | ~mask.any(-1, keepdim=True)
+        return mask, False
+
     # -- forward -------------------------------------------------------------
 
     def forward(
         self,
         input_ids: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
         labels: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
         **kwargs,
     ) -> CausalLMOutputWithPast:
+        # Training/eval passes `labels` and never reuses the cache, so don't pay
+        # for it unless the caller explicitly asks.
+        if use_cache is None:
+            use_cache = self.config.use_cache and labels is None
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+        T = input_ids.shape[1]
+        if position_ids is None:
+            position_ids = torch.arange(past_len, past_len + T, device=input_ids.device)
+        position_ids = position_ids.view(-1, T)
+
         x = self.transformer["wte"](input_ids)
+        cos, sin = self.rotary(position_ids)
+        attn_mask, is_causal = self._causal_mask(attention_mask, T, past_len, x.device)
+
         for block in self.transformer["h"]:
-            x = block(x)
+            x = block(x, cos, sin, attn_mask, is_causal, past_key_values)
         x = F.rms_norm(x, (x.size(-1),))
-        logits = self.lm_head(x).float()
+
+        # Only the last position matters while generating; computing the full
+        # [B, T, vocab_size] logits during prefill costs hundreds of MB.
+        keep = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(x[:, keep]).float()
 
         loss = None
         if labels is not None:
@@ -175,9 +275,8 @@ class PITForCausalLM(PreTrainedModel):
                 ignore_index=-100,
             )
 
-        return CausalLMOutputWithPast(loss=loss, logits=logits)
-
-    def prepare_inputs_for_generation(
-        self, input_ids: torch.Tensor, **kwargs
-    ) -> dict:
-        return {"input_ids": input_ids}
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values if use_cache else None,
+        )