r4

config.json

@@ -0,0 +1 @@
+{"model_type": "chronogpt", "architectures": ["ChronoGPTForCausalLM"], "auto_map": {"AutoConfig": "modeling_chronogpt.ChronoGPTConfig", "AutoModelForCausalLM": "modeling_chronogpt.ChronoGPTForCausalLM"}, "vocab_size": 50304, "num_layers": 52, "num_heads": 12, "model_dim": 1536}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c8d49269b96100140317ceaa2453ffcc3cfcb45c483f4e578fa2cd4473d4c82
+size 3717113244

modeling_chronogpt.py

@@ -0,0 +1,147 @@
+"""AutoModelForCausalLM-compatible wrapper for ChronoGPT (weights map 1:1 to manelalab/chrono-gpt-v1)."""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class CastedLinear(nn.Linear):
+    def __init__(self, in_features, out_features):
+        super().__init__(in_features, out_features, bias=False)
+    def forward(self, x):
+        return F.linear(x, self.weight.type_as(x))
+
+
+class Rotary(nn.Module):
+    def __init__(self, dim, max_seq_len=65536):
+        super().__init__()
+        angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=dim // 4, dtype=torch.float32)
+        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(dim // 4)])
+        t = torch.arange(max_seq_len, dtype=torch.float32)
+        theta = torch.einsum('i,j -> ij', t, angular_freq)
+        self.register_buffer('cos', theta.cos(), persistent=False)
+        self.register_buffer('sin', theta.sin(), persistent=False)
+    def forward(self, x):
+        cos, sin = self.cos[None, :x.size(-3), None, :], self.sin[None, :x.size(-3), None, :]
+        x1, x2 = x.float().chunk(2, dim=-1)
+        y1 = x1 * cos + x2 * sin
+        y2 = x1 * (-sin) + x2 * cos
+        return torch.cat((y1, y2), 3).type_as(x)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.c_q = CastedLinear(dim, dim)
+        self.c_k = CastedLinear(dim, dim)
+        self.c_v = CastedLinear(dim, dim)
+        self.lambdas = nn.Parameter(torch.tensor([0.5, 0.5]))
+        self.rotary = Rotary(self.head_dim)
+        self.c_proj = CastedLinear(dim, dim)
+    def forward(self, x, ve):
+        B, T = x.size(0), x.size(1)
+        q = self.c_q(x).view(B, T, self.num_heads, self.head_dim)
+        k = self.c_k(x).view(B, T, self.num_heads, self.head_dim)
+        v = self.c_v(x).view(B, T, self.num_heads, self.head_dim)
+        if ve is not None:
+            v = self.lambdas[0] * v + self.lambdas[1] * ve.view_as(v)
+        else:
+            v = self.lambdas[0] * v
+        q, k = norm(q), norm(k)
+        q, k = self.rotary(q), self.rotary(k)
+        y = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.c_proj(y)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.c_fc = CastedLinear(dim, 4 * dim)
+        self.c_proj = CastedLinear(4 * dim, dim)
+    def forward(self, x):
+        return self.c_proj(F.relu(self.c_fc(x)).square())
+
+
+class Block(nn.Module):
+    def __init__(self, model_dim, num_heads, use_attn=True):
+        super().__init__()
+        self.attn = CausalSelfAttention(model_dim, num_heads) if use_attn else None
+        self.mlp = MLP(model_dim)
+        self.lambdas = nn.Parameter(torch.tensor([1., 0.]))
+    def forward(self, x, ve, x0):
+        x = self.lambdas[0] * x + self.lambdas[1] * x0
+        if self.attn is not None:
+            x = x + self.attn(norm(x), ve)
+        x = x + self.mlp(norm(x))
+        return x
+
+
+class ValueEmbedding(nn.Module):
+    def __init__(self, vocab_size, model_dim, num_layers=52):
+        super().__init__()
+        self.num_layers = num_layers
+        self.embed = nn.ModuleList([nn.Embedding(vocab_size, model_dim) for _ in range(3)])
+    def forward(self, inputs):
+        base = [emb(inputs).bfloat16() for emb in self.embed]
+        L = self.num_layers; half = L // 2
+        encoder = [base[i] if i < 3 else None for i in range(half)]
+        decoder = [base[i - (half - 3)] if i >= (half - 3) else None for i in range(half)]
+        return encoder + decoder
+
+
+class ChronoGPTConfig(PretrainedConfig):
+    model_type = "chronogpt"
+    def __init__(self, vocab_size=50304, num_layers=52, num_heads=12, model_dim=1536, **kwargs):
+        self.vocab_size = vocab_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.model_dim = model_dim
+        super().__init__(**kwargs)
+
+
+class ChronoGPTForCausalLM(PreTrainedModel):
+    config_class = ChronoGPTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_heads = config.num_heads
+        self.vocab_size = config.vocab_size
+        self.embed = nn.Embedding(config.vocab_size, config.model_dim)
+        self.blocks = nn.ModuleList([Block(config.model_dim, config.num_heads, use_attn=True) for _ in range(config.num_layers)])
+        self.value_embeds = ValueEmbedding(config.vocab_size, config.model_dim, num_layers=config.num_layers)
+        self.lm_head = CastedLinear(config.model_dim, config.vocab_size)
+        self.num_encoder_layers = config.num_layers // 2
+        self.num_decoder_layers = config.num_layers - self.num_encoder_layers
+        self.skip_weights = nn.Parameter(torch.ones(self.num_decoder_layers))
+
+    @torch.inference_mode()
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)
+        B = input_ids.size(0)
+        x0 = norm(self.embed(input_ids).bfloat16())
+        x = x0
+        ve = [self.value_embeds(input_ids[i].view(-1)) for i in range(B)]
+        ve = [torch.stack([ve[b][i] for b in range(B)]) if ve[0][i] is not None else None for i in range(len(ve[0]))]
+        ve_enc, ve_dec = ve[:self.num_encoder_layers], ve[self.num_encoder_layers:]
+        skip_connections = []
+        for i in range(self.num_encoder_layers):
+            x = self.blocks[i](x, ve_enc[i], x0)
+            skip_connections.append(x)
+        for i in range(self.num_decoder_layers):
+            x = x + self.skip_weights[i] * skip_connections.pop()
+            x = self.blocks[self.num_encoder_layers + i](x, ve_dec[i], x0)
+        x = norm(x)
+        logits = self.lm_head(x)
+        logits = 15 * torch.tanh(logits / 15)
+        return CausalLMOutputWithPast(logits=logits.float())