#!/usr/bin/env python3 """ DeepSeek-V4 MoE Runner NaN Test Tests the Nvfp4MoE (grouped GEMM path) with real weights. The single-expert tests pass — this test exercises the FULL MoE runner with routing, padding, grouped GEMM, and combine. Usage (on B200): cd /root/nvfp4-megamoe-kernel PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_moe_runner_nan_b200.py """ import sys, os, json, torch, torch.nn.functional as F from safetensors import safe_open REPO = "/root/nvfp4-megamoe-kernel" sys.path.insert(0, REPO) MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" DEV = "cuda:0" H = 7168 INTERMEDIATE = 3072 NUM_EXPERTS = 384 TOPK = 6 EPS = 1e-6 _cache = {} def P(k, wm, md): if k in _cache: return _cache[k] with safe_open(os.path.join(md, wm[k]), framework="pt") as f: t = f.get_tensor(k) _cache[k] = t return t def rms(x, w, eps=1e-6): v = x.float().pow(2).mean(-1, keepdim=True) return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) def pack_expert_weights(wm, G, layer_id=2, num_local_experts=16): """Pack per-expert weights into stacked format for Nvfp4MoE. Only loads the first num_local_experts to fit in memory. """ m = f"model.layers.{layer_id}.mlp" # Load expert weights and stack (only first num_local_experts) gate_ws, gate_sfs, gate_gss = [], [], [] up_ws, up_sfs, up_gss = [], [], [] down_ws, down_sfs, down_gss = [], [], [] for i in range(num_local_experts): e = f"{m}.experts.{i}" gate_ws.append(G(f"{e}.gate_proj.weight")) gate_sfs.append(G(f"{e}.gate_proj.weight_scale")) gate_gs = G(f"{e}.gate_proj.weight_scale_2") gate_gss.append(gate_gs) up_ws.append(G(f"{e}.up_proj.weight")) up_sfs.append(G(f"{e}.up_proj.weight_scale")) up_gs = G(f"{e}.up_proj.weight_scale_2") up_gss.append(up_gs) down_ws.append(G(f"{e}.down_proj.weight")) down_sfs.append(G(f"{e}.down_proj.weight_scale")) down_gs = G(f"{e}.down_proj.weight_scale_2") down_gss.append(down_gs) if i % 50 == 0: print(f" Loaded expert {i}/{num_local_experts}") # Stack into (E, ...) tensors w13_w = torch.stack(gate_ws) # (E, 3072, 3584) w13_sf = torch.stack(gate_sfs) w13_gs = torch.stack(gate_gss) if gate_gss[0].dim() > 0 else torch.tensor([g.item() for g in gate_gss], device=DEV) # Actually w13 = stacked gate+up, w2 = down # But our runner expects separate L1 (gate+up) and L2 (down) # The w13 format is (E, 2*intermediate, hidden//2) with gate and up interleaved # For Nvfp4MoE, we stack gate and up side-by-side # Stack gate and up into w13 format: (E, 2*intermediate, hidden//2) w13_w = torch.cat([torch.stack(gate_ws), torch.stack(up_ws)], dim=1) # (E, 6144, 3584) w13_sf = torch.cat([torch.stack(gate_sfs), torch.stack(up_sfs)], dim=1) w13_gs = torch.cat([torch.stack(gate_gss), torch.stack(up_gss)], dim=0) w2_w = torch.stack(down_ws) w2_sf = torch.stack(down_sfs) w2_gs = torch.stack(down_gss) return w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs def test_moe_runner(layer_id=2): """Test the Nvfp4MoE with real weights.""" from dsv4.layers.moe import Nvfp4MoE torch.cuda.set_device(0) torch.manual_seed(42) torch.cuda.empty_cache() _cache.clear() with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: wm = json.load(f)["weight_map"] G = lambda k: P(k, wm, MODEL).to(DEV) p = f"model.layers.{layer_id}" m = f"{p}.mlp" emb = G("model.embed_tokens.weight") fnorm = G(f"{p}.post_attention_layernorm.weight") print(f" Packing expert weights (384 experts)...") # Test with fewer experts to fit in memory num_local_experts = 16 # Use 16 experts (out of 384) for testing # Create the runner first, then prepare weights intermediate_size = INTERMEDIATE # 3072 hidden_size = H # 7168 runner = Nvfp4MoE( num_experts=num_local_experts, hidden_size=hidden_size, intermediate_size=intermediate_size, max_num_tokens=8192, top_k=TOPK, device=str(DEV), ) # Load and pack weights print(f" Loading expert weights...") w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs = pack_expert_weights(wm, G, layer_id, num_local_experts) print(f" w13_w: {w13_w.shape}, w2_w: {w2_w.shape}") print(f" w13_gs: {w13_gs.shape}, w2_gs: {w2_gs.shape}") print(f" w13 NaN: {torch.isnan(w13_w.float()).any()}") print(f" w2 NaN: {torch.isnan(w2_w.float()).any()}") # Prepare weights for the runner l1_fp4 = w13_w.view(torch.float4_e2m1fn_x2) l2_fp4 = w2_w.view(torch.float4_e2m1fn_x2) l1_sf = w13_sf.to(torch.float8_e4m3fn) if w13_sf.dtype != torch.float8_e4m3fn else w13_sf l2_sf = w2_sf.to(torch.float8_e4m3fn) if w2_sf.dtype != torch.float8_e4m3fn else w2_sf runner.prepare_weights_from_stacked( l1_fp4, l1_sf, w13_gs.tolist() if w13_gs.dim() == 1 else w13_gs.flatten().tolist(), l2_fp4, l2_sf, w2_gs.tolist() if w2_gs.dim() == 1 else w2_gs.flatten().tolist(), ) # Test with various token counts for num_tokens in [1, 4, 8, 16]: token_ids = torch.randint(1, 1000, (num_tokens,), dtype=torch.long, device=DEV) hidden = emb[token_ids] normed = rms(hidden, fnorm, EPS) topk_ids = torch.randint(0, num_local_experts, (num_tokens, TOPK), device=DEV) print(f" {num_tokens} tokens: input amax={normed.amax():.4f} NaN={torch.isnan(normed).any()}") topk_weights = torch.softmax(torch.randn(num_tokens, TOPK, device=DEV), dim=-1) print(f" {num_tokens} tokens: input amax={normed.amax():.4f} NaN={torch.isnan(normed).any()}") with torch.no_grad(): result = runner.run(normed, topk_weights, topk_ids) result_nan = torch.isnan(result).any().item() result_amax = result.amax().item() if not result_nan else -1 print(f" {num_tokens} tokens: output amax={result_amax:.4f} NaN={result_nan}") if result_nan: nan_rows = torch.isnan(result).any(dim=1).sum().item() print(f" {num_tokens} tokens: {nan_rows}/{num_tokens} rows have NaN") del runner, w13_w, w13_sf, w13_gs, w2_w, w2_sf, w2_gs torch.cuda.empty_cache() _cache.clear() def main(): print("=" * 70) print(" DeepSeek-V4 MoE Runner NaN Test") print(" Tests Nvfp4MoE (grouped GEMM) with real weights") print("=" * 70) test_moe_runner(layer_id=2) print(f"\n{'='*70}") if __name__ == "__main__": main()