From 1434b35971a3b8e5e2c65586856c84c3e21bc6b3 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 09:21:41 +0000 Subject: [PATCH] =?UTF-8?q?Add=20residual=20diagnostic=20test=20=E2=80=94?= =?UTF-8?q?=20per-layer=20magnitude=20tracking?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/test_residual_diagnostic.py | 263 ++++++++++++++++++++++++++++++ 1 file changed, 263 insertions(+) create mode 100644 tests/test_residual_diagnostic.py diff --git a/tests/test_residual_diagnostic.py b/tests/test_residual_diagnostic.py new file mode 100644 index 00000000..5c42cc95 --- /dev/null +++ b/tests/test_residual_diagnostic.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +"""Diagnostic: track per-layer magnitudes to find residual explosion. + +Runs single token "The" through all 61 layers and prints: + |X_in|, |x_normed|, |F_attn|, |X_mid|, |F_ffn|, |X_next| + for each layer. + +This identifies WHERE the residual stream starts exploding. + +Usage (on B200): + source /root/dsv4-nvfp4-workspace/venv/bin/activate + cd /root/dsv4-nvfp4-workspace/kernel + python3 tests/test_residual_diagnostic.py +""" +import os, sys, math, json +import torch +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +NUM_GPUS = 8 + +# Import helpers from the main test +from tests.test_minimal_e2e import ( + FP4_LUT, dequant_nvfp4_weight, nvfp4_linear, RMSNorm, + build_rope_cache, apply_rope_partial, apply_inverse_rope, + load_weights_to_cpu, get_layer_weights +) +from single_shot_inference import mHCBlock + + +def main(): + device = 'cuda:0' + with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f: + cfg = json.load(f) + + n_layers = cfg["num_hidden_layers"] + H = cfg["hidden_size"] + n_h = cfg["num_attention_heads"] + hd = cfg["head_dim"] + rd = cfg.get("qk_rope_head_dim", cfg.get("rope_dim", 64)) + n_hc = 4 + o_rank = cfg.get("output_group_dim", 1024) + o_groups = cfg.get("num_output_groups", 16) + heads_per_group = n_h // o_groups + group_input_dim = heads_per_group * hd + + print("Loading weights to CPU...") + all_weights = load_weights_to_cpu(CHECKPOINT_DIR) + + embed_w = all_weights.get("model.embed_tokens.weight") + embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to(device)) + + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) + input_ids = tokenizer.encode("The") + tid = torch.tensor([input_ids[0]], dtype=torch.long, device=device) + positions = torch.tensor([0], dtype=torch.long, device=device) + + emb = embed(tid) + from dsv4.layers.mhc import mHCLayer + X = mHCLayer.init_state(emb, n_hc) + + print(f"\n{'L':>3} {'|X_in|':>10} {'|x_norm|':>10} {'|F_attn|':>10} {'|X_mid|':>10} {'|F_ffn|':>10} {'|X_out|':>10} nan? inf?") + + for li in range(n_layers): + gpu = li % NUM_GPUS + dev = f"cuda:{gpu}" + X = X.to(dev) + torch.cuda.set_device(gpu) + + w = get_layer_weights(all_weights, li, dev) + rope_cos, rope_sin = build_rope_cache(8192, rd, dev) + positions_dev = positions.to(dev) + + # Build per-layer components + attn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev) + attn_mhc.load_from_checkpoint( + all_weights[f"model.layers.{li}.attn_hc.fn"], + all_weights[f"model.layers.{li}.attn_hc.base"], + all_weights[f"model.layers.{li}.attn_hc.scale"]) + ffn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev) + ffn_mhc.load_from_checkpoint( + all_weights[f"model.layers.{li}.ffn_hc.fn"], + all_weights[f"model.layers.{li}.ffn_hc.base"], + all_weights[f"model.layers.{li}.ffn_hc.scale"]) + attn_norm = RMSNorm(H, eps=cfg.get('rms_norm_eps', 1e-6), device=dev) + attn_norm.weight = all_weights[f"model.layers.{li}.input_layernorm.weight"].to(device=dev, dtype=torch.float32) + ffn_norm = RMSNorm(H, eps=cfg.get('rms_norm_eps', 1e-6), device=dev) + ffn_norm.weight = all_weights[f"model.layers.{li}.post_attention_layernorm.weight"].to(device=dev, dtype=torch.float32) + + # Track magnitudes + x_in_mag = X.float().abs().max().item() + + # ATTENTION + x_in_attn, attn_ctx = attn_mhc.pre_block(X) + x_normed = attn_norm.forward(x_in_attn) + x_normed_mag = x_normed.float().abs().max().item() + + pre = f"model.layers.{li}.self_attn" + + # Q projection + c_Q = nvfp4_linear(x_normed, w[f"{pre}.q_a_proj.weight"], + w[f"{pre}.q_a_proj.weight_scale"], w[f"{pre}.q_a_proj.weight_scale_2"]) + q_norm_w = w.get(f"{pre}.q_a_norm.weight") + if q_norm_w is not None: + c_Q_f = c_Q.float() + c_Q = (c_Q_f * c_Q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() * q_norm_w.float()).bfloat16() + q = nvfp4_linear(c_Q, w[f"{pre}.q_b_proj.weight"], + w[f"{pre}.q_b_proj.weight_scale"], w[f"{pre}.q_b_proj.weight_scale_2"]) + + kv = nvfp4_linear(x_normed, w[f"{pre}.kv_proj.weight"], + w[f"{pre}.kv_proj.weight_scale"], w[f"{pre}.kv_proj.weight_scale_2"]) + kv_norm_w = w.get(f"{pre}.kv_norm.weight") + if kv_norm_w is not None: + kv_f = kv.float() + kv = (kv_f * kv_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() * kv_norm_w.float()).bfloat16() + + q_heads = q.reshape(1, n_h, hd) + kv_new = kv.reshape(1, 1, hd) + q_roped = apply_rope_partial(q_heads, positions_dev, rope_cos, rope_sin, hd, rd) + kv_roped = apply_rope_partial(kv_new, positions_dev, rope_cos, rope_sin, hd, rd) + + k_exp = kv_roped.expand(n_h, -1, -1).contiguous() + v_exp = kv_roped.expand(n_h, -1, -1).contiguous() + q_input = q_roped.permute(1, 0, 2) + scale = 1.0 / math.sqrt(hd) + + sink_key = f"{pre}.sinks" + if sink_key in w: + sinks = w[sink_key].to(device=dev) + sink_k = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=dev) + sink_v = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=dev) + k_ws = torch.cat([k_exp, sink_k], dim=1) + v_ws = torch.cat([v_exp, sink_v], dim=1) + sink_bias = torch.zeros(n_h, 1, 2, dtype=torch.bfloat16, device=dev) + for h in range(n_h): + sink_bias[h, :, -1] = sinks[h] + attn_out = torch.nn.functional.scaled_dot_product_attention( + q_input, k_ws, v_ws, attn_mask=sink_bias, scale=scale) + else: + attn_out = torch.nn.functional.scaled_dot_product_attention( + q_input, k_exp, v_exp, scale=scale, is_causal=False) + attn_out = attn_out.permute(1, 0, 2) + + attn_out = apply_inverse_rope(attn_out, positions_dev, rope_cos, rope_sin, hd, rd) + + # Output projection + attn_flat = attn_out.reshape(1, n_h * hd) + attn_grouped = attn_flat.reshape(1, o_groups, heads_per_group * hd) + oa_w = w[f"{pre}.o_a_proj.weight"].bfloat16() + oa_3d = oa_w.reshape(o_groups, o_rank, group_input_dim) + grouped_out = torch.bmm(attn_grouped.permute(1, 0, 2), oa_3d.transpose(1, 2)) + grouped_flat = grouped_out.permute(1, 0, 2).reshape(1, o_groups * o_rank) + F_attn = nvfp4_linear(grouped_flat, w[f"{pre}.o_b_proj.weight"], + w[f"{pre}.o_b_proj.weight_scale"], w[f"{pre}.o_b_proj.weight_scale_2"]) + f_attn_mag = F_attn.float().abs().max().item() + + X_mid = attn_mhc.post_block(X, F_attn, attn_ctx) + x_mid_mag = X_mid.float().abs().max().item() + + # FFN (shared expert + routed, using the FIXED MoE loop) + x_ffn, ffn_ctx = ffn_mhc.pre_block(X_mid) + x_ffn_normed = ffn_norm.forward(x_ffn) + + # Routed MoE + n_experts = cfg["n_routed_experts"] + top_k = cfg.get("num_experts_per_tok", 6) + routed_scaling = cfg.get("routed_scaling_factor", 2.5) + swiglu_limit = cfg.get("swiglu_limit", 10.0) + is_hash = li < 3 + + if is_hash: + tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid" + if tid2eid_key in w: + tid2eid = w[tid2eid_key] + tid_val = tid.item() if tid.device == dev else tid.to(dev).item() + expert_ids = tid2eid[tid_val] + expert_weights = torch.ones(top_k, dtype=torch.float32, device=dev) / top_k + else: + is_hash = False + + if not is_hash: + gate_w = w[f"model.layers.{li}.mlp.gate.weight"] + logits = torch.nn.functional.linear(x_ffn_normed, gate_w.bfloat16()) + activated = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6) + e_bias_key = f"model.layers.{li}.mlp.gate.e_bias" + if e_bias_key in w: + activated = activated + w[e_bias_key].float().unsqueeze(0) + scores, indices = activated.topk(top_k, dim=-1) + unbiased = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6) + unbiased_scores = torch.gather(unbiased, -1, indices) + expert_weights = unbiased_scores / unbiased_scores.sum(dim=-1, keepdim=True) + expert_ids = indices[0] + + expert_outputs = [] + for i, eid in enumerate(expert_ids): + eid_int = eid.item() + epre = f"model.layers.{li}.mlp.experts.{eid_int}" + gate = nvfp4_linear(x_ffn_normed, w[f"{epre}.gate_proj.weight"], + w[f"{epre}.gate_proj.weight_scale"], w[f"{epre}.gate_proj.weight_scale_2"]) + up = nvfp4_linear(x_ffn_normed, w[f"{epre}.up_proj.weight"], + w[f"{epre}.up_proj.weight_scale"], w[f"{epre}.up_proj.weight_scale_2"]) + silu_out = torch.nn.functional.silu(gate.float()) + if swiglu_limit is not None: + silu_out = silu_out.clamp(-swiglu_limit, swiglu_limit) + up_clamped = up.float().clamp(-swiglu_limit, swiglu_limit) + else: + up_clamped = up.float() + hidden = (silu_out * up_clamped).bfloat16() + down = nvfp4_linear(hidden, w[f"{epre}.down_proj.weight"], + w[f"{epre}.down_proj.weight_scale"], w[f"{epre}.down_proj.weight_scale_2"]) + expert_outputs.append(down) + + routed_out = torch.zeros_like(x_ffn_normed) + for i, (out, wt) in enumerate(zip(expert_outputs, expert_weights)): + routed_out = routed_out + (out.float() * wt).bfloat16() + routed_out = (routed_out.float() * routed_scaling).bfloat16() + + # Shared expert + se_pre = f"model.layers.{li}.mlp.shared_experts" + gate = nvfp4_linear(x_ffn_normed, w[f"{se_pre}.gate_proj.weight"], + w[f"{se_pre}.gate_proj.weight_scale"], w[f"{se_pre}.gate_proj.weight_scale_2"]) + up = nvfp4_linear(x_ffn_normed, w[f"{se_pre}.up_proj.weight"], + w[f"{se_pre}.up_proj.weight_scale"], w[f"{se_pre}.up_proj.weight_scale_2"]) + hidden = (torch.nn.functional.silu(gate.float()).clamp(-10, 10) * up.float().clamp(-10, 10)).bfloat16() + shared_out = nvfp4_linear(hidden, w[f"{se_pre}.down_proj.weight"], + w[f"{se_pre}.down_proj.weight_scale"], w[f"{se_pre}.down_proj.weight_scale_2"]) + + F_ffn = routed_out + shared_out + f_ffn_mag = F_ffn.float().abs().max().item() + + X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) + x_out_mag = X_next.float().abs().max().item() + + has_nan = torch.isnan(X_next).any().item() + has_inf = torch.isinf(X_next).any().item() + + print(f"{li:3d} {x_in_mag:10.3f} {x_normed_mag:10.3f} {f_attn_mag:10.3f} {x_mid_mag:10.3f} {f_ffn_mag:10.3f} {x_out_mag:10.3f} {'NaN' if has_nan else ''} {'INF' if has_inf else ''}") + + X = X_next + del w + torch.cuda.empty_cache() + + # Final logits + X = X.to('cuda:0') + lm_w = all_weights.get("lm_head.weight", embed_w).bfloat16().to('cuda:0') + final_norm_w = all_weights.get("model.norm.weight") + 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() + x_out = (xf * rms * final_norm_w.to('cuda:0').float()).bfloat16() + logits = torch.nn.functional.linear(x_out, lm_w) + top5_vals, top5_ids = torch.topk(logits[0], 5) + top5_str = ' '.join([f'{tokenizer.decode([t.item()])}({v.item():.1f})' + for t, v in zip(top5_ids, top5_vals)]) + print(f"\nTop-5: {top5_str}") + + +if __name__ == "__main__": + main()