Create modeling_snowflake_core.py

modeling_snowflake_core.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from typing import Optional, Tuple
+
+# Optional: import custom config if present
+try:
+    from .configuration_snowflake_core import SnowflakeCoreConfig
+except ImportError:
+    SnowflakeCoreConfig = PretrainedConfig
+
+class FusedSelfAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+    def forward(self, x, attn_mask=None, key_padding_mask=None):
+        B, T, C = x.size()
+        qkv = self.qkv_proj(x)  # [B, T, 3 * C]
+        qkv = qkv.reshape(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # Each: [B, num_heads, T, head_dim]
+
+        attn_scores = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)  # [B, num_heads, T, T]
+        if attn_mask is not None:
+            attn_scores = attn_scores + attn_mask.unsqueeze(0).unsqueeze(0).to(attn_scores.dtype)
+        if key_padding_mask is not None:
+            attn_scores = attn_scores.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+        attn_probs = F.softmax(attn_scores, dim=-1)
+        attn_output = attn_probs @ v  # [B, num_heads, T, head_dim]
+        attn_output = attn_output.transpose(1, 2).reshape(B, T, C)
+        return self.out_proj(attn_output)
+
+class GPTBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.1):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(embed_dim)
+        self.attn = FusedSelfAttention(embed_dim, num_heads)
+        self.dropout1 = nn.Dropout(dropout)
+        self.ln2 = nn.LayerNorm(embed_dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, 4 * embed_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(4 * embed_dim, embed_dim),
+        )
+        self.dropout2 = nn.Dropout(dropout)
+    def forward(self, x, attn_mask=None, key_padding_mask=None):
+        h = self.ln1(x)
+        attn_output = self.attn(h, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
+        x = x + self.dropout1(attn_output)
+        x = x + self.dropout2(self.mlp(self.ln2(x)))
+        return x
+
+class SnowflakeCoreG1(PreTrainedModel):
+    config_class = SnowflakeCoreConfig
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.vocab_size = config.vocab_size
+        self.embed_dim = config.embed_dim
+        self.num_heads = config.num_heads
+        self.num_layers = config.num_layers
+        self.max_length = config.max_length
+        self.ffn_dim = getattr(config, 'ffn_dim', 4 * config.embed_dim)
+        self.dropout = getattr(config, 'dropout', 0.1)
+
+        self.embed = nn.Embedding(self.vocab_size, self.embed_dim)
+        self.pos_embed = nn.Embedding(self.max_length, self.embed_dim)
+        self.dropout_layer = nn.Dropout(self.dropout)
+        self.blocks = nn.ModuleList([
+            GPTBlock(self.embed_dim, self.num_heads, self.dropout) for _ in range(self.num_layers)
+        ])
+        self.ln_f = nn.LayerNorm(self.embed_dim)
+        self.lm_head = nn.Linear(self.embed_dim, self.vocab_size, bias=False)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value):
+        self.embed = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs
+    ) -> Tuple:
+        B, T = input_ids.size()
+        pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
+        x = self.embed(input_ids) + self.pos_embed(pos)
+        x = self.dropout_layer(x)
+        causal_mask = torch.triu(torch.ones(T, T, device=input_ids.device), diagonal=1).bool()
+        causal_mask = causal_mask.masked_fill(causal_mask, float('-inf'))
+        key_padding_mask = None
+        if attention_mask is not None:
+            key_padding_mask = attention_mask == 0
+        for block in self.blocks:
+            x = block(x, attn_mask=causal_mask, key_padding_mask=key_padding_mask)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous().view(-1, self.vocab_size)
+            shift_labels = labels[:, 1:].contiguous().view(-1)
+            loss = F.cross_entropy(shift_logits, shift_labels, ignore_index=self.config.pad_token_id)
+        if loss is not None:
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, **kwargs):
+        return super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)