DZER-Studios
/

Create_Vexion-LM

@@ -1,350 +0,0 @@
-# Copyright 2026 Dmitry
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import math
-import torch.utils.checkpoint as cp
-class GPTConfig:
-    def __init__(self, vocab_size=40960, embed_dim=1024, n_layers=16, n_heads=16,
-                 n_kv_heads=None, intermediate_size=2560, max_seq_len=2048,
-                 dropout=0.0, use_lora=False, num_experts=4, top_k=2, tie_word_embeddings=False, window_size=1024, anchor_size=64):
-        self.vocab_size = vocab_size
-        self.embed_dim = embed_dim
-        self.n_layers = n_layers
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
-        self.intermediate_size = intermediate_size
-        self.max_seq_len = max_seq_len
-        self.dropout = dropout
-        self.use_lora = use_lora
-        self.num_experts = num_experts
-        self.top_k = top_k
-        self.tie_word_embeddings = tie_word_embeddings
-        self.window_size = window_size
-        self.anchor_size = anchor_size
-def precompute_freqs_cis(dim: int, end: int, anchor_size: int = 64, 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.view(1, freqs_cis.shape[0], 1, freqs_cis.shape[1])
-    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 VexNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-    def forward(self, x):
-        abs_mean = x.to(torch.float32).abs().mean(dim=-1, keepdim=True).to(x.dtype)
-        norm_x = x / (abs_mean + self.eps)
-        return norm_x.type_as(x) * self.weight.type_as(x)
-class DoRALinear(nn.Module):
-    def __init__(self, linear_layer, rank=32, alpha=32, dropout=0.05):
-        super().__init__()
-        self.linear = linear_layer
-        in_features = linear_layer.weight.shape[1]
-        out_features = linear_layer.weight.shape[0]
-        self.lora_A = nn.Parameter(torch.zeros(in_features, rank))
-        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))
-        self.scaling = alpha / rank
-        self.dropout = nn.Dropout(dropout)
-        nn.init.normal_(self.lora_A, std=1 / rank)
-        nn.init.zeros_(self.lora_B)
-        self.lora_m = nn.Parameter(self.linear.weight.data.norm(p=2, dim=1, keepdim=True))
-    def forward(self, x):
-        W = self.linear.weight
-        lora_weight = (self.lora_A @ self.lora_B).T * self.scaling
-        W_modified = W + lora_weight
-        norm = W_modified.to(torch.float32).norm(p=2, dim=1, keepdim=True).to(W_modified.dtype)
-        W_dora = self.lora_m * (W_modified / norm)
-        return F.linear(self.dropout(x), W_dora, self.linear.bias)
-class DiffCausalSelfAttention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.n_heads = config.n_heads
-        self.n_embd = config.embed_dim
-        self.head_dim = self.n_embd // self.n_heads
-        self.window_size = getattr(config, 'window_size', 512)
-        self.anchor_size = getattr(config, 'anchor_size', 64)
-        self.c_attn = nn.Linear(self.n_embd, 5 * self.n_embd, bias=False)
-        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
-        self.lambda_noise = nn.Parameter(torch.zeros(self.n_heads, 1, 1))
-        self.diff_norm = VexNorm(self.n_embd)
-    def forward(self, x, freqs_cis=None, use_cache=False, past_kv=None):
-        B, T, C = x.size()
-        qkv = self.c_attn(x)
-        q1, q2, k1, k2, v = qkv.split(self.n_embd, dim=2)
-        q1 = q1.view(B, T, self.n_heads, self.head_dim)
-        q2 = q2.view(B, T, self.n_heads, self.head_dim)
-        k1 = k1.view(B, T, self.n_heads, self.head_dim)
-        k2 = k2.view(B, T, self.n_heads, self.head_dim)
-        v  =  v.view(B, T, self.n_heads, self.head_dim)
-        if freqs_cis is not None:
-            q1, k1 = apply_rotary_emb(q1, k1, freqs_cis)
-            q2, k2 = apply_rotary_emb(q2, k2, freqs_cis)
-        q1, q2 = q1.transpose(1, 2), q2.transpose(1, 2)
-        k1, k2 = k1.transpose(1, 2), k2.transpose(1, 2)
-        v  =  v.transpose(1, 2)
-        if use_cache:
-            if past_kv is not None:
-                past_k1, past_k2, past_v = past_kv
-                k1 = torch.cat([past_k1, k1], dim=2)
-                k2 = torch.cat([past_k2, k2], dim=2)
-                v  = torch.cat([past_v, v], dim=2)
-            past_kv = (k1, k2, v)
-        else:
-            past_kv = None
-        seq_len_kv = k1.size(2)
-        att1 = (q1 @ k1.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
-        att2 = (q2 @ k2.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
-        if seq_len_kv > 1:
-            q_pos = torch.arange(seq_len_kv - T, seq_len_kv, device=x.device).unsqueeze(1)
-            k_pos = torch.arange(seq_len_kv, device=x.device).unsqueeze(0)
-            window_mask = (k_pos >= q_pos - self.window_size) | (k_pos < self.anchor_size)
-            causal_mask = torch.triu(torch.full((T, seq_len_kv), float("-inf"), device=x.device), diagonal=seq_len_kv - T + 1)
-            final_mask = causal_mask.masked_fill(~window_mask, float("-inf"))
-            att1 = att1 + final_mask
-            att2 = att2 + final_mask
-        att1 = F.softmax(att1, dim=-1)
-        att2 = F.softmax(att2, dim=-1)
-        noise_canceller = torch.exp(self.lambda_noise)
-        diff_att = att1 - (noise_canceller * att2)
-        y = diff_att @ v
-        y = y.transpose(1, 2).contiguous().view(B, T, C)
-        y = self.diff_norm(y)
-        return self.c_proj(y), past_kv
-class SwiGLU(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.w1 = nn.Linear(config.embed_dim, config.intermediate_size, bias=False)
-        self.w2 = nn.Linear(config.embed_dim, config.intermediate_size, bias=False)
-        self.w3 = nn.Linear(config.intermediate_size, config.embed_dim, bias=False)
-        self.w3.GPT_SCALE_INIT = True
-        if config.use_lora:
-            self.w1 = DoRALinear(self.w1)
-            self.w2 = DoRALinear(self.w2)
-            self.w3 = DoRALinear(self.w3)
-    def forward(self, x):
-        return self.w3(F.silu(self.w1(x)) * self.w2(x))
-class VexionMoE(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.num_experts = config.num_experts
-        self.top_k = getattr(config, 'top_k', 2)
-        self.experts = nn.ModuleList([SwiGLU(config) for _ in range(self.num_experts)])
-        self.router = nn.Linear(config.embed_dim, self.num_experts, bias=False)
-        self.register_buffer('expert_usage_tracker', torch.zeros(self.num_experts))
-    def forward(self, x):
-        B, T, C = x.shape
-        x_flat = x.view(-1, C)
-        router_logits = self.router(x_flat)
-        routing_weights = F.softmax(router_logits, dim=-1)
-        if self.training:
-            with torch.no_grad():
-                batch_usage = routing_weights.sum(dim=0)
-                self.expert_usage_tracker += batch_usage
-        topk_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
-        tokens_per_expert = torch.bincount(selected_experts.flatten(), minlength=self.num_experts)
-        route_fraction = tokens_per_expert.float() / selected_experts.numel()
-        mean_probs = routing_weights.mean(dim=0)
-        aux_loss = self.num_experts * torch.sum(mean_probs * route_fraction)
-        final_output = torch.zeros_like(x_flat)
-        for i, expert in enumerate(self.experts):
-            expert_mask = (selected_experts == i).any(dim=-1)
-            if not expert_mask.any():
-                continue
-            expert_tokens = x_flat[expert_mask]
-            expert_out = expert(expert_tokens)
-            idx_in_topk = (selected_experts[expert_mask] == i).nonzero(as_tuple=True)[1]
-            token_weights = topk_weights[expert_mask, idx_in_topk].unsqueeze(-1)
-            temp_output = torch.zeros_like(x_flat)
-            temp_output[expert_mask] = expert_out * token_weights
-            final_output = final_output + temp_output
-        return final_output.view(B, T, C), aux_loss
-    @torch.no_grad()
-    def mutate_dead_experts(self, optimizer=None, threshold_ratio=0.01, noise_factor=0.05):
-        total_tokens = self.expert_usage_tracker.sum().item()
-        if total_tokens == 0: return 0
-        best_expert_idx = torch.argmax(self.expert_usage_tracker).item()
-        best_expert = self.experts[best_expert_idx]
-        mutated_count = 0
-        for i in range(self.num_experts):
-            usage_ratio = self.expert_usage_tracker[i].item() / total_tokens
-            if usage_ratio < threshold_ratio and i != best_expert_idx:
-                dead_expert = self.experts[i]
-                for dead_param, best_param in zip(dead_expert.parameters(), best_expert.parameters()):
-                    dead_param.data.copy_(best_param.data)
-                    if optimizer is not None and dead_param in optimizer.state:
-                        del optimizer.state[dead_param]
-                    noise = torch.randn_like(dead_param.data)
-                    scaled_noise = noise * noise_factor * best_param.data.std()
-                    dead_param.data.add_(scaled_noise)
-                self.router.weight.data[i].zero_()
-                mutated_count += 1
-        self.expert_usage_tracker.zero_()
-        return mutated_count
-class Block(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.ln_1 = VexNorm(config.embed_dim)
-        self.attn = DiffCausalSelfAttention(config)
-        self.ln_2 = VexNorm(config.embed_dim)
-        self.mlp = VexionMoE(config)
-    def forward(self, x, freqs_cis, use_cache=False, past_kv=None):
-        attn_out, past_kv_out = self.attn(self.ln_1(x), freqs_cis, use_cache, past_kv)
-        x = x + attn_out
-        mlp_out, aux_loss = self.mlp(self.ln_2(x))
-        x = x + mlp_out
-        return x, aux_loss, past_kv_out
-class GPT(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.transformer = nn.ModuleDict(dict(
-            wte = nn.Embedding(config.vocab_size, config.embed_dim),
-            drop = nn.Dropout(config.dropout),
-            h = nn.ModuleList([Block(config) for _ in range(config.n_layers)]),
-            ln_f = VexNorm(config.embed_dim),
-        ))
-        self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
-        if config.tie_word_embeddings:
-            self.lm_head.weight = self.transformer.wte.weight
-        self.register_buffer(
-            "freqs_cis",
-            precompute_freqs_cis(
-                config.embed_dim // config.n_heads,
-                config.max_seq_len * 2,
-                anchor_size=config.anchor_size
-            ),
-            persistent=False
-        )
-        self.apply(self._init_weights)
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            std = 0.02
-            if hasattr(module, 'GPT_SCALE_INIT') and module.GPT_SCALE_INIT:
-                std *= (2 * self.config.n_layers) ** -0.5
-            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
-            if module.bias is not None:
-                torch.nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
-    def forward(self, idx, targets=None, use_cache=False, past_key_values=None):
-        b, t = idx.size()
-        start_pos = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-        freqs_cis = self.freqs_cis[start_pos : start_pos + t]
-        tok_emb = self.transformer.wte(idx)
-        x = self.transformer.drop(tok_emb)
-        if self.training and not x.requires_grad:
-            x.requires_grad_(True)
-        total_aux_loss = 0.0
-        new_past_key_values = () if use_cache else None
-        for i, block in enumerate(self.transformer.h):
-            past_kv = past_key_values[i] if past_key_values is not None else None
-            if self.training:
-                x, aux_loss, _ = cp.checkpoint(block, x, freqs_cis, False, None, use_reentrant=False)
-            else:
-                x, aux_loss, past_kv_out = block(x, freqs_cis, use_cache, past_kv)
-                if use_cache:
-                    new_past_key_values += (past_kv_out,)
-            total_aux_loss += aux_loss
-        x = self.transformer.ln_f(x)
-        if targets is not None:
-            logits = self.lm_head(x)
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)).float(), targets.view(-1))
-            if self.training:
-                loss = loss + 0.01 * total_aux_loss
-        else:
-            logits = self.lm_head(x[:, [-1], :])
-            loss = None
-        if use_cache:
-            return logits, loss, new_past_key_values
-        return logits, loss