Upload modeling_dolphy.py

modeling_dolphy.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        norm = x.pow(2).mean(-1, keepdim=True)
+        return self.weight * x * torch.rsqrt(norm + self.eps)
+
+class MLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.gelu(self.gate_proj(x)) * self.up_proj(x))
+
+class DolphyBlock(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, num_heads, fused=False):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size)
+        self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
+        self.norm2 = RMSNorm(hidden_size)
+        self.mlp = None if fused else MLP(hidden_size, intermediate_size)
+
+    def forward(self, x, attn_mask=None):
+        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x), attn_mask=attn_mask)[0]
+        if self.mlp:
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+class Dolphy1ForCausalLM(nn.Module):
+    def __init__(self, vocab_size=32000, hidden_size=4096, intermediate_size=16384, num_layers=32, num_heads=32, moe_fused=True):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, hidden_size)
+        self.blocks = nn.ModuleList([
+            DolphyBlock(hidden_size, intermediate_size, num_heads, fused=moe_fused) for _ in range(num_layers)
+        ])
+        self.norm = RMSNorm(hidden_size)
+        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        x = self.embed(input_ids)
+        for block in self.blocks:
+            x = block(x, attention_mask)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        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=None,
+            hidden_states=None,
+            attentions=None,
+        )