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 — DeepInfer (DeepSeek-V4.1) Architectural Modules | |
| (`phobos-1/models/phobos_scratch_v1/deepinfer_modules.py`) | |
| ============================================================================= | |
| Pure-PyTorch, A100 (`SM80`) & training-compatible implementations of the 5 core | |
| architectural innovations from `https://github.com/deepseek-ai/DeepInfer`: | |
| 1. `DeepInferNoAuxGate` (`v4_modules/moe.py` -> `TopKGate`): | |
| Uses `sqrtsoftplus` scoring + `noaux_tc` (`e_score_correction_bias`) dynamic | |
| auxiliary-loss-free load balancing and normalized top-k probabilities. | |
| 2. `DeepInferEngramMemory` (`v4_modules/engram.py` & `engram_metadata.py`): | |
| Prime-modulo causal 2-gram & 3-gram hash lookup tables with gated KV projection | |
| fused directly into transformer hidden states. | |
| 3. `DeepInferHyperConnection` (`v4_modules/hc_blocks.py` & `mhc_modeling.py`): | |
| Multi-Head Hyper-Connections (`hc_mult = 4` parallel residual streams) with | |
| learned dynamic sigmoid pre-mix, post-mix, and residual stream coupling. | |
| 4. `DeepInferMarkovMTPHead` (`v4_modules/heads.py` & `mtp/dflash.py`): | |
| `MarkovGPTHead` low-rank first-order bigram logit correction (`markov_w1/w2`) | |
| + `DFlash` next-token draft projection (`MTP`) + `ConfidenceHead`. | |
| 5. `DeepInferStructuredMasker` (`structure_output/mask.py` & `matcher.py`): | |
| Step-wise token logit bitmasking guaranteeing balanced `<think>...</think>` | |
| and `<tool_call>...</tool_call>` blocks and preventing premature `<|eos|>`. | |
| Strictly <= 320 lines. | |
| """ | |
| from __future__ import annotations | |
| from typing import Dict, List, Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| def sqrtsoftplus(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: | |
| """DeepInfer `TopKGate` scoring function: `sqrt(softplus(x) + eps)`.""" | |
| return torch.sqrt(F.softplus(x.float()) + eps).to(x.dtype) | |
| class DeepInferNoAuxGate(nn.Module): | |
| """ | |
| DeepInfer `TopKGate` (`v4_modules/moe.py`) with `sqrtsoftplus` activation and | |
| `noaux_tc` (`e_score_correction_bias`) auxiliary-loss-free load balancing. | |
| """ | |
| def __init__(self, num_branches: int = 4, num_leaves: int = 8, bias_lr: float = 0.0015) -> None: | |
| super().__init__() | |
| self.num_branches = num_branches | |
| self.num_leaves = num_leaves | |
| self.bias_lr = bias_lr | |
| self.register_buffer("branch_correction_bias", torch.zeros(num_branches, dtype=torch.float32)) | |
| self.register_buffer("leaf_correction_bias", torch.zeros(num_leaves, dtype=torch.float32)) | |
| def score_and_balance( | |
| self, branch_logits: torch.Tensor, leaf_logits: torch.Tensor, top_k_leaves: int = 2 | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """ | |
| Computes `sqrtsoftplus` scores, applies `e_score_correction_bias` for top-k choice | |
| (`noaux_tc`), updates correction biases during training, and returns normalized weights. | |
| """ | |
| b_scores = sqrtsoftplus(branch_logits) | |
| l_scores = sqrtsoftplus(leaf_logits) | |
| b_choice = b_scores + self.branch_correction_bias.to(b_scores.dtype) | |
| l_choice = l_scores + self.leaf_correction_bias.to(l_scores.dtype) | |
| _, top_l_idx = torch.topk(l_choice, k=top_k_leaves, dim=-1) | |
| top_l_raw = torch.gather(l_scores, dim=-1, index=top_l_idx) | |
| top_l_weights = top_l_raw / (top_l_raw.sum(dim=-1, keepdim=True) + 1e-6) | |
| if self.training: | |
| with torch.no_grad(): | |
| b_top1 = b_choice.argmax(dim=-1).reshape(-1) | |
| b_freq = torch.bincount(b_top1, minlength=self.num_branches).float() / max(1, b_top1.numel()) | |
| b_target = 1.0 / float(self.num_branches) | |
| self.branch_correction_bias.add_(self.bias_lr * torch.sign(b_target - b_freq)).clamp_(-0.4, 0.4) | |
| l_flat = top_l_idx.reshape(-1) | |
| l_freq = torch.bincount(l_flat, minlength=self.num_leaves).float() / max(1, l_flat.numel()) | |
| l_target = 1.0 / float(self.num_leaves) | |
| self.leaf_correction_bias.add_(self.bias_lr * torch.sign(l_target - l_freq)).clamp_(-0.4, 0.4) | |
| b_probs = b_scores / (b_scores.sum(dim=-1, keepdim=True) + 1e-6) | |
| l_probs = l_scores / (l_scores.sum(dim=-1, keepdim=True) + 1e-6) | |
| return b_probs, l_probs, top_l_idx, top_l_weights | |
| class DeepInferEngramMemory(nn.Module): | |
| """ | |
| DeepInfer `Engram` (`v4_modules/engram.py` & `engram_metadata.py`): | |
| Hashes causal 2-grams `(x_{t-1}, x_t)` and 3-grams `(x_{t-2}, x_{t-1}, x_t)` into | |
| distinct prime-sized tables, projects through a fused KV linear layer, and gates into `h`. | |
| """ | |
| def __init__( | |
| self, | |
| d_model: int, | |
| layer_id: int, | |
| ngram_dim: int = 48, | |
| prime_2gram: int = 8191, | |
| prime_3gram: int = 8209, | |
| ) -> None: | |
| super().__init__() | |
| self.prime_2gram = prime_2gram | |
| self.prime_3gram = prime_3gram | |
| # Deterministic odd multipliers per layer (`NgramHashMapping` convention) | |
| m0 = 2 * (10007 * (layer_id + 1) + 31) + 1 | |
| m1 = 2 * (20011 * (layer_id + 1) + 97) + 1 | |
| m2 = 2 * (30013 * (layer_id + 1) + 193) + 1 | |
| self.register_buffer("multipliers", torch.tensor([m0, m1, m2], dtype=torch.long), persistent=False) | |
| self.embed_2gram = nn.Embedding(prime_2gram, ngram_dim) | |
| self.embed_3gram = nn.Embedding(prime_3gram, ngram_dim) | |
| nn.init.normal_(self.embed_2gram.weight, mean=0.0, std=0.02) | |
| nn.init.normal_(self.embed_3gram.weight, mean=0.0, std=0.02) | |
| self.fused_kv_projection = nn.Linear(ngram_dim * 2, d_model * 2, bias=False) | |
| self.weight_hidden = nn.Parameter(torch.ones(d_model) * 0.5) | |
| self.weight_embed = nn.Parameter(torch.ones(d_model) * 0.5) | |
| self.res_scale = nn.Parameter(torch.tensor(0.15)) | |
| def forward(self, h: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: | |
| """Applies causal 2-gram & 3-gram hash lookup and gated fusion to `h` (`[B, T, D]`).""" | |
| ids_64 = input_ids.long() | |
| prev1 = F.pad(ids_64[:, :-1], (1, 0), value=0) | |
| prev2 = F.pad(ids_64[:, :-2], (2, 0), value=0) | |
| m0, m1, m2 = self.multipliers[0], self.multipliers[1], self.multipliers[2] | |
| h2 = torch.remainder(ids_64 * m0 + prev1 * m1, self.prime_2gram) | |
| h3 = torch.remainder(ids_64 * m0 + prev1 * m1 + prev2 * m2, self.prime_3gram) | |
| ngram_feat = torch.cat([self.embed_2gram(h2), self.embed_3gram(h3)], dim=-1).to(h.dtype) | |
| k_ng, v_ng = self.fused_kv_projection(ngram_feat).chunk(2, dim=-1) | |
| h_norm = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + 1e-6) | |
| gate = torch.sigmoid(h_norm * self.weight_hidden.to(h.dtype) + k_ng * self.weight_embed.to(h.dtype)) | |
| return h + self.res_scale.to(h.dtype) * (gate * v_ng) | |
| class DeepInferHyperConnection(nn.Module): | |
| """ | |
| DeepInfer `SigmoidHCBlockMixin` (`v4_modules/hc_blocks.py` & `mhc_modeling.py`): | |
| Maintains `hc_mult = 4` parallel residual streams `[B, T, M, D]` and learns dynamic | |
| pre-layer mixing, post-layer writeback, and cross-stream residual coupling. | |
| """ | |
| def __init__(self, d_model: int, layer_id: int, hc_mult: int = 4) -> None: | |
| super().__init__() | |
| self.hc_mult = hc_mult | |
| self.d_model = d_model | |
| # Dynamic summary projection from each stream `D -> 2+M` scores with symmetry-breaking init | |
| self.stream_score = nn.Linear(d_model, 2 + hc_mult, bias=True) | |
| nn.init.normal_(self.stream_score.weight, mean=0.0, std=0.02 / (d_model ** 0.5)) | |
| nn.init.zeros_(self.stream_score.bias) | |
| base_pre = torch.linspace(-1.0, 1.0, steps=hc_mult) | |
| base_pre[layer_id % hc_mult] = 1.5 | |
| self.base_pre = nn.Parameter(base_pre) | |
| base_post = torch.full((hc_mult,), -1.0) | |
| base_post[layer_id % hc_mult] = 1.5 | |
| self.base_post = nn.Parameter(base_post) | |
| self.base_res = nn.Parameter(torch.eye(hc_mult) * 2.5 - 1.25) | |
| self.dynamic_scale = nn.Parameter(torch.tensor(0.25)) | |
| def pre_mix(self, trunk: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """Mixes `[B, T, M, D]` trunk into sub-layer input `[B, T, D]` and returns `(x_in, post_w, res_mat)`.""" | |
| raw = self.stream_score(trunk) * self.dynamic_scale.to(trunk.dtype) # [B, T, M, 2+M] | |
| pre_logits = raw[..., 0] + self.base_pre.to(trunk.dtype) # [B, T, M] | |
| post_logits = raw[..., 1] + self.base_post.to(trunk.dtype) # [B, T, M] | |
| res_logits = raw[..., 2:] + self.base_res.to(trunk.dtype) # [B, T, M, M] | |
| pre_w = F.softmax(pre_logits.float(), dim=-1).to(trunk.dtype) | |
| post_w = torch.sigmoid(post_logits) | |
| res_mat = F.softmax(res_logits.float(), dim=-1).to(trunk.dtype) | |
| x_in = torch.einsum("btm,btmd->btd", pre_w, trunk) | |
| return x_in, post_w, res_mat | |
| def post_mix( | |
| self, sublayer_out: torch.Tensor, trunk: torch.Tensor, post_w: torch.Tensor, res_mat: torch.Tensor | |
| ) -> torch.Tensor: | |
| """Writes `sublayer_out` (`[B, T, D]`) back into the `M`-stream residual trunk (`[B, T, M, D]`).""" | |
| mixed_trunk = torch.einsum("btij,btjd->btid", res_mat, trunk) | |
| return mixed_trunk + post_w.unsqueeze(-1) * sublayer_out.unsqueeze(-2) | |
| class DeepInferMarkovMTPHead(nn.Module): | |
| """ | |
| DeepInfer `MarkovGPTHead` + `DFlash` Next-Token MTP & `ConfidenceHead` | |
| (`v4_modules/heads.py` & `models/mtp/dflash.py`). | |
| """ | |
| def __init__(self, vocab_size: int, d_model: int, markov_rank: int = 64) -> None: | |
| super().__init__() | |
| self.markov_rank = markov_rank | |
| self.markov_w1 = nn.Embedding(vocab_size, markov_rank) | |
| self.markov_w2 = nn.Linear(markov_rank, vocab_size, bias=False) | |
| nn.init.normal_(self.markov_w1.weight, mean=0.0, std=0.02) | |
| nn.init.zeros_(self.markov_w2.weight) | |
| self.markov_scale = nn.Parameter(torch.tensor(0.25)) | |
| # DFlash lightweight next-2 token (MTP) projection + ConfidenceHead | |
| self.mtp_proj = nn.Linear(d_model + markov_rank, d_model, bias=False) | |
| self.confidence_head = nn.Linear(d_model + markov_rank, 1, bias=False) | |
| def forward_logits_and_mtp( | |
| self, | |
| h_norm: torch.Tensor, | |
| input_ids: torch.Tensor, | |
| lm_head_weight: torch.Tensor, | |
| pad_id: int = 0, | |
| compute_mtp_loss: bool = False, | |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """ | |
| Computes main logits with `MarkovGPTHead` first-order bigram correction, token | |
| confidence probability (`ConfidenceHead`), and optional `DFlash` 2-step MTP loss. | |
| """ | |
| m_emb = self.markov_w1(input_ids).to(h_norm.dtype) # [B, T, markov_rank] | |
| conf_in = torch.cat([h_norm, m_emb], dim=-1) | |
| fused_w = torch.cat( | |
| [lm_head_weight, self.markov_scale.to(lm_head_weight.dtype) * self.markov_w2.weight], | |
| dim=-1, | |
| ) | |
| logits = F.linear(conf_in, fused_w) | |
| confidence = torch.sigmoid(self.confidence_head(conf_in)).squeeze(-1) | |
| mtp_loss = h_norm.new_zeros(()) | |
| if compute_mtp_loss and input_ids.size(1) > 2: | |
| # Predict token t+2 from (h_t, markov_emb(x_{t+1})) only on non-pad targets | |
| next1_emb = m_emb[:, 1:-1, :] | |
| mtp_in = torch.cat([h_norm[:, :-2, :], next1_emb], dim=-1) | |
| mtp_h = F.silu(self.mtp_proj(mtp_in)).reshape(-1, h_norm.size(-1)) | |
| targets_t2 = input_ids[:, 2:].reshape(-1) | |
| valid_mask = targets_t2 != pad_id | |
| if valid_mask.any(): | |
| valid_h = mtp_h[valid_mask] | |
| valid_t = targets_t2[valid_mask] | |
| chunk_sz = 16384 | |
| total_loss = h_norm.new_zeros((), dtype=torch.float32) | |
| for s in range(0, valid_h.size(0), chunk_sz): | |
| e = min(valid_h.size(0), s + chunk_sz) | |
| c_logits = F.linear(valid_h[s:e], lm_head_weight).float() | |
| total_loss = total_loss + F.cross_entropy(c_logits, valid_t[s:e], reduction="sum") | |
| mtp_loss = total_loss / float(valid_h.size(0)) | |
| return logits, confidence, mtp_loss | |
| def forward_fused_training_losses( | |
| self, | |
| h_norm: torch.Tensor, | |
| input_ids: torch.Tensor, | |
| lm_head_weight: torch.Tensor, | |
| pad_id: int = 0, | |
| assistant_mask: Optional[torch.Tensor] = None, | |
| ) -> Tuple[torch.Tensor, float, torch.Tensor, float]: | |
| """ | |
| Memory-efficient chunked training loss for large A100 batches (`65K` tokens/step): | |
| Computes assistant-salience next-token CE, raw CE, and 2-step DFlash MTP loss only on | |
| non-pad tokens in 16,384-token chunks without allocating full `[B, T, V]` tensors. | |
| """ | |
| m_emb = self.markov_w1(input_ids).to(h_norm.dtype) | |
| conf_in = torch.cat([h_norm, m_emb], dim=-1) | |
| mean_conf = float(torch.sigmoid(self.confidence_head(conf_in)).detach().mean().item()) | |
| fused_w = torch.cat( | |
| [lm_head_weight, self.markov_scale.to(lm_head_weight.dtype) * self.markov_w2.weight], | |
| dim=-1, | |
| ) | |
| shift_in = conf_in[:, :-1, :].reshape(-1, conf_in.size(-1)) | |
| shift_t = input_ids[:, 1:].reshape(-1) | |
| shift_asst = assistant_mask[:, 1:].reshape(-1).float() if assistant_mask is not None else None | |
| valid_1 = shift_t != pad_id | |
| salience_ce = h_norm.new_zeros((), dtype=torch.float32) | |
| raw_ce_val = 0.0 | |
| chunk_sz = 16384 | |
| if valid_1.any(): | |
| vin = shift_in[valid_1] | |
| vt = shift_t[valid_1] | |
| vasst = shift_asst[valid_1] if shift_asst is not None else None | |
| w_num = h_norm.new_zeros((), dtype=torch.float32) | |
| w_den = h_norm.new_zeros((), dtype=torch.float32) | |
| raw_sum = h_norm.new_zeros((), dtype=torch.float32) | |
| for s in range(0, vin.size(0), chunk_sz): | |
| e = min(vin.size(0), s + chunk_sz) | |
| c_logits = F.linear(vin[s:e], fused_w).float() | |
| pt_ce = F.cross_entropy(c_logits, vt[s:e], reduction="none") | |
| with torch.no_grad(): | |
| sw = torch.ones_like(pt_ce) | |
| sw = torch.where(pt_ce >= 2.2, torch.full_like(sw, 2.0), sw) | |
| sw = torch.where(pt_ce <= 0.25, torch.full_like(sw, 0.25), sw) | |
| if vasst is not None: | |
| sw = sw * (0.20 + 1.80 * vasst[s:e]) | |
| w_num = w_num + (pt_ce * sw).sum() | |
| w_den = w_den + sw.sum() | |
| raw_sum = raw_sum + pt_ce.detach().sum() | |
| salience_ce = w_num / w_den.clamp_min(1.0) | |
| raw_ce_val = float((raw_sum / float(vin.size(0))).item()) | |
| mtp_loss = h_norm.new_zeros((), dtype=torch.float32) | |
| if input_ids.size(1) > 2: | |
| mtp_in = torch.cat([h_norm[:, :-2, :], m_emb[:, 1:-1, :]], dim=-1) | |
| mtp_h = F.silu(self.mtp_proj(mtp_in)).reshape(-1, h_norm.size(-1)) | |
| targets_t2 = input_ids[:, 2:].reshape(-1) | |
| valid_2 = targets_t2 != pad_id | |
| if valid_2.any(): | |
| vh2 = mtp_h[valid_2] | |
| vt2 = targets_t2[valid_2] | |
| t2_sum = h_norm.new_zeros((), dtype=torch.float32) | |
| for s in range(0, vh2.size(0), chunk_sz): | |
| e = min(vh2.size(0), s + chunk_sz) | |
| c_logits2 = F.linear(vh2[s:e], lm_head_weight).float() | |
| t2_sum = t2_sum + F.cross_entropy(c_logits2, vt2[s:e], reduction="sum") | |
| mtp_loss = t2_sum / float(vh2.size(0)) | |
| return salience_ce, raw_ce_val, mtp_loss, mean_conf | |
| class DeepInferStructuredMasker: | |
| """ | |
| DeepInfer `XGrammarUtils` / `XGrammarMatcher` (`structure_output/mask.py`): | |
| Applies step-wise logit bitmasking during generation to enforce well-formed | |
| `<think>...</think>` and `<tool_call>...</tool_call>` blocks and block premature `<|eos|>`. | |
| """ | |
| def apply_step_constraints_inplace( | |
| logits: torch.Tensor, | |
| generated_seq: List[int], | |
| eos_id: Optional[int], | |
| think_start_id: int = 5, | |
| think_end_id: int = 6, | |
| tool_start_id: int = 7, | |
| tool_end_id: int = 8, | |
| min_tokens_before_eos: int = 8, | |
| steps_generated: int = 0, | |
| ) -> None: | |
| """Modifies `logits` (`[B, V]`) in-place to enforce structural tag state machine rules.""" | |
| n_think_open = sum(1 for t in generated_seq if t == think_start_id) | |
| n_think_close = sum(1 for t in generated_seq if t == think_end_id) | |
| in_think = n_think_open > n_think_close | |
| n_tool_open = sum(1 for t in generated_seq if t == tool_start_id) | |
| n_tool_close = sum(1 for t in generated_seq if t == tool_end_id) | |
| in_tool = n_tool_open > n_tool_close | |
| if in_think: | |
| # Cannot open nested <think> or emit <|eos|> before closing </think> | |
| logits[:, think_start_id] = -1e9 | |
| if eos_id is not None: | |
| logits[:, eos_id] = -1e9 | |
| else: | |
| # Cannot emit orphan </think> when not inside <think> | |
| logits[:, think_end_id] = -1e9 | |
| if in_tool: | |
| logits[:, tool_start_id] = -1e9 | |
| if eos_id is not None: | |
| logits[:, eos_id] = -1e9 | |
| else: | |
| logits[:, tool_end_id] = -1e9 | |
| if eos_id is not None and steps_generated < min_tokens_before_eos: | |
| logits[:, eos_id] = -1e9 | |