#!/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" @dataclass 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 @property def num_leaves(self) -> int: return self.num_branches * self.leaves_per_branch @property 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) @staticmethod 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] = {} @property 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, } @torch.no_grad() 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]} @torch.no_grad() 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