""" PIT (Point-In-Time) GPT model — self-contained for trust_remote_code=True loading. 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 GenerationMixin, PreTrainedModel from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration_pit import PITConfig # --------------------------------------------------------------------------- # Architecture (mirrors models/GPT.py exactly) # --------------------------------------------------------------------------- class Rotary(nn.Module): def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0): super().__init__() self.dim = dim self.base = base * scaling_factor 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 )) 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: d = x.shape[3] // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat([x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos], dim=3).type_as(x) class CausalSelfAttention(nn.Module): 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 self.c_q = nn.Linear(config.n_embd, config.n_embd, bias=False) self.c_k = nn.Linear(config.n_embd, config.n_embd, bias=False) 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_() 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) 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) # [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)) class MLP(nn.Module): def __init__(self, config: PITConfig): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False) self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False) self.c_proj.weight.data.zero_() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.c_proj(F.relu(self.c_fc(x)).square()) class Block(nn.Module): def __init__(self, config: PITConfig, layer_idx: int = 0): super().__init__() self.attn = CausalSelfAttention(config, layer_idx) self.mlp = MLP(config) 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 # --------------------------------------------------------------------------- # HuggingFace PreTrainedModel wrapper # --------------------------------------------------------------------------- class PITForCausalLM(PreTrainedModel, GenerationMixin): """ Point-In-Time GPT wrapped as a HuggingFace CausalLM. Supports AutoModelForCausalLM, generate(), and pipeline("text-generation"). Loading ------- >>> from transformers import AutoTokenizer, AutoModelForCausalLM >>> tokenizer = AutoTokenizer.from_pretrained("Diamegs/PIT-4B-FT-202012") >>> model = AutoModelForCausalLM.from_pretrained( ... "Diamegs/PIT-4B-FT-202012", ... trust_remote_code=True, ... torch_dtype=torch.bfloat16, ... device_map="auto", ... ) """ config_class = PITConfig _no_split_modules = ["Block"] _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, 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() # -- weight tying hooks required by PreTrainedModel ---------------------- def get_input_embeddings(self) -> nn.Embedding: return self.transformer["wte"] def set_input_embeddings(self, value: nn.Embedding) -> None: self.transformer["wte"] = value def get_output_embeddings(self) -> nn.Linear: return self.lm_head 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, cos, sin, attn_mask, is_causal, past_key_values) x = F.rms_norm(x, (x.size(-1),)) # 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: loss = F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_key_values if use_cache else None, )