Release MarsNext Phobos-1 (100% 35-Dataset Base Step 33977 + 60-Book 50/50 Replay Post-Trained)
07d2bfc verified | #!/usr/bin/env python3 | |
| """ | |
| SynapCore Phobos-Scratch v1 — Pure Causal Architecture + DeepInfer (DeepSeek-V4.1) Integration | |
| (`phobos-1/models/phobos_scratch_v1/architecture.py`) | |
| ============================================================================================== | |
| 100% From-Scratch Causal Language Model combining: | |
| 1. Exact 1D Rotary Position Embeddings (RoPE) + Multi-Token Causal QK Horizon (`W=8`). | |
| 2. Strict Causal Grouped-Query FlashAttention (`is_causal=True`, zero future-token leak). | |
| 3. Token-Level Hierarchical Cortical MoE (`Shared E00 + 4 Branches -> 8 Specialist Leaves`) | |
| with Multi-Scale Causal Router Horizon (`W=16`) + DeepInfer `TopKGate` (`sqrtsoftplus` | |
| scoring & `noaux_tc` auxiliary-loss-free bias balancing). | |
| 4. DeepInfer `Engram` Prime-Modulo 2-Gram/3-Gram Hash Memory Tables on even layers. | |
| 5. DeepInfer `mHC` Multi-Head Hyper-Connections (`hc_mult = 4` parallel residual streams). | |
| 6. DeepInfer `MarkovGPTHead` + `DFlash` 2-Step Multi-Token Prediction (`MTP`) & `ConfidenceHead`. | |
| 7. DeepInfer `XGrammar` Step-Wise Structured Output Bitmasking during `generate()`. | |
| Strictly <= 475 lines. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import asdict, dataclass | |
| import math | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from deepinfer_modules import ( | |
| DeepInferEngramMemory, | |
| DeepInferHyperConnection, | |
| DeepInferMarkovMTPHead, | |
| DeepInferNoAuxGate, | |
| DeepInferStructuredMasker, | |
| ) | |
| SCRATCH_WEIGHTS_DIR: Path = Path(__file__).resolve().parent / "weights" | |
| class PhobosScratchConfig: | |
| """Configuration for the From-Scratch SynapCore Phobos + DeepInfer causal language model.""" | |
| vocab_size: int = 32768 | |
| d_model: int = 512 | |
| num_layers: int = 8 | |
| num_heads: int = 8 | |
| num_kv_heads: int = 4 | |
| max_seq_len: int = 2048 | |
| shared_mlp_neurons: int = 1024 | |
| leaf_mlp_neurons: int = 256 | |
| num_branches: int = 4 | |
| leaves_per_branch: int = 2 | |
| top_k_leaves: int = 2 | |
| rope_theta: float = 10000.0 | |
| rms_norm_eps: float = 1e-6 | |
| router_temp: float = 0.15 | |
| aux_loss_coef: float = 0.02 | |
| causal_qk_window: int = 8 | |
| router_context_window: int = 16 | |
| tie_word_embeddings: bool = True | |
| use_deepinfer_noaux_gate: bool = True | |
| use_deepinfer_engram: bool = True | |
| use_deepinfer_hc: bool = True | |
| hc_mult: int = 4 | |
| use_deepinfer_markov_mtp: bool = True | |
| markov_rank: int = 64 | |
| use_deepinfer_structured_mask: bool = True | |
| def num_leaves(self) -> int: return self.num_branches * self.leaves_per_branch | |
| def head_dim(self) -> int: return self.d_model // self.num_heads | |
| class RMSNorm(nn.Module): | |
| """Root Mean Square Layer Normalization.""" | |
| def __init__(self, dim: int, eps: float = 1e-6) -> None: | |
| super().__init__() | |
| self.eps, self.weight = eps, nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return (x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps).to(x.dtype)) * self.weight | |
| class RotaryPositionEmbedding(nn.Module): | |
| """Exact Rotary Position Embedding (RoPE) rotating Query and Key vectors in 2D subspaces.""" | |
| def __init__(self, head_dim: int, max_seq_len: int = 2048, theta: float = 10000.0) -> None: | |
| super().__init__() | |
| self.head_dim = head_dim | |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) | |
| self.register_buffer("inv_freq", inv_freq, persistent=False) | |
| self._build_cache(max_seq_len) | |
| def _build_cache(self, seq_len: int) -> None: | |
| freqs = torch.outer(torch.arange(seq_len, dtype=torch.float32, device=self.inv_freq.device), self.inv_freq) | |
| emb = torch.cat((freqs, freqs), dim=-1) | |
| self.register_buffer("cos_cached", emb.cos(), persistent=False) | |
| self.register_buffer("sin_cached", emb.sin(), persistent=False) | |
| def _rotate_half(x: torch.Tensor) -> torch.Tensor: | |
| half = x.shape[-1] // 2 | |
| return torch.cat((-x[..., half:], x[..., :half]), dim=-1) | |
| def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| seq_len = q.shape[2] | |
| if seq_len > self.cos_cached.shape[0]: | |
| self._build_cache(seq_len) | |
| cos = self.cos_cached[:seq_len].to(q.dtype).unsqueeze(0).unsqueeze(0) | |
| sin = self.sin_cached[:seq_len].to(q.dtype).unsqueeze(0).unsqueeze(0) | |
| return (q * cos) + (self._rotate_half(q) * sin), (k * cos) + (self._rotate_half(k) * sin) | |
| class CausalGroupedQueryAttention(nn.Module): | |
| """Strict Causal Grouped-Query Attention (GQA) with Causal Multi-Token QK Horizon.""" | |
| def __init__(self, cfg: PhobosScratchConfig) -> None: | |
| super().__init__() | |
| self.num_heads = cfg.num_heads | |
| self.num_kv_heads = cfg.num_kv_heads | |
| self.num_kv_groups = cfg.num_heads // cfg.num_kv_heads | |
| self.head_dim = cfg.head_dim | |
| self.qk_window = max(1, int(cfg.causal_qk_window)) | |
| self.qk_horizon_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) | |
| nn.init.normal_(self.qk_horizon_proj.weight, mean=0.0, std=0.02 / math.sqrt(cfg.d_model)) | |
| self.q_proj = nn.Linear(cfg.d_model, cfg.num_heads * self.head_dim, bias=False) | |
| self.k_proj = nn.Linear(cfg.d_model, cfg.num_kv_heads * self.head_dim, bias=False) | |
| self.v_proj = nn.Linear(cfg.d_model, cfg.num_kv_heads * self.head_dim, bias=False) | |
| self.o_proj = nn.Linear(cfg.num_heads * self.head_dim, cfg.d_model, bias=False) | |
| self.rope = RotaryPositionEmbedding(self.head_dim, cfg.max_seq_len, cfg.rope_theta) | |
| self._qk_cache: Dict[Tuple[int, str], torch.Tensor] = {} | |
| def _causal_sliding_mean(self, x: torch.Tensor) -> torch.Tensor: | |
| _, seq_len, _ = x.shape | |
| key = (seq_len, str(x.device)) | |
| with torch.amp.autocast(device_type=x.device.type, enabled=False): | |
| weights = self._qk_cache.get(key) | |
| if weights is None: | |
| idx = torch.arange(seq_len, device=x.device) | |
| diff = idx.unsqueeze(1) - idx.unsqueeze(0) | |
| mask = ((diff >= 0) & (diff < self.qk_window)).float() | |
| weights = (mask / mask.sum(dim=-1, keepdim=True).clamp(min=1.0)).unsqueeze(0) | |
| self._qk_cache[key] = weights | |
| out = torch.matmul(weights, x.float()) | |
| return out.to(x.dtype) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| bsz, seq_len, _ = x.shape | |
| x_qk = x + 0.35 * self.qk_horizon_proj(self._causal_sliding_mean(x)) if self.qk_window > 1 else x | |
| q = self.q_proj(x_qk).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) | |
| k = self.k_proj(x_qk).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) | |
| v = self.v_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) | |
| q, k = self.rope(q, k) | |
| if self.num_kv_groups > 1: | |
| k = k.repeat_interleave(self.num_kv_groups, dim=1) | |
| v = v.repeat_interleave(self.num_kv_groups, dim=1) | |
| attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True) | |
| return self.o_proj(attn_out.transpose(1, 2).contiguous().view(bsz, seq_len, -1)) | |
| class SwiGLUExpertMLP(nn.Module): | |
| """Single SwiGLU MLP expert (`w_down(silu(w_gate(x)) * w_up(x))`).""" | |
| def __init__(self, d_model: int, num_neurons: int) -> None: | |
| super().__init__() | |
| self.d_model = d_model | |
| self.num_neurons = num_neurons | |
| self.w_gate = nn.Linear(d_model, num_neurons, bias=False) | |
| self.w_up = nn.Linear(d_model, num_neurons, bias=False) | |
| self.w_down = nn.Linear(num_neurons, d_model, bias=False) | |
| nn.init.normal_(self.w_gate.weight, mean=0.0, std=0.02) | |
| nn.init.normal_(self.w_up.weight, mean=0.0, std=0.02) | |
| nn.init.normal_(self.w_down.weight, mean=0.0, std=0.02 / math.sqrt(max(1, num_neurons))) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) | |
| class HierarchicalCorticalMoELayer(nn.Module): | |
| """Token-Level Hierarchical Cortical MoE with DeepInfer `TopKGate` (`sqrtsoftplus` + `noaux_tc`).""" | |
| def __init__(self, cfg: PhobosScratchConfig) -> None: | |
| super().__init__() | |
| self.d_model = cfg.d_model | |
| self.num_branches = cfg.num_branches | |
| self.leaves_per_branch = cfg.leaves_per_branch | |
| self.top_k = min(cfg.top_k_leaves, cfg.num_leaves) | |
| self.router_temp = cfg.router_temp | |
| self.router_window = max(1, int(cfg.router_context_window)) | |
| self.shared_e00 = SwiGLUExpertMLP(cfg.d_model, cfg.shared_mlp_neurons) | |
| self.branch_centroids = nn.Parameter(F.normalize(torch.randn(cfg.num_branches, cfg.d_model), p=2, dim=-1)) | |
| self.leaf_centroids = nn.Parameter(F.normalize(torch.randn(cfg.num_leaves, cfg.d_model), p=2, dim=-1)) | |
| self.leaf_parent_branch = nn.Buffer( | |
| torch.arange(cfg.num_leaves, dtype=torch.long) // cfg.leaves_per_branch, | |
| persistent=True, | |
| ) | |
| self.leaf_origin_id = nn.Buffer(torch.arange(cfg.num_leaves, dtype=torch.long), persistent=True) | |
| self.leaves = nn.ModuleList([SwiGLUExpertMLP(cfg.d_model, cfg.leaf_mlp_neurons) for _ in range(cfg.num_leaves)]) | |
| self.noaux_gate = DeepInferNoAuxGate(cfg.num_branches, cfg.num_leaves) if cfg.use_deepinfer_noaux_gate else None | |
| self.register_buffer("leaf_token_ema", torch.ones(cfg.num_leaves) / float(cfg.num_leaves)) | |
| self.register_buffer("leaf_dispersion_ema", torch.zeros(cfg.num_leaves)) | |
| self._router_cache: Dict[Tuple[int, str], torch.Tensor] = {} | |
| def num_leaves(self) -> int: | |
| return len(self.leaves) | |
| def build_causal_router_context(self, x: torch.Tensor) -> torch.Tensor: | |
| bsz, seq_len, dim = x.shape | |
| key = (seq_len, str(x.device)) | |
| with torch.amp.autocast(device_type=x.device.type, enabled=False): | |
| r_mat = self._router_cache.get(key) | |
| if r_mat is None: | |
| idx = torch.arange(seq_len, device=x.device) | |
| diff = idx.unsqueeze(1) - idx.unsqueeze(0) | |
| eye_m = (diff == 0).float() | |
| win_m = ((diff >= 0) & (diff < self.router_window)).float() | |
| win_m = win_m / win_m.sum(dim=-1, keepdim=True).clamp(min=1.0) | |
| pre_m = (diff >= 0).float() | |
| pre_m = pre_m / pre_m.sum(dim=-1, keepdim=True).clamp(min=1.0) | |
| r_mat = (0.25 * eye_m + 0.35 * win_m + 0.40 * pre_m).unsqueeze(0) | |
| self._router_cache[key] = r_mat | |
| r = torch.matmul(r_mat, x.float()) | |
| return r.to(x.dtype).reshape(bsz * seq_len, dim) | |
| def compute_hierarchical_routing( | |
| self, | |
| x_flat: torch.Tensor, | |
| target_branches: Optional[torch.Tensor] = None, | |
| target_leaves: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, float, float]: | |
| x_unit = F.normalize(x_flat.float(), p=2, dim=-1) | |
| b_unit = F.normalize(self.branch_centroids.float(), p=2, dim=-1) | |
| l_unit = F.normalize(self.leaf_centroids.float(), p=2, dim=-1) | |
| branch_logits = (x_unit @ b_unit.T) / self.router_temp | |
| leaf_raw_logits = (x_unit @ l_unit.T) / self.router_temp | |
| combined_logits = 0.5 * branch_logits[:, self.leaf_parent_branch] + 0.5 * leaf_raw_logits | |
| if self.noaux_gate is not None: | |
| _, probs, top_k_idx, top_k_weights = self.noaux_gate.score_and_balance( | |
| branch_logits, combined_logits, top_k_leaves=self.top_k | |
| ) | |
| else: | |
| probs = F.softmax(combined_logits, dim=-1) | |
| top_k_vals, top_k_idx = torch.topk(probs, k=self.top_k, dim=-1) | |
| top_k_weights = top_k_vals / top_k_vals.sum(dim=-1, keepdim=True).clamp(min=1e-8) | |
| n_leaves = self.num_leaves | |
| frac_tokens = F.one_hot(top_k_idx[:, 0], num_classes=n_leaves).float().mean(dim=0) | |
| aux_loss = float(n_leaves) * torch.sum(frac_tokens * probs.mean(dim=0)) | |
| sup_loss = torch.tensor(0.0, device=x_flat.device) | |
| branch_acc, leaf_acc = 0.0, 0.0 | |
| if target_branches is not None and target_leaves is not None: | |
| valid = target_branches != -100 | |
| if valid.any(): | |
| if n_leaves == 8: | |
| orig_logits = combined_logits | |
| else: | |
| orig_cols = [combined_logits[:, self.leaf_origin_id == oid].max(dim=-1).values for oid in range(8)] | |
| orig_logits = torch.stack(orig_cols, dim=-1) | |
| b_loss = F.cross_entropy(branch_logits, target_branches, ignore_index=-100) | |
| l_loss = F.cross_entropy(orig_logits, target_leaves, ignore_index=-100) | |
| sup_loss = 0.5 * b_loss + 0.5 * l_loss | |
| with torch.no_grad(): | |
| branch_acc = float((branch_logits.argmax(dim=-1)[valid] == target_branches[valid]).float().mean().item()) | |
| leaf_acc = float((orig_logits.argmax(dim=-1)[valid] == target_leaves[valid]).float().mean().item()) | |
| if self.training: | |
| with torch.no_grad(): | |
| self.leaf_token_ema.mul_(0.95).add_(frac_tokens * 0.05) | |
| for l_idx in range(n_leaves): | |
| mask = top_k_idx[:, 0] == l_idx | |
| if mask.any(): | |
| disp = 1.0 - (x_unit[mask] @ l_unit[l_idx]).mean() | |
| self.leaf_dispersion_ema[l_idx] = 0.95 * self.leaf_dispersion_ema[l_idx] + 0.05 * disp | |
| return top_k_weights.to(x_flat.dtype), top_k_idx, aux_loss, sup_loss, branch_acc, leaf_acc | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| target_branches: Optional[torch.Tensor] = None, | |
| target_leaves: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, float]: | |
| bsz, seq_len, dim = x.shape | |
| x_flat = x.view(-1, dim) | |
| r_flat = self.build_causal_router_context(x) | |
| shared_out = self.shared_e00(x_flat) | |
| top_w, top_idx, aux_loss, sup_loss, b_acc, l_acc = self.compute_hierarchical_routing( | |
| r_flat, target_branches, target_leaves | |
| ) | |
| routed_out = torch.zeros_like(shared_out) | |
| for l_idx, leaf_mlp in enumerate(self.leaves): | |
| token_mask = top_idx == l_idx | |
| if not token_mask.any(): | |
| continue | |
| token_positions, k_slot = torch.where(token_mask) | |
| leaf_act = leaf_mlp(x_flat[token_positions]) | |
| weights = top_w[token_positions, k_slot].unsqueeze(-1).to(leaf_act.dtype) | |
| routed_out.index_add_(0, token_positions, (leaf_act * weights).to(routed_out.dtype)) | |
| return (shared_out + routed_out).view(bsz, seq_len, dim), aux_loss, sup_loss, b_acc, l_acc | |
| class PhobosScratchBlock(nn.Module): | |
| """Single Pre-Norm Transformer + Hierarchical Cortical MoE Block with DeepInfer `Engram` & `mHC`.""" | |
| def __init__(self, cfg: PhobosScratchConfig, layer_id: int = 0) -> None: | |
| super().__init__() | |
| self.layer_id = layer_id | |
| self.attn_norm, self.attn = RMSNorm(cfg.d_model, cfg.rms_norm_eps), CausalGroupedQueryAttention(cfg) | |
| self.moe_norm, self.moe = RMSNorm(cfg.d_model, cfg.rms_norm_eps), HierarchicalCorticalMoELayer(cfg) | |
| self.engram = DeepInferEngramMemory(cfg.d_model, layer_id) if (cfg.use_deepinfer_engram and layer_id % 2 == 0) else None | |
| self.hc_attn = DeepInferHyperConnection(cfg.d_model, layer_id, cfg.hc_mult) if cfg.use_deepinfer_hc else None | |
| self.hc_moe = DeepInferHyperConnection(cfg.d_model, layer_id, cfg.hc_mult) if cfg.use_deepinfer_hc else None | |
| def forward( | |
| self, state: torch.Tensor, input_ids: Optional[torch.Tensor] = None, | |
| target_branches: Optional[torch.Tensor] = None, target_leaves: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, float]: | |
| if self.hc_attn is not None and self.hc_moe is not None: | |
| x_attn_in, post_w1, res_m1 = self.hc_attn.pre_mix(state) | |
| if self.engram is not None and input_ids is not None: | |
| x_attn_in = self.engram(x_attn_in, input_ids) | |
| state = self.hc_attn.post_mix(self.attn(self.attn_norm(x_attn_in)), state, post_w1, res_m1) | |
| x_moe_in, post_w2, res_m2 = self.hc_moe.pre_mix(state) | |
| moe_out, aux_loss, sup_loss, b_acc, l_acc = self.moe(self.moe_norm(x_moe_in), target_branches, target_leaves) | |
| return self.hc_moe.post_mix(moe_out, state, post_w2, res_m2), aux_loss, sup_loss, b_acc, l_acc | |
| if self.engram is not None and input_ids is not None: | |
| state = self.engram(state, input_ids) | |
| state = state + self.attn(self.attn_norm(state)) | |
| moe_out, aux_loss, sup_loss, b_acc, l_acc = self.moe(self.moe_norm(state), target_branches, target_leaves) | |
| return state + moe_out, aux_loss, sup_loss, b_acc, l_acc | |
| class PhobosScratchLM(nn.Module): | |
| """Complete From-Scratch Causal SynapCore Phobos Language Model with DeepInfer V4.1 Modules.""" | |
| def __init__(self, cfg: PhobosScratchConfig) -> None: | |
| super().__init__() | |
| self.cfg = cfg | |
| self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.d_model) | |
| self.layers = nn.ModuleList([PhobosScratchBlock(cfg, layer_id=i) for i in range(cfg.num_layers)]) | |
| self.hc_head_mix = nn.Parameter(torch.zeros(cfg.hc_mult)) if cfg.use_deepinfer_hc else None | |
| self.final_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps) | |
| self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) | |
| self.markov_mtp_head = DeepInferMarkovMTPHead(cfg.vocab_size, cfg.d_model, cfg.markov_rank) if cfg.use_deepinfer_markov_mtp else None | |
| nn.init.normal_(self.embed_tokens.weight, mean=0.0, std=0.02) | |
| if cfg.tie_word_embeddings: | |
| self.lm_head.weight = self.embed_tokens.weight | |
| def _init_state(self, x: torch.Tensor) -> torch.Tensor: | |
| if not self.cfg.use_deepinfer_hc: | |
| return x | |
| scales = torch.linspace(0.97, 1.03, steps=self.cfg.hc_mult, device=x.device, dtype=x.dtype).view(1, 1, -1, 1) | |
| return x.unsqueeze(-2) * scales | |
| def forward( | |
| self, input_ids: torch.Tensor, target_branches: Optional[torch.Tensor] = None, | |
| target_leaves: Optional[torch.Tensor] = None, pad_id: int = 0, return_fused_loss: bool = False, | |
| assistant_mask: Optional[torch.Tensor] = None, | |
| ) -> Dict[str, Any]: | |
| state = self._init_state(self.embed_tokens(input_ids)) | |
| tb_flat = target_branches.view(-1) if target_branches is not None else None | |
| tl_flat = target_leaves.view(-1) if target_leaves is not None else None | |
| total_aux, total_sup = torch.tensor(0.0, device=input_ids.device), torch.tensor(0.0, device=input_ids.device) | |
| sum_b_acc, sum_l_acc, n_layers = 0.0, 0.0, float(len(self.layers)) | |
| for layer in self.layers: | |
| state, l_aux, l_sup, b_acc, l_acc = layer(state, input_ids, tb_flat, tl_flat) | |
| total_aux, total_sup, sum_b_acc, sum_l_acc = total_aux + l_aux, total_sup + l_sup, sum_b_acc + b_acc, sum_l_acc + l_acc | |
| x_out = torch.einsum("m,btmd->btd", F.softmax(self.hc_head_mix.float(), dim=-1).to(state.dtype), state) if (self.cfg.use_deepinfer_hc and self.hc_head_mix is not None) else state | |
| h_norm = self.final_norm(x_out) | |
| if return_fused_loss and self.markov_mtp_head is not None: | |
| sal_ce, raw_ce, mtp_loss, mean_conf = self.markov_mtp_head.forward_fused_training_losses(h_norm, input_ids, self.lm_head.weight, pad_id=pad_id, assistant_mask=assistant_mask) | |
| return {"salience_ce": sal_ce, "raw_ce": raw_ce, "aux_loss": total_aux / n_layers, "sup_router_loss": total_sup / n_layers, "mtp_loss": mtp_loss, "mean_confidence": mean_conf, "branch_routing_acc": sum_b_acc / n_layers, "leaf_routing_acc": sum_l_acc / n_layers} | |
| if self.markov_mtp_head is not None: | |
| logits, conf, mtp_loss = self.markov_mtp_head.forward_logits_and_mtp(h_norm, input_ids, self.lm_head.weight, pad_id=pad_id, compute_mtp_loss=self.training) | |
| mean_conf = float(conf.detach().mean().item()) | |
| else: | |
| logits, mtp_loss, mean_conf = self.lm_head(h_norm), torch.tensor(0.0, device=input_ids.device), 1.0 | |
| return { | |
| "logits": logits, "aux_loss": total_aux / n_layers, "sup_router_loss": total_sup / n_layers, | |
| "mtp_loss": mtp_loss, "mean_confidence": mean_conf, | |
| "branch_routing_acc": sum_b_acc / n_layers, "leaf_routing_acc": sum_l_acc / n_layers, | |
| } | |
| def inspect_natural_routing(self, input_ids: torch.Tensor) -> Dict[str, Any]: | |
| """Inspects which Branch (0..3) and Leaf (0..7) the model naturally routes a prompt to.""" | |
| self.eval() | |
| state = self._init_state(self.embed_tokens(input_ids)) | |
| layer_routes: List[Dict[str, int]] = [] | |
| for idx, layer in enumerate(self.layers): | |
| if layer.hc_attn is not None and layer.hc_moe is not None: | |
| x_attn_in, post_w1, res_m1 = layer.hc_attn.pre_mix(state) | |
| if layer.engram is not None: | |
| x_attn_in = layer.engram(x_attn_in, input_ids) | |
| state = layer.hc_attn.post_mix(layer.attn(layer.attn_norm(x_attn_in)), state, post_w1, res_m1) | |
| x_moe_in, post_w2, res_m2 = layer.hc_moe.pre_mix(state) | |
| r_flat = layer.moe.build_causal_router_context(layer.moe_norm(x_moe_in)) | |
| _, top_idx, _, _, _, _ = layer.moe.compute_hierarchical_routing(r_flat) | |
| moe_out, _, _, _, _ = layer.moe(layer.moe_norm(x_moe_in)) | |
| state = layer.hc_moe.post_mix(moe_out, state, post_w2, res_m2) | |
| else: | |
| if layer.engram is not None: | |
| state = layer.engram(state, input_ids) | |
| x_attn = state + layer.attn(layer.attn_norm(state)) | |
| r_flat = layer.moe.build_causal_router_context(layer.moe_norm(x_attn)) | |
| _, top_idx, _, _, _, _ = layer.moe.compute_hierarchical_routing(r_flat) | |
| moe_out, _, _, _, _ = layer.moe(layer.moe_norm(x_attn)) | |
| state = x_attn + moe_out | |
| tail_idx = top_idx[max(0, len(top_idx) // 2) :, 0] | |
| vals, counts = torch.unique(tail_idx, return_counts=True) | |
| dom_phys = int(vals[counts.argmax()].item()) | |
| dom_leaf = int(layer.moe.leaf_origin_id[dom_phys].item()) | |
| dom_branch = int(layer.moe.leaf_parent_branch[dom_phys].item()) | |
| layer_routes.append({"layer": idx, "branch_id": dom_branch, "leaf_id": dom_leaf, "physical_leaf_id": dom_phys}) | |
| return {"layer_routes": layer_routes, "final_layer": layer_routes[-1]} | |
| def generate( | |
| self, | |
| input_ids: torch.Tensor, | |
| max_new_tokens: int = 32, | |
| temperature: float = 0.0, | |
| eos_token_id: Optional[int] = None, | |
| seed: Optional[int] = 42, | |
| repetition_penalty: float = 1.15, | |
| ) -> torch.Tensor: | |
| """Causal autoregressive token generation with DeepInfer structured bitmasking & seed control.""" | |
| self.eval() | |
| if seed is not None: | |
| torch.manual_seed(seed) | |
| seq = input_ids.clone() | |
| prompt_len = seq.shape[1] | |
| for step_idx in range(max_new_tokens): | |
| ctx = seq[:, -self.cfg.max_seq_len :] | |
| logits = self(ctx)["logits"][:, -1, :].float() | |
| if self.cfg.use_deepinfer_structured_mask: | |
| DeepInferStructuredMasker.apply_step_constraints_inplace( | |
| logits, seq[0, prompt_len:].tolist(), eos_token_id, steps_generated=step_idx | |
| ) | |
| if repetition_penalty > 1.0 and seq.shape[1] > prompt_len: | |
| for b in range(seq.shape[0]): | |
| seen = torch.unique(seq[b, max(prompt_len, seq.shape[1] - 24) :]) | |
| scores = logits[b, seen] | |
| logits[b, seen] = torch.where(scores > 0, scores / repetition_penalty, scores * repetition_penalty) | |
| next_tok = ( | |
| torch.argmax(logits, dim=-1, keepdim=True) | |
| if temperature <= 1e-5 | |
| else torch.multinomial(F.softmax(logits / temperature, dim=-1), num_samples=1) | |
| ) | |
| seq = torch.cat([seq, next_tok], dim=1) | |
| if eos_token_id is not None and int(next_tok[0, 0].item()) == eos_token_id: | |
| break | |
| return seq | |
| def parameter_summary(self) -> Dict[str, Any]: | |
| total = sum(p.numel() for p in self.parameters()) | |
| leaf_total = sum(p.numel() for layer in self.layers for leaf in layer.moe.leaves for p in leaf.parameters()) | |
| active_leaf = sum(sum(p.numel() for p in layer.moe.leaves[0].parameters()) * layer.moe.top_k for layer in self.layers) | |
| return { | |
| "total_params": total, "active_params_per_token": total - leaf_total + active_leaf, | |
| "num_layers": len(self.layers), "leaves_per_layer": [l.moe.num_leaves for l in self.layers], | |
| } | |
| def save_checkpoint(self, filename: str = "stage0_verified.pt", extra_meta: Optional[Dict[str, Any]] = None) -> Path: | |
| SCRATCH_WEIGHTS_DIR.mkdir(parents=True, exist_ok=True) | |
| out_path = SCRATCH_WEIGHTS_DIR / filename | |
| torch.save({ | |
| "config": asdict(self.cfg), "leaves_per_layer": [l.moe.num_leaves for l in self.layers], | |
| "leaf_parent_branches": [l.moe.leaf_parent_branch.cpu() for l in self.layers], | |
| "state_dict": self.state_dict(), "meta": extra_meta or {}, | |
| }, out_path) | |
| return out_path | |