From 9d57b0453b78433d58d5bd605cc936c3158b7efe Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 15:04:46 +0000 Subject: [PATCH] auto: pre-test commit --- tests/unit/test_layer_comparison.py | 124 ++++++++++++++++++++ tests/unit/test_mhc_comparison.py | 169 ++++++++++++++++++++++++++++ 2 files changed, 293 insertions(+) create mode 100644 tests/unit/test_layer_comparison.py create mode 100644 tests/unit/test_mhc_comparison.py diff --git a/tests/unit/test_layer_comparison.py b/tests/unit/test_layer_comparison.py new file mode 100644 index 00000000..deeda5fa --- /dev/null +++ b/tests/unit/test_layer_comparison.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Layer-by-layer comparison: production kernel vs PyTorch reference. + +This test loads both pipelines, runs the same input, and compares +hidden states after each layer to find where the residual diverges. +""" +import os, sys, json, time, math, torch, torch.nn.functional as F +from pathlib import Path + +CHECKPOINT_DIR = os.environ.get("CHECKPOINT_DIR", "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4") +DEVICE = "cuda:0" + +def main(): + torch.manual_seed(42) + + # Load config + 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"] + hd = cfg["head_dim"] + n_hc = cfg.get("n_hc", 4) + print(f"Model: {n_layers} layers, {H} hidden, {hd} head_dim, {n_hc} mHC streams") + + # --- Load production pipeline --- + print("\nLoading production pipeline...") + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from single_shot_inference import DSV4Model + prod_model = DSV4Model(CHECKPOINT_DIR, device=DEVICE) + print("Production pipeline loaded.") + + # --- Load PyTorch reference pipeline --- + print("\nLoading PyTorch reference pipeline...") + from single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights, forward_layer, rmsnorm + all_w = load_weights(CHECKPOINT_DIR) + print("Reference pipeline loaded.") + + # --- Same input for both --- + # Use the DeepSeek prompt + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR, trust_remote_code=True) + prompt = "The capital of France is" + ids = tokenizer.encode(prompt, add_special_tokens=False) + # Add chat template + user_token = 128803 + asst_token = 128804 + chat_ids = [user_token] + ids + [asst_token] + print(f"Input: {len(chat_ids)} tokens: {chat_ids}") + + # --- Run production pipeline: prefill --- + print("\n=== Production Pipeline: Prefill ===") + prod_model.kv_cache.reset() + prod_X = None + prod_layer_states = [] # (X_l, X_mid, X_next) per layer + + # Process tokens one at a time (decode style) + for ti, tid in enumerate(chat_ids): + token_id = torch.tensor([[tid]], dtype=torch.int32, device=DEVICE) + if ti == len(chat_ids) - 1: + # Save layer states for the last token + # We need to modify the production pipeline to capture per-layer states + # For now, just run and capture the final output + pass + prod_model.decode_step(token_id, position_offset=ti) + + print("Production prefill done.") + + # --- Run reference pipeline: prefill --- + print("\n=== Reference Pipeline: Prefill ===") + # Initialize mHC state + emb_w = all_w.get("model.embed_tokens.weight") + emb_ref = torch.nn.Embedding(emb_w.shape[0], emb_w.shape[1]) + emb_ref.weight.data = emb_w.bfloat16().to(DEVICE) + + ref_X = mHCBlock.init_state(emb_ref(torch.tensor(chat_ids, device=DEVICE)), n_hc=n_hc) + + # Build mHC blocks and norms for reference + attn_mhcs, ffn_mhcs = [], [] + attn_norms, ffn_norms = [], [] + for li in range(n_layers): + a_mhc = mHCBlock(H, n_hc, device=DEVICE) + a_mhc.load(all_w[f"model.layers.{li}.attn_hc.fn"], + all_w[f"model.layers.{li}.attn_hc.base"], + all_w[f"model.layers.{li}.attn_hc.scale"]) + attn_mhcs.append(a_mhc) + + f_mhc = mHCBlock(H, n_hc, device=DEVICE) + f_mhc.load(all_w[f"model.layers.{li}.ffn_hc.fn"], + all_w[f"model.layers.{li}.ffn_hc.base"], + all_w[f"model.layers.{li}.ffn_hc.scale"]) + ffn_mhcs.append(f_mhc) + + attn_norms.append(all_w[f"model.layers.{li}.input_layernorm.weight"].bfloat16().to(DEVICE)) + ffn_norms.append(all_w[f"model.layers.{li}.post_attention_layernorm.weight"].bfloat16().to(DEVICE)) + + # Run reference layer by layer + print("Running reference layer by layer...") + ref_kv_cache = {} + for li in range(n_layers): + w = all_w + X_before = ref_X.clone() + ref_X = forward_layer(ref_X, w, li, cfg, None, None, + attn_mhcs[li], ffn_mhcs[li], + attn_norms[li], ffn_norms[li], + ref_kv_cache, torch.arange(len(chat_ids), device=DEVICE), + 0) + x_max = ref_X.abs().max().item() + if li % 10 == 0 or li >= 55: + print(f" Ref L{li}: |X|={x_max:.1f}") + + print("Reference prefill done.") + print(f" Final |X|: {ref_X.abs().max().item():.1f}") + + # Compare + # We can't easily compare per-layer because the production pipeline + # doesn't expose intermediate states. But we can compare the final + # hidden state and the decoded token. + + print("\n=== Summary ===") + print(f"Production final |X|: N/A (need to instrument)") + print(f"Reference final |X|: {ref_X.abs().max().item():.1f}") + +if __name__ == "__main__": + main() diff --git a/tests/unit/test_mhc_comparison.py b/tests/unit/test_mhc_comparison.py new file mode 100644 index 00000000..7ed825d6 --- /dev/null +++ b/tests/unit/test_mhc_comparison.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +"""Focused comparison: production MoE vs PyTorch reference MoE at specific layers. + +This test: +1. Loads both pipelines +2. Processes the same input token through 1 layer +3. Compares F_attn and F_ffn magnitudes between production and reference +4. Identifies where the magnitude diverges +""" +import os, sys, json, time, math, torch, torch.nn.functional as F +from pathlib import Path + +CHECKPOINT_DIR = os.environ.get("CHECKPOINT_DIR", "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4") +DEVICE = "cuda:0" +HC_EPS = 1e-6 + +def sinkhorn_knopp(logits, t_max=20, eps=HC_EPS): + M = torch.softmax(logits, -1) + eps + M = M / (M.sum(-2, keepdim=True) + eps) + for _ in range(t_max - 1): + M = M / (M.sum(-1, keepdim=True) + eps) + M = M / (M.sum(-2, keepdim=True) + eps) + return M + +def unweighted_rmsnorm(x, eps=1e-6): + x_f = x.float() + rms = x_f.pow(2).mean(-1, keepdim=True).add(eps).rsqrt() + return (x_f * rms).to(x.dtype) + +def rmsnorm(x, w, eps=1e-6): + x_f = x.float() + rms = x_f.pow(2).mean(-1, keepdim=True).add(eps).rsqrt() + return (x_f * rms * w.float()).to(x.dtype) + +FP4_LUT = torch.tensor([0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + +def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None): + O, I2 = weight.shape; I = I2 * 2 + lo = (weight & 0x0F).to(torch.int8); hi = (weight >> 4).to(torch.int8) + lut = FP4_LUT.to(device=weight.device, dtype=torch.float32) + lo_f = lut[(lo & 0x07).long()] * torch.where((lo >> 3).bool(), -1., 1.) + hi_f = lut[(hi & 0x07).long()] * torch.where((hi >> 3).bool(), -1., 1.) + w = torch.stack([lo_f, hi_f], -1).reshape(O, I) + s = weight_scale.float().repeat_interleave(16, 1) + if weight_scale_2 is not None: s = s * weight_scale_2.float() + return (w * s).bfloat16() + +def main(): + torch.manual_seed(42) + + with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f: + cfg = json.load(f) + H = cfg["hidden_size"] + n_hc = cfg.get("n_hc", 4) + n_layers = cfg["num_hidden_layers"] + n_experts = cfg["n_routed_experts"] + top_k = cfg.get("num_experts_per_tok", 6) + intermediate = cfg.get("intermediate_size", 18432) + print(f"Model: {n_layers} layers, {H} hidden, {n_experts} experts, top-{top_k}") + + # Load weights + print("Loading weights...") + from safetensors.torch import load_file + cdir = Path(CHECKPOINT_DIR); wmap = {} + idx = cdir / "model.safetensors.index.json" + if idx.exists(): + with open(idx) as f: wmap = json.load(f).get("weight_map", {}) + shards = set(wmap.values()) if wmap else set(); all_w = {} + for sn in sorted(shards): + if (cdir / sn).exists(): all_w.update(load_file(str(cdir / sn))) + print(f"Loaded {len(all_w)} tensors") + + # Create a realistic hidden state (simulate running through a few layers) + # Use token embedding + a few layers of mHC + from single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights as ref_load_weights, forward_layer + ref_all_w = ref_load_weights(CHECKPOINT_DIR) + + # Build mHC blocks for first 3 layers + attn_mhcs, ffn_mhcs = [], [] + attn_norms, ffn_norms = [], [] + for li in range(min(5, n_layers)): + a_mhc = mHCBlock(H, n_hc, device=DEVICE) + a_mhc.load(ref_all_w[f"model.layers.{li}.attn_hc.fn"], + ref_all_w[f"model.layers.{li}.attn_hc.base"], + ref_all_w[f"model.layers.{li}.attn_hc.scale"]) + attn_mhcs.append(a_mhc) + f_mhc = mHCBlock(H, n_hc, device=DEVICE) + f_mhc.load(ref_all_w[f"model.layers.{li}.ffn_hc.fn"], + ref_all_w[f"model.layers.{li}.ffn_hc.base"], + ref_all_w[f"model.layers.{li}.ffn_hc.scale"]) + ffn_mhcs.append(f_mhc) + attn_norms.append(ref_all_w[f"model.layers.{li}.input_layernorm.weight"].bfloat16().to(DEVICE)) + ffn_norms.append(ref_all_w[f"model.layers.{li}.post_attention_layernorm.weight"].bfloat16().to(DEVICE)) + + # Process one token through first 3 layers to get a realistic X state + emb_w = ref_all_w["model.embed_tokens.weight"] + emb = torch.nn.Embedding(emb_w.shape[0], emb_w.shape[1]) + emb.weight.data = emb_w.bfloat16().to(DEVICE) + + # "The" token + tid = 455 + X = mHCBlock.init_state(emb(torch.tensor([tid], device=DEVICE)), n_hc=n_hc) + print(f"\nInitial |X| = {X.abs().max().item():.2f}") + + # Run through first 3 layers using reference + kv_cache = {} + for li in range(3): + X = forward_layer(X, ref_all_w, li, cfg, None, None, + attn_mhcs[li], ffn_mhcs[li], + attn_norms[li], ffn_norms[li], + kv_cache, torch.tensor([3], device=DEVICE), + tid) + print(f" Ref L{li}: |X| = {X.abs().max().item():.2f}") + + # Now X is a realistic hidden state after 3 layers + # Save it for both production and reference comparison + X_ref = X.clone() + X_prod = X.clone() + print(f"\nAfter 3 layers: |X| = {X_ref.abs().max().item():.2f}") + + # --- Compare mHC at L3 --- + li = 3 + print(f"\n=== Comparing mHC at L{li} ===") + + # Reference mHC + a_mhc = attn_mhcs[3] # Already loaded + x_in_ref, ctx_ref = a_mhc.pre_block(X_ref) + print(f" Ref x_in: |x| = {x_in_ref.abs().max().item():.4f}") + print(f" Ref A: {ctx_ref['A'][0].tolist()}") + print(f" Ref C: {ctx_ref['C'][0].tolist()}") + print(f" Ref B row_sums: {ctx_ref['B'][0].sum(-1).tolist()}") + + # Production mHC + from dsv4.layers.mhc import mHCLayer + prod_mhc = mHCLayer(hidden_dim=H, n_hc=n_hc, device=DEVICE) + # Load weights + fn = ref_all_w[f"model.layers.{li}.attn_hc.fn"].to(DEVICE, torch.float32) + base = ref_all_w[f"model.layers.{li}.attn_hc.base"].to(DEVICE) + scale = ref_all_w[f"model.layers.{li}.attn_hc.scale"].to(DEVICE) + n = n_hc + prod_mhc.load_weights( + W_pre=fn[0:n], W_post=fn[n:2*n], W_comb=fn[2*n:], + S_pre=base[0:n].reshape(1, n), S_post=base[n:2*n].reshape(n, 1), + S_comb=base[2*n:].reshape(n, n), + alpha_pre=scale[0].item(), alpha_post=scale[1].item(), alpha_comb=scale[2].item() + ) + x_in_prod, ctx_prod = prod_mhc.pre_block(X_prod) + print(f" Prod x_in: |x| = {x_in_prod.abs().max().item():.4f}") + A_prod = ctx_prod.A_l + C_prod = ctx_prod.C_l + B_prod = ctx_prod.B_l + print(f" Prod A: {A_prod[0].tolist()}") + print(f" Prod C: {C_prod[0].tolist()}") + print(f" Prod B row_sums: {B_prod[0].sum(-1).tolist()}") + + # Compare + cos_xin = F.cosine_similarity(x_in_ref.flatten().float(), x_in_prod.flatten().float(), dim=0).item() + cos_A = F.cosine_similarity(ctx_ref['A'].flatten().float(), A_prod.flatten().float(), dim=0).item() + cos_C = F.cosine_similarity(ctx_ref['C'].flatten().float(), C_prod.flatten().float(), dim=0).item() + cos_B = F.cosine_similarity(ctx_ref['B'].flatten().float(), B_prod.flatten().float(), dim=0).item() + print(f"\n cos(x_in): {cos_xin:.6f}") + print(f" cos(A): {cos_A:.6f}") + print(f" cos(C): {cos_C:.6f}") + print(f" cos(B): {cos_B:.6f}") + + print("\nDone.") + +if __name__ == "__main__": + main()