Rorical/logos-1b-base
283
1"""Linear (Kimi Delta Attention) decoder-only transformer.2 3Pure-PyTorch chunkwise-parallel KDA scan.4"""5 6import math7from dataclasses import dataclass8from typing import List, Optional, Tuple, Dict, Any9 10import torch11import torch.nn as nn12import torch.nn.functional as F13 14from einops import rearrange15 16from .lm_loss import (17 lm_cross_entropy_from_logits,18 token_superposition_attention_mask,19 token_superposition_embeddings,20)21from .baseline import (22 BaselineConfig,23 RMSNorm,24 SwiGLU,25 MoELayer,26 combine_lm_and_aux_loss,27 init_moe_router_weights,28 _validate_moe_config,29 count_parameters,30 model_summary,31)32 33 34class _ShortConvolution(nn.Module):35 """Causal depthwise 1-D conv with optional cached state for O(1) decode."""36 37 def __init__(38 self,39 hidden_size: int,40 kernel_size: int,41 activation: str = "silu",42 bias: bool = False,43 ):44 super().__init__()45 self.kernel_size = kernel_size46 self.conv = nn.Conv1d(47 hidden_size,48 hidden_size,49 kernel_size=kernel_size,50 groups=hidden_size,51 padding=kernel_size - 1,52 bias=bias,53 )54 self.activation = activation55 56 def forward(57 self,58 x: torch.Tensor,59 cache: Optional[torch.Tensor] = None,60 return_cache: bool = False,61 ):62 T = x.size(1)63 K = self.kernel_size64 65 if cache is None:66 y = self.conv(x.transpose(1, 2))[..., :T].transpose(1, 2)67 else:68 x_full = torch.cat([cache, x], dim=1)69 y = F.conv1d(70 x_full.transpose(1, 2),71 self.conv.weight,72 self.conv.bias,73 stride=1,74 padding=0,75 groups=self.conv.groups,76 ).transpose(1, 2)77 78 if self.activation == "silu":79 y = F.silu(y)80 81 if not return_cache:82 return y83 84 if K <= 1:85 new_cache = x.new_zeros(x.size(0), 0, x.size(-1))86 else:87 combined = torch.cat([cache, x], dim=1) if cache is not None else x88 if combined.size(1) >= K - 1:89 new_cache = combined[:, -(K - 1):].contiguous()90 else:91 pad = combined.new_zeros(92 combined.size(0), (K - 1) - combined.size(1), combined.size(-1)93 )94 new_cache = torch.cat([pad, combined], dim=1)95 return y, new_cache96 97 98class _RMSNormGatedSigmoid(nn.Module):99 def __init__(self, dim: int, eps: float = 1e-5):100 super().__init__()101 self.weight = nn.Parameter(torch.ones(dim))102 self.eps = eps103 104 def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:105 dtype = x.dtype106 x_f = x.float()107 rms_inv = x_f.pow(2).mean(dim=-1, keepdim=True).add_(self.eps).rsqrt()108 y = (x_f * rms_inv).to(dtype) * self.weight109 return y * torch.sigmoid(gate.to(dtype))110 111 112def _kda_gate(113 g: torch.Tensor,114 A_log: torch.Tensor,115 dt_bias: torch.Tensor,116) -> torch.Tensor:117 """Log-space decay gate: ``-exp(A_log) * softplus(g + dt_bias)``."""118 H, K = g.shape[-2], g.shape[-1]119 g = g.float() + dt_bias.float().view(H, K)120 dt = F.softplus(g)121 A = A_log.float().view(1, 1, H, 1)122 return -A.exp() * dt123 124 125def _kda_chunk_scan(126 q: torch.Tensor,127 k: torch.Tensor,128 v: torch.Tensor,129 log_g: torch.Tensor,130 beta: torch.Tensor,131 chunk_size: int = 64,132 use_qk_l2norm: bool = True,133 initial_state: Optional[torch.Tensor] = None,134 output_final_state: bool = False,135):136 """Chunkwise-parallel KDA scan in pure PyTorch.137 138 Recurrence: ``S_i = (I - beta_i k_i k_i^T) D_i S_{i-1} + beta_i k_i v_i^T``,139 ``o_i = q_i @ S_i``, with ``D_i = diag(exp(log_g_i))``. Chunk-level140 parallelism comes from the similarity transform ``~S_i = W_i^{-1} S_i``141 plus a single triangular solve per chunk.142 """143 B, T, H, K = q.shape144 V = v.shape[-1]145 orig_dtype = v.dtype146 device = q.device147 148 # The body runs in fp32: per-channel decays accumulate aggressively and149 # CUDA's triangular_solve has no bf16/fp16 kernel.150 with torch.autocast(device_type=device.type, enabled=False):151 if use_qk_l2norm:152 q = F.normalize(q, dim=-1)153 k = F.normalize(k, dim=-1)154 scale = K ** -0.5155 156 q = q.float() * scale157 k = k.float()158 v = v.float()159 log_g = log_g.float()160 beta = beta.float()161 162 pad = (chunk_size - T % chunk_size) % chunk_size163 if pad > 0:164 q = F.pad(q, (0, 0, 0, 0, 0, pad))165 k = F.pad(k, (0, 0, 0, 0, 0, pad))166 v = F.pad(v, (0, 0, 0, 0, 0, pad))167 log_g = F.pad(log_g, (0, 0, 0, 0, 0, pad))168 beta = F.pad(beta, (0, 0, 0, pad))169 Nc = (T + pad) // chunk_size170 C = chunk_size171 172 q = rearrange(q, "b (n c) h k -> b h n c k", c=C)173 k = rearrange(k, "b (n c) h k -> b h n c k", c=C)174 v = rearrange(v, "b (n c) h v -> b h n c v", c=C)175 log_g = rearrange(log_g, "b (n c) h k -> b h n c k", c=C)176 beta = rearrange(beta, "b (n c) h -> b h n c", c=C)177 178 # Clamp the cumulative log-decay to [-15, 0]: at default A/dt_bias179 # ranges a 64-token cumsum can drop below -80, and exp(-cum) then180 # overflows fp32 and NaNs the triangular solve.181 cum_log_g = log_g.cumsum(dim=-2).clamp(min=-15.0)182 W = cum_log_g.exp()183 W_inv = (-cum_log_g).exp()184 185 u_mat = k * W_inv186 w_mat = k * W187 q_tilde = q * W188 189 beta_e = beta.unsqueeze(-1)190 beta_w = beta_e * w_mat191 beta_v = beta_e * v192 193 L = torch.einsum("bhnik,bhnjk->bhnij", beta_w, u_mat)194 upper_incl_diag = torch.triu(195 torch.ones(C, C, dtype=torch.bool, device=device), diagonal=0196 )197 L = L.masked_fill(upper_incl_diag, 0)198 199 I_plus_L = L + torch.eye(C, dtype=L.dtype, device=device)200 effective_v = torch.linalg.solve_triangular(201 I_plus_L, beta_v, upper=False, unitriangular=True202 )203 effective_w = torch.linalg.solve_triangular(204 I_plus_L, beta_w, upper=False, unitriangular=True205 )206 207 intra_attn = torch.einsum("bhnik,bhnjk->bhnij", q_tilde, u_mat)208 strict_upper = torch.triu(209 torch.ones(C, C, dtype=torch.bool, device=device), diagonal=1210 )211 intra_attn = intra_attn.masked_fill(strict_upper, 0)212 213 if initial_state is not None:214 S = initial_state.to(dtype=q.dtype, device=q.device)215 else:216 S = q.new_zeros(B, H, K, V)217 outputs: List[torch.Tensor] = []218 for n in range(Nc):219 delta = effective_v[:, :, n] - effective_w[:, :, n] @ S220 o_inter = q_tilde[:, :, n] @ S221 o_chunk = o_inter + intra_attn[:, :, n] @ delta222 outputs.append(o_chunk)223 224 state_update = torch.einsum(225 "bhck,bhcv->bhkv", u_mat[:, :, n], delta226 )227 S = W[:, :, n, -1].unsqueeze(-1) * (S + state_update)228 229 out = torch.stack(outputs, dim=2)230 out = rearrange(out, "b h n c v -> b (n c) h v")231 if pad > 0:232 out = out[:, :T]233 out = out.to(orig_dtype)234 235 if output_final_state:236 # State stays fp32 so cached decode preserves precision.237 return out, S238 return out239 240 241def _kda_recurrent_step(242 q: torch.Tensor,243 k: torch.Tensor,244 v: torch.Tensor,245 log_g: torch.Tensor,246 beta: torch.Tensor,247 state: torch.Tensor,248 use_qk_l2norm: bool = True,249) -> Tuple[torch.Tensor, torch.Tensor]:250 """Single-token KDA step matching ``_kda_chunk_scan`` for ``T == 1``."""251 assert q.size(1) == 1 and k.size(1) == 1 and v.size(1) == 1252 orig_dtype = v.dtype253 K = q.size(-1)254 device = q.device255 256 with torch.autocast(device_type=device.type, enabled=False):257 if use_qk_l2norm:258 q = F.normalize(q, dim=-1)259 k = F.normalize(k, dim=-1)260 scale = K ** -0.5261 262 q_t = (q[:, 0].float()) * scale263 k_t = k[:, 0].float()264 v_t = v[:, 0].float()265 g_t = log_g[:, 0].float()266 b_t = beta[:, 0].float()267 268 S = state.to(torch.float32)269 S = S * g_t.exp().unsqueeze(-1)270 kS = torch.einsum("bhk,bhkv->bhv", k_t, S)271 update = torch.einsum(272 "bhk,bhv->bhkv", (b_t.unsqueeze(-1) * k_t), (v_t - kS)273 )274 S = S + update275 o = torch.einsum("bhk,bhkv->bhv", q_t, S).unsqueeze(1)276 return o.to(orig_dtype), S277 278 279@dataclass280class LinearConfig(BaselineConfig):281 head_dim: int = 64282 conv_size: int = 4283 chunk_size: int = 64284 A_init_range: Tuple[float, float] = (1, 16)285 286 expand: int = 2287 rope_base: float = 10000.0288 289 def __post_init__(self):290 if self.d_model % self.num_heads != 0:291 raise ValueError("d_model must be divisible by num_heads")292 if self.partial_rope_dim is not None:293 if self.partial_rope_dim % 2 != 0:294 raise ValueError(295 f"partial_rope_dim ({self.partial_rope_dim}) must be even"296 )297 if self.head_dim < 1:298 raise ValueError("head_dim must be >= 1")299 if self.chunk_size < 1:300 raise ValueError("chunk_size must be >= 1")301 if self.conv_size < 1:302 raise ValueError("conv_size must be >= 1")303 _validate_moe_config(self)304 305 306class KimiDeltaAttention(nn.Module):307 def __init__(self, config: LinearConfig):308 super().__init__()309 self.hidden_size = config.d_model310 self.num_heads = config.num_heads311 self.head_dim = config.head_dim312 self.head_k_dim = self.head_dim313 self.conv_size = config.conv_size314 self.chunk_size = config.chunk_size315 316 projection_size = self.num_heads * self.head_dim317 318 self.q_proj = nn.Linear(self.hidden_size, projection_size, bias=False)319 self.k_proj = nn.Linear(self.hidden_size, projection_size, bias=False)320 self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)321 322 self.q_conv1d = _ShortConvolution(projection_size, self.conv_size, "silu")323 self.k_conv1d = _ShortConvolution(projection_size, self.conv_size, "silu")324 self.v_conv1d = _ShortConvolution(projection_size, self.conv_size, "silu")325 326 A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(327 *config.A_init_range328 )329 self.A_log = nn.Parameter(torch.log(A))330 self.A_log._no_weight_decay = True331 332 self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32))333 self.dt_bias._no_weight_decay = True334 335 self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)336 self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)337 338 self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)339 340 self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)341 self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=True)342 343 self.o_norm = _RMSNormGatedSigmoid(self.head_dim, eps=config.norm_eps)344 self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)345 346 self._reset_parameters()347 348 def _reset_parameters(self):349 # Inverse-softplus init (Mamba-2 / KDA scheme).350 dt = torch.exp(351 torch.rand(self.num_heads * self.head_dim)352 * (math.log(0.1) - math.log(0.001))353 + math.log(0.001)354 )355 dt = torch.clamp(dt, min=1e-4)356 inv_dt = dt + torch.log(-torch.expm1(-dt))357 with torch.no_grad():358 self.dt_bias.copy_(inv_dt)359 360 def forward(361 self,362 x: torch.Tensor,363 attention_mask: Optional[torch.Tensor] = None,364 cache: Optional[Dict[str, Optional[torch.Tensor]]] = None,365 ) -> torch.Tensor:366 use_cache = cache is not None367 368 q_in = self.q_proj(x)369 k_in = self.k_proj(x)370 v_in = self.v_proj(x)371 if use_cache:372 q, cache["conv_state_q"] = self.q_conv1d(373 q_in, cache=cache.get("conv_state_q"), return_cache=True374 )375 k, cache["conv_state_k"] = self.k_conv1d(376 k_in, cache=cache.get("conv_state_k"), return_cache=True377 )378 v, cache["conv_state_v"] = self.v_conv1d(379 v_in, cache=cache.get("conv_state_v"), return_cache=True380 )381 else:382 q = self.q_conv1d(q_in)383 k = self.k_conv1d(k_in)384 v = self.v_conv1d(v_in)385 386 q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)387 k = rearrange(k, "... (h d) -> ... h d", d=self.head_dim)388 v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim)389 390 g_raw = self.f_b_proj(self.f_a_proj(x))391 g_raw = rearrange(g_raw, "... (h d) -> ... h d", d=self.head_dim)392 log_g = _kda_gate(g_raw, self.A_log, self.dt_bias)393 394 beta = self.b_proj(x).float().sigmoid()395 396 # Zero q/k/v, log_g, beta at padded positions so they contribute no397 # content and no decay to the recurrent state.398 if attention_mask is not None:399 mask_4d = attention_mask.unsqueeze(-1).unsqueeze(-1)400 q = q * mask_4d.to(q.dtype)401 k = k * mask_4d.to(k.dtype)402 v = v * mask_4d.to(v.dtype)403 log_g = log_g * mask_4d.to(log_g.dtype)404 beta = beta * attention_mask.unsqueeze(-1).to(beta.dtype)405 406 if use_cache:407 prev_state = cache.get("recurrent_state")408 if prev_state is not None and x.size(1) == 1:409 o, new_state = _kda_recurrent_step(410 q, k, v, log_g, beta, prev_state, use_qk_l2norm=True411 )412 else:413 o, new_state = _kda_chunk_scan(414 q=q, k=k, v=v, log_g=log_g, beta=beta,415 chunk_size=self.chunk_size,416 use_qk_l2norm=True,417 initial_state=prev_state,418 output_final_state=True,419 )420 cache["recurrent_state"] = new_state421 else:422 o = _kda_chunk_scan(423 q=q, k=k, v=v, log_g=log_g, beta=beta,424 chunk_size=self.chunk_size,425 use_qk_l2norm=True,426 )427 428 gate = self.g_b_proj(self.g_a_proj(x))429 gate = rearrange(gate, "... (h d) -> ... h d", d=self.head_dim)430 o = self.o_norm(o, gate)431 432 o = rearrange(o, "b t h d -> b t (h d)")433 return self.o_proj(o)434 435 436class LinearTransformerBlock(nn.Module):437 def __init__(self, config: LinearConfig):438 super().__init__()439 self.use_moe = config.use_moe440 441 self.kda_norm = RMSNorm(config.d_model, eps=config.norm_eps)442 self.kda = KimiDeltaAttention(config)443 444 self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)445 if config.use_moe:446 self.ffn = MoELayer(config)447 else:448 self.ffn = SwiGLU(config.d_model, config.d_ff)449 450 def forward(451 self,452 x: torch.Tensor,453 attention_mask: Optional[torch.Tensor] = None,454 is_causal: bool = True,455 cache: Optional[Dict[str, Optional[torch.Tensor]]] = None,456 ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:457 x = x + self.kda(self.kda_norm(x), attention_mask=attention_mask, cache=cache)458 459 if self.use_moe:460 ffn_out, aux_loss, topk_indices = self.ffn(self.ffn_norm(x))461 x = x + ffn_out462 return x, aux_loss, topk_indices463 else:464 x = x + self.ffn(self.ffn_norm(x))465 return x, torch.zeros((), device=x.device, dtype=x.dtype), None466 467 468class LinearTransformer(nn.Module):469 def __init__(self, config: LinearConfig):470 super().__init__()471 self.config = config472 473 self.token_emb = nn.Embedding(config.vocab_size, config.d_model)474 475 self.layers = nn.ModuleList([476 LinearTransformerBlock(config) for _ in range(config.num_layers)477 ])478 479 self.final_norm = RMSNorm(config.d_model, eps=config.norm_eps)480 self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)481 self.lm_head.weight = self.token_emb.weight482 483 self._init_weights()484 485 def _init_weights(self):486 for module in self.modules():487 if isinstance(module, nn.Linear):488 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)489 if module.bias is not None:490 torch.nn.init.zeros_(module.bias)491 elif isinstance(module, nn.Embedding):492 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)493 init_moe_router_weights(self, self.config.router_init_std)494 495 def forward(496 self,497 input_ids: torch.Tensor,498 attention_mask: Optional[torch.Tensor] = None,499 labels: Optional[torch.Tensor] = None,500 is_causal: bool = True,501 caches: Optional[List[Dict[str, Optional[torch.Tensor]]]] = None,502 token_superposition_bag_size: int = 1,503 ) -> Dict[str, Any]:504 x = token_superposition_embeddings(505 self.token_emb, input_ids, token_superposition_bag_size,506 )507 attention_mask = token_superposition_attention_mask(508 attention_mask, token_superposition_bag_size,509 )510 511 aux_loss = torch.zeros((), device=input_ids.device, dtype=x.dtype)512 topk_indices_list: List[Optional[torch.Tensor]] = []513 for i, layer in enumerate(self.layers):514 layer_cache = caches[i] if caches is not None else None515 x, layer_aux, layer_topk = layer(516 x, attention_mask=attention_mask,517 is_causal=is_causal, cache=layer_cache,518 )519 aux_loss = aux_loss + layer_aux520 topk_indices_list.append(layer_topk)521 522 x = self.final_norm(x)523 logits = self.lm_head(x)524 525 lm_loss: Optional[torch.Tensor] = None526 if labels is not None:527 lm_loss = lm_cross_entropy_from_logits(528 logits,529 labels,530 token_superposition_bag_size=token_superposition_bag_size,531 ignore_index=-100,532 )533 loss = combine_lm_and_aux_loss(534 lm_loss,535 aux_loss if self.config.use_moe else None,536 self.training,537 )538 539 return {540 "logits": logits,541 "loss": loss,542 "lm_loss": lm_loss,543 "aux_loss": aux_loss if self.config.use_moe else None,544 "topk_indices": topk_indices_list if self.config.use_moe else None,545 }546 547 def update_router_biases(self, topk_indices_list: List[Optional[torch.Tensor]]) -> None:548 if not self.config.use_moe:549 return550 for layer, topk_indices in zip(self.layers, topk_indices_list):551 if topk_indices is not None and isinstance(layer.ffn, MoELayer):552 layer.ffn.update_bias(topk_indices)553 554 @torch.no_grad()555 def get_balance_stats(self) -> Dict[str, float]:556 if not self.config.use_moe:557 return {}558 stats = {}559 for idx, layer in enumerate(self.layers):560 if hasattr(layer.ffn, "bias"):561 bias = layer.ffn.bias562 stats[f"layer{idx}_bias_mean"] = bias.abs().mean().item()563 stats[f"layer{idx}_bias_max"] = bias.abs().max().item()564 return stats565 566 @torch.no_grad()567 def generate(568 self,569 input_ids: torch.Tensor,570 max_new_tokens: int = 100,571 temperature: float = 1.0,572 top_k: Optional[int] = None,573 attention_mask: Optional[torch.Tensor] = None,574 eos_token_id: Optional[int] = None,575 ) -> torch.Tensor:576 self.train(False)577 578 caches: List[Dict[str, Optional[torch.Tensor]]] = [579 {580 "recurrent_state": None,581 "conv_state_q": None,582 "conv_state_k": None,583 "conv_state_v": None,584 }585 for _ in self.layers586 ]587 588 def _sample(logits: torch.Tensor) -> torch.Tensor:589 logits = logits / temperature590 if top_k is not None:591 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))592 logits = logits.masked_fill(logits < v[:, [-1]], float("-inf"))593 probs = F.softmax(logits, dim=-1)594 return torch.multinomial(probs, num_samples=1)595 596 outputs = self.forward(input_ids, is_causal=True, caches=caches)597 next_token = _sample(outputs["logits"][:, -1, :])598 input_ids = torch.cat([input_ids, next_token], dim=-1)599 600 if eos_token_id is not None and (next_token == eos_token_id).all():601 return input_ids602 603 for _ in range(max_new_tokens - 1):604 outputs = self.forward(next_token, is_causal=True, caches=caches)605 next_token = _sample(outputs["logits"][:, -1, :])606 input_ids = torch.cat([input_ids, next_token], dim=-1)607 if eos_token_id is not None and (next_token == eos_token_id).all():608 break609 610 return input_ids611