Update modeling.py

modeling.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import sdpa_kernel, SDPBackend
+
+class RotaryPositionalEncoding(nn.Module):
+    """
+    Rotary Position Embeddings (RoPE) - efficient implementation
+    """
+    def __init__(self, d_head, max_seq_len=8192, base=10000.0):
+        super().__init__()
+        self.d_head = d_head
+        self.max_seq_len = max_seq_len
+        self.base = base
+
+        # Precompute inverse frequencies
+        inv_freq = 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head))
+        self.register_buffer('inv_freq', inv_freq, persistent=False)
+
+        # Precompute cos and sin for maximum sequence length
+        self._precompute_freqs(max_seq_len)
+
+    def _precompute_freqs(self, seq_len):
+        """Precompute cos and sin values for positions"""
+        t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
+        freqs = torch.outer(t, self.inv_freq)  # (seq_len, d_head/2)
+
+        # Create cos and sin embeddings
+        freqs_cos = torch.cos(freqs)
+        freqs_sin = torch.sin(freqs)
+
+        # Interleave to match the dimension (seq_len, d_head)
+        self.register_buffer('freqs_cos', freqs_cos.repeat_interleave(2, dim=-1), persistent=False)
+        self.register_buffer('freqs_sin', freqs_sin.repeat_interleave(2, dim=-1), persistent=False)
+
+    def rotate_half(self, x):
+        """Rotate half the hidden dims of the input"""
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+        return torch.stack([-x2, x1], dim=-1).flatten(-2)
+
+    def forward(self, q, k, start_pos=0):
+        """
+        Apply rotary embeddings to query and key tensors
+        Args:
+            q: (batch_size, n_heads, seq_len, d_head)
+            k: (batch_size, n_heads, seq_len, d_head)
+            start_pos: starting position for caching scenarios
+        Returns:
+            q_rot, k_rot with rotary embeddings applied
+        """
+        seq_len = q.shape[2]
+
+        # Get the precomputed frequencies for this sequence length
+        freqs_cos = self.freqs_cos[start_pos:start_pos + seq_len]
+        freqs_sin = self.freqs_sin[start_pos:start_pos + seq_len]
+
+        # Apply rotary embeddings
+        q_rot = q * freqs_cos + self.rotate_half(q) * freqs_sin
+        k_rot = k * freqs_cos + self.rotate_half(k) * freqs_sin
+
+        return q_rot, k_rot
+
+class Attention(nn.Module):
+    def __init__(self, d_model, n_heads, d_head):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_head
+
+        self.Wq = nn.Linear(d_model, n_heads * d_head, bias=False)
+        self.Wk = nn.Linear(d_model, n_heads * d_head, bias=False)
+        self.Wv = nn.Linear(d_model, n_heads * d_head, bias=False)
+        self.Wo = nn.Linear(n_heads * d_head, d_model, bias=False)
+
+        # Initialize RoPE
+        self.rope = RotaryPositionalEncoding(d_head)
+
+    def forward(self, x):
+        # x is shape batch_size, seq_len, d_model
+        batch_size, seq_len, d_model = x.shape
+        q = self.Wq(x) # q is shape batch_size, seq_len, n_heads * d_head
+        k = self.Wk(x)
+        v = self.Wv(x)
+
+        # reshape to batch_size, n_heads, seq_len, d_head
+        q = q.reshape(batch_size, seq_len, self.n_heads, self.d_head).transpose(1,2)
+        k = k.reshape(batch_size, seq_len, self.n_heads, self.d_head).transpose(1,2)
+        v = v.reshape(batch_size, seq_len, self.n_heads, self.d_head).transpose(1,2)
+
+        q, k = self.rope(q, k)
+        with sdpa_kernel(SDPBackend.FLASH_ATTENTION): # ensure use flash attention
+            a = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True)# a is (batch_size, n_heads, seq_len, d_head)
+        a = a.transpose(1,2) # change a to (batch_size, seq_len, n_heads, d_head)
+        a = a.reshape(batch_size, seq_len, self.n_heads * self.d_head)
+        out = self.Wo(a) # out is (batch_size, seq_len, d_model)
+        return out
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, d_head):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_head
+
+        self.attn = Attention(d_model, n_heads, d_head)
+        self.mlp = nn.Sequential(nn.Linear(d_model, 4*d_model), nn.ReLU(), nn.Linear(4*d_model, d_model))
+
+        self.norm1 = nn.RMSNorm(d_model)
+        self.norm2 = nn.RMSNorm(d_model)
+
+    def forward(self, x):
+        x = self.attn(self.norm1(x)) + x
+        x = self.mlp(self.norm2(x)) + x
+        return x
+
+class GPT(nn.Module):
+    def __init__(self, d_model, n_heads, d_head, n_vocab, n_layers):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_head
+        self.n_vocab = n_vocab
+
+        self.embed = nn.Embedding(n_vocab, d_model)
+
+        self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads, d_head) for _ in range(n_layers)])
+
+        self.norm = nn.RMSNorm(d_model)
+        self.out_head = nn.Linear(d_model, n_vocab)
+
+    def forward(self, x):
+        x = self.embed(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.out_head(self.norm(x))
+        return x
+
+class CustomModel(PreTrainedModel):
+    config_class = CustomConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GPT(config.d_model, config.n_heads, config.d_head, config.n_vocab, config.n_layers)
+
+    def forward(self, tensor):
+        return self.model(tensor)