Create modeling_easyformer.py

modeling_easyformer.py

@@ -0,0 +1,81 @@
+
+import math, torch, torch.nn as nn, torch.nn.functional as F
+from transformers import PreTrainedModel, GenerationMixin
+from configuration_easyformer import EasyFormerConfig
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d)); self.eps = eps
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w
+
+class EasyFormerAttention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.register_buffer("mask", torch.tril(torch.ones(cfg.ctx, cfg.ctx)).bool())
+    def forward(self, x):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(C)
+        att = att.masked_fill(~self.mask[:T, :T], float("-inf"))
+        att = self.drop(F.softmax(att, dim=-1))
+        return self.proj(att @ v)
+
+class EasyFormerFFN(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.fc1 = nn.Linear(cfg.d_model, 2 * cfg.d_model)
+        self.fc2 = nn.Linear(2 * cfg.d_model, cfg.d_model)
+        self.drop = nn.Dropout(cfg.dropout)
+    def forward(self, x):
+        return self.drop(self.fc2(F.relu(self.fc1(x))))
+
+class EasyFormerBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.ln1 = RMSNorm(cfg.d_model); self.attn = EasyFormerAttention(cfg)
+        self.ln2 = RMSNorm(cfg.d_model); self.ffn  = EasyFormerFFN(cfg)
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+class EasyFormerPreTrainedModel(PreTrainedModel):
+    config_class = EasyFormerConfig
+    base_model_prefix = "easyformer"
+    supports_gradient_checkpointing = False
+
+class EasyFormerLMHeadModel(EasyFormerPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        self.cfg = config
+        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
+        self.pos_emb = nn.Embedding(config.ctx, config.d_model)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList(
+            [EasyFormerBlock(config) for _ in range(config.n_layer)]
+        )
+        self.ln_f = RMSNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.tok_emb.weight
+
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        B, T = input_ids.shape
+        pos = torch.arange(T, device=input_ids.device)
+        x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos))
+        for b in self.blocks: x = b(x)
+        logits = self.lm_head(self.ln_f(x))
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.cfg.vocab_size),
+                labels.view(-1), ignore_index=-100,
+            )
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}