diff --git a/single_shot_inference.py b/single_shot_inference.py index af7a5b23..fc8f553d 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -49,8 +49,7 @@ def parse_args(): p.add_argument('--no-inverse-rope', action='store_true', help='Skip inverse RoPE on attention output') p.add_argument('--skip-moe', action='store_true', help='Only use shared expert (skip routed)') p.add_argument('--skip-mhc', action='store_true', help='Bypass mHC, use simple residual (diagnostic)') - p.add_argument('--no-thinking', action='store_true', help='Force model to skip thinking (use <|EOT|> instead of thinking tokens)') - p.add_argument('--max-tokens', type=int, default=512, help='Max new tokens to generate') + p.add_argument('--max-tokens', type=int, default=4096, help='Max new tokens to generate') p.add_argument('--prompt', type=str, default=None, help='Override prompt') return p.parse_args() @@ -64,12 +63,10 @@ NUM_GPUS = 8 SKIP_ROUTED_MOE = _args.skip_moe # If True, only use shared expert (debug) INVERSE_ROPE = not _args.no_inverse_rope # If False, skip inverse RoPE on attention output (diagnostic) SKIP_MHC = _args.skip_mhc # If True, bypass mHC and use simple residual connections (diagnostic) +# NO --no-thinking flag: this is a reasoning model, it MUST think. Do not skip thinking tokens. MHC_DIAG = False # If True, print per-layer mHC diagnostics (B_l row/col sums, C_l values) GROWTH_DIAG = True # If True, print compact per-layer residual trace ATTN_DIAG = False # If True, print per-layer attention entropy (expensive) -# When True: applies inverse RoPE at query position → converts absolute→relative -# When False: leaves relative position encoding intact for output projection -# DSV4 partial RoPE only affects last 64/512 dims; first 448 are always un-RoPE'd print(f"Config: INVERSE_ROPE={INVERSE_ROPE}, SKIP_ROUTED_MOE={SKIP_ROUTED_MOE}, MAX_NEW_TOKENS={MAX_NEW_TOKENS}") # ===================================================================== @@ -301,6 +298,67 @@ class SimpleKVCache: return self.k[:, :self.len], self.v[:, :self.len] +class HcHead: + """Final readout from the mHC residual state. + + Matches HuggingFace DeepseekV4HyperHead: + 1. UnweightedRMSNorm on flattened residual + 2. Linear projection: fn (4, 28672) → (T, 4) pre-mixing coefficients + 3. sigmoid(pre * scale + base) + eps + 4. Weighted sum: (pre.unsqueeze(-1) * X).sum(dim=1) → (T, H) + + This is NOT stream 0! The model uses a learned projection to collapse + the 4 hyper-connection streams into the final hidden state. + """ + def __init__(self, hidden_dim=7168, n_hc=4, device='cuda:0'): + self.d = hidden_dim + self.n_hc = n_hc + self.K = n_hc * hidden_dim # 28672 + self.device = device + self.fn = None # (4, 28672) FP32 + self.base = None # (4,) BF16 + self.scale = None # (1,) FP32 scalar + self.eps = 1e-6 + + def load_from_checkpoint(self, fn, base, scale=None): + """Load hc_head weights. + fn: (4, 28672) — maps flattened residual to 4 mixing coefficients + base: (4,) — bias + scale: (1,) — sigmoid scaling (optional, defaults to 1.0) + """ + self.fn = fn.to(device=self.device, dtype=torch.float32).contiguous() + self.base = base.to(device=self.device, dtype=torch.bfloat16).contiguous() + if scale is not None: + self.scale = scale.to(device=self.device, dtype=torch.float32).contiguous() + else: + self.scale = torch.tensor(1.0, dtype=torch.float32, device=self.device) + + def forward(self, X_L): + """Read out final hidden state from mHC residual. + + X_L: (T, n_hc, H) BF16 — final layer residual state + Returns: (T, H) BF16 + """ + T = X_L.shape[0] + # Flatten: (T, n_hc * H) + X_flat = X_L.reshape(T, self.K).bfloat16() + + # Unweighted RMSNorm (same as in mHC) + X_f = X_flat.float() + rms_inv = X_f.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt() + X_normed = (X_f * rms_inv).bfloat16() + + # Linear projection: (T, K) @ (4, K).T → (T, 4) + mixes = torch.nn.functional.linear(X_normed, self.fn.bfloat16()).float() + + # Apply scale + bias + sigmoid + eps + pre = torch.sigmoid(mixes * self.scale + self.base.float().unsqueeze(0)) + self.eps + + # Weighted sum: (T, n_hc) * (T, n_hc, H) → (T, H) + x_out = (pre.unsqueeze(-1) * X_L.float()).sum(dim=1) + return x_out.to(torch.bfloat16) + + # ===================================================================== # Weight loading — streams safetensors shards, distributes to 8 GPUs # ===================================================================== @@ -825,6 +883,21 @@ def main(): final_norm_w = all_weights.get("model.norm.weight") if final_norm_w is not None: final_norm_w = final_norm_w.to('cuda:0') + + # ==== hc_head — final readout from mHC residual (NOT stream 0!) ==== + # The model uses DeepseekV4HyperHead to project from the 4-stream mHC + # residual to the final hidden state. Just taking stream 0 is WRONG. + hc_head = HcHead(hidden_dim=H, n_hc=n_hc, device='cuda:0') + hc_fn = all_weights.get("model.hc_head.hc_fn") + hc_base = all_weights.get("model.hc_head.hc_base") + hc_scale = all_weights.get("model.hc_head.hc_scale") + if hc_fn is not None and hc_base is not None: + hc_head.load_from_checkpoint(hc_fn, hc_base, hc_scale) + print(f" hc_head loaded: fn={hc_fn.shape} base={hc_base.shape}") + else: + print(" WARNING: hc_head weights not found in checkpoint, falling back to stream 0") + hc_head = None + # Build RoPE caches with YaRN scaling from model config rope_params = cfg.get("rope_parameters", {}) rope_type = rope_params.get("rope_type", "default") @@ -875,7 +948,7 @@ def main(): tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) minimal_e2e_test(layer_weights, cfg, rope_caches, attn_mhc_blocks, ffn_mhc_blocks, attn_norms, ffn_norms, embed, lm_w, - final_norm_w, tokenizer) + final_norm_w, tokenizer, hc_head=hc_head) # ==== Phase 2.6: Single-layer trace ==== if True: # Always run the trace @@ -1093,8 +1166,11 @@ def main(): X = X.to('cuda:0') torch.cuda.set_device(0) - # Read out stream 0 → RMSNorm → lm_head - x_out = X[:, 0, :] # (1, H) + # Read out final hidden state via hc_head (learned projection from mHC residual) + if hc_head is not None: + x_out = hc_head.forward(X) # (1, H) — proper mHC readout + else: + x_out = X[:, 0, :] # (1, H) — fallback if final_norm_w is not None: xf = x_out.float() rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() @@ -1109,15 +1185,6 @@ def main(): thinking_in_top20 = any(tid.item() in [128821, 128822] for tid in top20_ids) top20_ids_set = set(top20_ids.tolist()) next_id = torch.argmax(logits, dim=-1).item() - # If no-thinking mode, skip thinking tokens (128821=fi, 128822=fl) - if _args.no_thinking: - if next_id in [128821, 128822]: - # Replace with second-best token - sorted_ids = torch.argsort(logits[0], descending=True) - for alt_id in sorted_ids: - if alt_id.item() not in [128821, 128822]: - next_id = alt_id.item() - break generated.append(next_id) all_tokens.append(next_id) @@ -1158,7 +1225,7 @@ def main(): def minimal_e2e_test(layer_weights, cfg, rope_caches, attn_mhc_blocks, ffn_mhc_blocks, attn_norms, ffn_norms, embed, lm_w, - final_norm_w, tokenizer): + final_norm_w, tokenizer, hc_head=None): """Process a single token 'The' through the model and check output logits. This is a focused diagnostic: if the model can't even produce reasonable @@ -1234,7 +1301,10 @@ def minimal_e2e_test(layer_weights, cfg, rope_caches, attn_mhc_blocks, torch.cuda.set_device(0) # Final norm + lm_head - x_out = X[:, 0, :] + if hc_head is not None: + x_out = hc_head.forward(X) # proper mHC readout + else: + x_out = X[:, 0, :] if final_norm_w is not None: xf = x_out.float() rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()