Upload train_henyo.py with huggingface_hub

train_henyo.py

@@ -0,0 +1,138 @@
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer, Trainer, TrainingArguments, PreTrainedModel, PretrainedConfig
+from datasets import load_dataset, IterableDataset
+
+# Configuration
+class ModelConfig(PretrainedConfig):
+    model_type = "custom_henyo_culturax"
+    def __init__(
+        self,
+        vocab_size=50257,
+        dim=768,
+        n_layers=12,
+        n_heads=12,
+        n_kv_heads=4,
+        multiple_of=256,
+        max_seq_len=1024,
+        dropout=0.05,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.multiple_of = multiple_of
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+        self.head_dim = dim // n_heads
+
+# Architecture Components
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+    def forward(self, x):
+        return self._norm(x.float()).type_as(x) * self.weight
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+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 = freqs_cis.unsqueeze(0).unsqueeze(0)
+    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)
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, args: ModelConfig):
+        super().__init__()
+        self.n_heads = args.n_heads
+        self.n_kv_heads = args.n_kv_heads
+        self.head_dim = args.head_dim
+        self.n_rep = self.n_heads // args.n_kv_heads
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        self.dropout = nn.Dropout(args.dropout)
+
+    def forward(self, x, freqs_cis, mask=None):
+        b, s, _ = x.shape
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        xq = xq.view(b, s, self.n_heads, self.head_dim).transpose(1, 2)
+        xk = xk.view(b, s, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        xv = xv.view(b, s, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+        if self.n_rep > 1:
+            xk = xk.repeat_interleave(self.n_rep, dim=1)
+            xv = xv.repeat_interleave(self.n_rep, dim=1)
+        output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=mask, dropout_p=self.dropout.p if self.training else 0.0, is_causal=True)
+        return self.wo(output.transpose(1, 2).contiguous().view(b, s, -1))
+
+class SwiGLU(nn.Module):
+    def __init__(self, args: ModelConfig):
+        super().__init__()
+        hidden_dim = 4 * args.dim
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = args.multiple_of * ((hidden_dim + args.multiple_of - 1) // args.multiple_of)
+        self.w1 = nn.Linear(args.dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, args.dim, bias=False)
+        self.w3 = nn.Linear(args.dim, hidden_dim, bias=False)
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelConfig):
+        super().__init__()
+        self.attention_norm = RMSNorm(args.dim)
+        self.attention = GroupedQueryAttention(args)
+        self.ffn_norm = RMSNorm(args.dim)
+        self.feed_forward = SwiGLU(args)
+    def forward(self, x, freqs_cis, mask=None):
+        x = x + self.attention(self.attention_norm(x), freqs_cis, mask)
+        x = x + self.feed_forward(self.ffn_norm(x))
+        return x
+
+class HenyoModel(PreTrainedModel):
+    config_class = ModelConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.norm = RMSNorm(config.dim)
+        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+        self.output.weight = self.tok_embeddings.weight
+        self.freqs_cis = precompute_freqs_cis(config.dim // config.n_heads, config.max_seq_len * 2)
+
+    def forward(self, input_ids, labels=None, **kwargs):
+        b, s = input_ids.shape
+        h = self.tok_embeddings(input_ids)
+        freqs_cis = self.freqs_cis[:s].to(h.device)
+        mask = None
+        if not hasattr(F, 'scaled_dot_product_attention'):
+             mask = torch.triu(torch.full((s, s), float("-inf"), device=h.device), diagonal=1)
+        for layer in self.layers:
+            h = layer(h, freqs_cis, mask)
+        h = self.norm(h)
+        logits = self.output(h)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous().view(-1, self.config.vocab_size)
+            shift_labels = labels[..., 1:].contiguous().view(-1)
+            loss = F.cross_entropy(shift_logits, shift_labels)
+        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}