From 11dce13afe451c6da1be2161faf834b214d2754e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 22:53:28 +0000 Subject: [PATCH] Add multi-layer pipeline test to check error accumulation --- tests/test_multilayer.py | 154 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 tests/test_multilayer.py diff --git a/tests/test_multilayer.py b/tests/test_multilayer.py new file mode 100644 index 00000000..6a34ffd1 --- /dev/null +++ b/tests/test_multilayer.py @@ -0,0 +1,154 @@ +"""Extended pipeline test: simulate multi-layer MoE to check for error accumulation. +Uses same config as vLLM: max_num_tokens=8192, max_chunks=8, 48 experts.""" +import torch +import torch.nn.functional as F +import sys, os, glob + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')) + +MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +LAYER_IDX = 0 # Use layer 0 weights for all layers (just testing accumulation) +NUM_EXPERTS = 48 +HIDDEN_SIZE = 7168 +INTERMEDIATE_SIZE = 3072 +NUM_TOKENS = 5 # "The capital of France is" +TOP_K = 6 +SWIGLU_LIMIT = 10.0 +DEVICE = "cuda" +NUM_LAYERS = 3 # Test error accumulation over multiple layers + + +def load_layer_tensors(model_dir, layer_idx): + tensors = {} + for sf in glob.glob(os.path.join(model_dir, "*.safetensors")): + from safetensors.torch import load_file + data = load_file(sf) + for k, v in data.items(): + if f"layers.{layer_idx}." in k and "mlp.experts" in k: + tensors[k.removeprefix("model.")] = v + return tensors + + +def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): + lut = torch.tensor([0.,0.5,1.,1.5,2.,3.,4.,6.,-0.,-0.5,-1.,-1.5,-2.,-3.,-4.,-6.], + dtype=torch.float32, device=packed_uint8.device) + lower = lut[(packed_uint8 & 0x0F).long()] + upper = lut[((packed_uint8 >> 4) & 0x0F).long()] + N, K = packed_uint8.shape[0], packed_uint8.shape[1] * 2 + bf16 = torch.stack([lower, upper], dim=-1).reshape(N, K) + K_sf = scale_e4m3.shape[1] + scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1) + return (bf16 * scale_2d * global_scale).to(torch.bfloat16) + + +def main(): + torch.cuda.set_device(0) + torch.manual_seed(42) + + print(f"=== Multi-Layer Pipeline Test ({NUM_LAYERS} layers) ===") + nvfp4_tensors = load_layer_tensors(MODEL_PATH, LAYER_IDX) + expert_indices = list(range(NUM_EXPERTS)) + + # Start with random hidden states (like after embedding + first attention) + hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 + + topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE) + for i in range(NUM_TOKENS): + topk_ids[i] = torch.randperm(NUM_EXPERTS)[:TOP_K] + topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K + + # Setup runner + from vllm.nvfp4_cutedsl import CuTeDSLMoERunner + from cutedsl.bridge import assemble_scales_3d_side, make_b_k_major + + l1_fp4, l1_sf, l1_gs_list = [], [], [] + l2_fp4, l2_sf, l2_gs_list = [], [], [] + for e in expert_indices: + gw = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) + uw = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) + gsf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) + usf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) + ggs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() + ugs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale_2"].item() + fw = torch.cat([gw, uw], dim=0).view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() + fsf = torch.cat([gsf, usf], dim=0).permute(1,0).contiguous() + mgs = max(ggs, ugs) + if ggs != ugs: + sf32 = fsf.float() + sf32[:, :INTERMEDIATE_SIZE] *= (ggs / mgs) + sf32[:, INTERMEDIATE_SIZE:] *= (ugs / mgs) + fsf = sf32.to(torch.float8_e4m3fn) + l1_fp4.append(fw); l1_sf.append(fsf); l1_gs_list.append(mgs) + dk = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" + if dk in nvfp4_tensors: + dw = nvfp4_tensors[dk].to(DEVICE) + dsf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) + dgs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale_2"].item() + l2_fp4.append(dw.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()) + l2_sf.append(dsf.permute(1,0).contiguous()); l2_gs_list.append(dgs) + else: + l2_fp4.append(torch.zeros(INTERMEDIATE_SIZE//2, HIDDEN_SIZE, dtype=torch.float4_e2m1fn_x2, device=DEVICE)) + l2_sf.append(torch.ones(INTERMEDIATE_SIZE//16, HIDDEN_SIZE, dtype=torch.float8_e4m3fn, device=DEVICE)) + l2_gs_list.append(1.0) + + runner = CuTeDSLMoERunner( + num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS, + top_k=TOP_K, device=DEVICE, + ) + runner.l1_fp4 = l1_fp4; runner.l1_sf = l1_sf; runner.l1_gs = l1_gs_list + runner.l2_fp4 = l2_fp4; runner.l2_sf = l2_sf; runner.l2_gs = l2_gs_list + runner.set_swiglu_limit(SWIGLU_LIMIT) + + # Warmup + with torch.no_grad(): + runner.compute_activation_global_scales(hidden, topk_weights, topk_ids) + + # Run multiple layers (using same weights, but hidden evolves) + ref_hidden = hidden.clone() + run_hidden = hidden.clone() + + for layer in range(NUM_LAYERS): + with torch.no_grad(): + # Runner + runner.compute_activation_global_scales(run_hidden, topk_weights, topk_ids) + run_out = runner.run(run_hidden, topk_weights, topk_ids) + run_hidden = run_out # MoE output replaces hidden (simplified, no residual) + + # BF16 reference + ref_out = torch.zeros(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) + for i, e in enumerate(expert_indices): + dk = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" + gk = f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight" + uk = f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight" + if dk not in nvfp4_tensors: + continue + gate_bf16 = dequantize_nvfp4_weight(nvfp4_tensors[gk].to(DEVICE), nvfp4_tensors[gk.replace('.weight', '.weight_scale')].to(DEVICE), nvfp4_tensors[gk.replace('.weight', '.weight_scale_2')].item()) + up_bf16 = dequantize_nvfp4_weight(nvfp4_tensors[uk].to(DEVICE), nvfp4_tensors[uk.replace('.weight', '.weight_scale')].to(DEVICE), nvfp4_tensors[uk.replace('.weight', '.weight_scale_2')].item()) + down_bf16 = dequantize_nvfp4_weight(nvfp4_tensors[dk].to(DEVICE), nvfp4_tensors[dk.replace('.weight', '.weight_scale')].to(DEVICE), nvfp4_tensors[dk.replace('.weight', '.weight_scale_2')].item()) + for t in range(NUM_TOKENS): + for k in range(TOP_K): + if topk_ids[t, k].item() != i: + continue + w = topk_weights[t, k].item() + x = ref_hidden[t] + gate = x @ gate_bf16.T + up = x @ up_bf16.T + gate_silu = F.silu(gate).clamp(max=SWIGLU_LIMIT) + up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT) + act = gate_silu * up + ref_out[t] += w * (act @ down_bf16.T) + ref_hidden = ref_out + + cos = F.cosine_similarity(ref_hidden.flatten().unsqueeze(0), run_hidden.flatten().unsqueeze(0)).item() + has_nan = torch.isnan(run_hidden).any().item() + has_inf = torch.isinf(run_hidden).any().item() + print(f"Layer {layer}: cosine={cos:.6f} ref_amax={ref_hidden.amax().item():.4f} run_amax={run_hidden.amax().item():.4f} NaN={has_nan} Inf={has_inf}") + + if has_nan: + print(f" ❌ NaN detected after layer {layer}! Stopping.") + break + + +if __name__ == "__main__": + main()