biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix critical bug: add hc_head for final mHC readout (was using stream 0)

Browse Source 
The model uses DeepseekV4HyperHead to project from the 4-stream mHC
residual to the final hidden state. Just taking stream 0 (X[:,0,:])
is WRONG — the hc_head learns how to combine the 4 streams.

Also:
- Remove --no-thinking mode (this is a reasoning model, it MUST think)
- Increase default max_tokens from 512 to 4096
- Load hc_head weights (fn, base, scale) from checkpoint
This commit is contained in:
biondizzle
2026-05-31 21:13:02 +00:00
parent baee36e728
commit 274ea13251
1 changed files with 89 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
single_shot_inference.py
View File

@@ -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()