|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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):
|
| input_dtype = x.dtype
|
| x = x.to(torch.float32)
|
| variance = x.abs().mean(-1, keepdim=True)
|
| x = x / (variance + self.eps)
|
| return (x.to(input_dtype) * self.weight)
|
|
|
| class DoRALinear(nn.Module):
|
| def __init__(self, linear_layer, rank=32, alpha=16, 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)
|
|
|
| max_len = config.max_seq_len
|
| mask = torch.tril(torch.ones(max_len, max_len)).view(1, 1, max_len, max_len)
|
| self.register_buffer("causal_mask", mask, persistent=False)
|
|
|
| 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)
|
|
|
| precomputed_causal = self.causal_mask[:, :, :T, :seq_len_kv]
|
| valid_mask = (precomputed_causal > 0.5) & window_mask
|
|
|
| final_mask = torch.zeros_like(att1).masked_fill_(~valid_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).clamp(max=2.0)
|
| 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
|
|
|
