Create model.py

model.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from torch.utils.checkpoint import checkpoint
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        mean_square = torch.mean(x ** 2, dim=-1, keepdim=True)
+        normalized_x = x / torch.sqrt(mean_square + self.eps)
+        return self.weight * normalized_x
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        max_len = x.size(1)
+        freqs = torch.arange(0, self.dim // 2, dtype=torch.float32).to(x.device)
+        inv_freq = 1.0 / (10000 ** (freqs / (self.dim // 2)))
+        t = torch.arange(max_len, dtype=torch.float32).to(x.device)
+        sinusoid_inp = torch.outer(t, inv_freq)
+        sin_inp = sinusoid_inp.sin()
+        cos_inp = sinusoid_inp.cos()
+        emb_sin_cos = torch.stack((sin_inp, cos_inp), dim=-1).view(max_len, -1)
+        return x + emb_sin_cos[:max_len, :self.dim].unsqueeze(0)
+
+def apply_rotary_emb(xq, xk, freqs_cis):
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+def reshape_for_broadcast(freqs_cis, x):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+class SwiGLU(nn.Module):
+    def __init__(self, embed_size, expansion_factor=4):
+        super().__init__()
+        self.fc1 = nn.Linear(embed_size, expansion_factor * embed_size)
+        self.fc2 = nn.Linear(expansion_factor * embed_size, embed_size)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.silu(x) * x
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
+
+class SelfAttention(nn.Module):
+    def __init__(self, embed_size, heads):
+        super().__init__()
+        self.embed_size = embed_size
+        self.heads = heads
+        self.head_dim = embed_size // heads
+
+        assert embed_size % heads == 0, "Embed size must be divisible by heads"
+
+        self.values = nn.Linear(embed_size, embed_size, bias=False)
+        self.keys = nn.Linear(embed_size, embed_size, bias=False)
+        self.queries = nn.Linear(embed_size, embed_size, bias=False)
+        self.fc_out = nn.Linear(embed_size, embed_size)
+
+    def forward(self, values, keys, queries, mask=None):
+        N = queries.shape[0]
+        value_len, key_len, query_len = values.shape[1], keys.shape[1], queries.shape[1]
+
+        values = self.values(values).view(N, value_len, self.heads, self.head_dim).transpose(1, 2)
+        keys = self.keys(keys).view(N, key_len, self.heads, self.head_dim).transpose(1, 2)
+        queries = self.queries(queries).view(N, query_len, self.heads, self.head_dim).transpose(1, 2)
+
+        energy = torch.einsum("bthd,bshd->bhts", [queries, keys])
+
+        if mask is not None:
+            energy = energy.masked_fill(mask == 0, float('-inf'))
+
+        attention = torch.softmax(energy / (self.head_dim ** 0.5), dim=-1)
+
+        out = torch.einsum("bhts,bshd->bthd", [attention, values]).transpose(1, 2).reshape(N, query_len, self.embed_size)
+        return self.fc_out(out)
+
+class TransformerBlock(nn.Module):
+    def __init__(self, embed_size, heads, expansion_factor=4, dropout=0.1, checkpoint=False):
+        super().__init__()
+        self.attention = SelfAttention(embed_size, heads)
+        self.feed_forward = SwiGLU(embed_size, expansion_factor)
+        self.norm1 = RMSNorm(embed_size)
+        self.norm2 = RMSNorm(embed_size)
+        self.rotary_pos_emb = RotaryPositionalEmbedding(embed_size)
+        self.checkpoint = checkpoint
+
+    def forward(self, value, mask=None):
+        def forward_fn(value, mask):
+            value = self.rotary_pos_emb(value)
+            attention = self.attention(value, value, value, mask)
+            x = self.norm1(attention + value)
+            forward = self.feed_forward(x)
+            out = self.norm2(forward + x)
+            return out
+
+        if self.checkpoint:
+            return checkpoint(forward_fn, value, mask)
+        else:
+            return forward_fn(value, mask)
+
+class GPT(nn.Module):
+    def __init__(self, vocab_size, embed_size, num_layers, heads, max_length, expansion_factor=4, dropout=0.1, checkpoint=False):
+        super().__init__()
+        self.word_embedding = nn.Embedding(vocab_size, embed_size)
+        self.position_embedding = nn.Embedding(max_length, embed_size)
+
+        self.src_vocab_size = vocab_size
+
+        self.layers = nn.ModuleList(
+            [TransformerBlock(embed_size, heads, expansion_factor, dropout, checkpoint)
+             for _ in range(num_layers)]
+        )
+        self.norm = RMSNorm(embed_size)
+        self.fc_out = nn.Linear(embed_size, vocab_size)
+
+    def forward(self, x, mask=None):
+        positions = torch.arange(0, x.size(1)).unsqueeze(0).to(x.device)
+        x = self.word_embedding(x) + self.position_embedding(positions)
+
+        for layer in self.layers:
+            x = layer(x, mask)
+
+        x = self.norm(x)
+        return self.fc_out(x)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = GPT(vocab_size=10000, embed_size=768, num_layers=20, heads=16, max_length=512, checkpoint=True)
+model.to(device)
+
+inputs = torch.randint(0, 10000, (1, 100), device=device)
+outputs = model(inputs)
+print(outputs.shape)  # Should output: [1, 100, 10000]