From 631e6ea3e46f14e98466347559d43c80e609912f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 11:33:41 +0000 Subject: [PATCH] Add --skip-mhc flag for simple residual diagnostic When enabled, bypasses mHC pre/post blocks and uses direct residual connections with 0.1 scaling. This helps isolate whether the mHC implementation is causing the garbage output. --- single_shot_inference.py | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index aacc7d99..a1a67e20 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -48,6 +48,7 @@ def parse_args(): p = argparse.ArgumentParser(description='DSV4 Single-Shot Inference') 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('--max-tokens', type=int, default=512, help='Max new tokens to generate') p.add_argument('--prompt', type=str, default=None, help='Override prompt') return p.parse_args() @@ -61,7 +62,7 @@ PROMPT = _args.prompt or "The capital of France is" 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) -LAYER_TRACE = False # If True, trace every computation step for first layer of first prefill token +SKIP_MHC = _args.skip_mhc # If True, bypass mHC and use simple residual connections (diagnostic) MHC_DIAG = True # If True, print per-layer mHC diagnostics (B_l row/col sums, C_l values) # When True: applies inverse RoPE at query position → converts absolute→relative # When False: leaves relative position encoding intact for output projection @@ -398,8 +399,13 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # ATTENTION SUB-BLOCK # ================================================================== - # -- mHC pre_block (attention) -- - x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H) + if SKIP_MHC: + # Simple residual: skip mHC, use direct input + x_in = X_l[:, 0, :] # Just take stream 0 + attn_ctx = None + else: + # -- mHC pre_block (attention) -- + x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H) if MHC_DIAG: # mHC diagnostics B_l, C_l = attn_ctx.B_l, attn_ctx.C_l print(f" L{li} pre_attn: |X_l|={X_l.abs().max().item():.2f} |x_in|={x_in.abs().max().item():.2f}", flush=True) @@ -546,8 +552,11 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, w[f"{pre}.o_b_proj.weight_scale"], w[f"{pre}.o_b_proj.weight_scale_2"]) # (T, H) - # -- mHC post_block (attention) -- - X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H) + if SKIP_MHC: + X_mid = X_l[:, 0, :].unsqueeze(1).expand(-1, 4, -1) + F_attn.unsqueeze(1) * 0.1 + else: + # -- mHC post_block (attention) -- + X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H) # Diagnostic: check mHC is stabilizing the residual if MHC_DIAG: # mHC diagnostics B_l, C_l = attn_ctx.B_l, attn_ctx.C_l @@ -562,8 +571,12 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # FFN SUB-BLOCK # ================================================================== - # -- mHC pre_block (FFN) -- - x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid) # (T, H) + if SKIP_MHC: + x_ffn = X_mid[:, 0, :] # Just take stream 0 + ffn_ctx = None + else: + # -- mHC pre_block (FFN) -- + x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid) # (T, H) # -- RMSNorm (pre-norm before FFN) -- x_ffn_normed = ffn_norm.forward(x_ffn) # (T, H) BF16 @@ -571,8 +584,11 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # -- MoE + shared expert -- F_ffn = moe_forward(x_ffn_normed, w, li, cfg, token_id, device) - # -- mHC post_block (FFN) -- - X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) # (T, n_hc, H) + if SKIP_MHC: + X_next = X_mid + F_ffn.unsqueeze(1) * 0.1 + else: + # -- mHC post_block (FFN) -- + X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) # (T, n_hc, H) if MHC_DIAG: # ffn mHC diagnostics B_l_ffn, C_l_ffn = ffn_ctx.B_l, ffn_ctx.C_l print(f" L{li} ffn: |X_mid|={X_mid.abs().max().item():.2f} |F_ffn|={F_ffn.abs().max().item():.2f} |B|={B_l_ffn.abs().max().item():.4f} |C|={C_l_ffn.abs().max().item():.4f} |X_next|={X_next.abs().max().item():.2f}", flush=True)