From e38d60a6e888abcb02f758f26d8fc1c1bf467360 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 18:07:44 +0000 Subject: [PATCH] Add pipeline test with real model weights, add swiglu_limit to reference moe_pipeline --- cutedsl/moe_pipeline.py | 7 +- tests/test_pipeline_real_weights.py | 191 ++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 tests/test_pipeline_real_weights.py diff --git a/cutedsl/moe_pipeline.py b/cutedsl/moe_pipeline.py index 899a4708..ac819a6e 100644 --- a/cutedsl/moe_pipeline.py +++ b/cutedsl/moe_pipeline.py @@ -126,6 +126,7 @@ def run_nvfp4_moe( expert_weights, # (num_tokens, top_k) float32 weights, # dict from prepare_nvfp4_moe_weights expert_indices, # list of expert IDs + swiglu_limit=None, # Optional clamp for SiLU output ): """Run the full NVFP4 MoE forward pass. @@ -209,7 +210,11 @@ def run_nvfp4_moe( up = l1_out[:, intermediate_size:] print(f" gate: shape={gate.shape}, amax={gate.abs().amax().item():.4f}", flush=True) print(f" up: shape={up.shape}, amax={up.abs().amax().item():.4f}", flush=True) - activated = torch.nn.functional.silu(gate) * up # (num_slots, intermediate) BF16 + gate_silu = torch.nn.functional.silu(gate) + if swiglu_limit is not None: + gate_silu = gate_silu.clamp(max=swiglu_limit) + up = up.clamp(min=-swiglu_limit, max=swiglu_limit) + activated = gate_silu * up # (num_slots, intermediate) BF16 print(f" After SiLU(gate)*up: shape={activated.shape}, amax={activated.abs().amax().item():.4f}", flush=True) # ════════════════════════════════════════════════════════════════ diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py new file mode 100644 index 00000000..0f44fd09 --- /dev/null +++ b/tests/test_pipeline_real_weights.py @@ -0,0 +1,191 @@ +"""Test #2: End-to-end single-layer test with real model weights. + +Loads layer 0 from the DeepSeek-V4-Pro-NVFP4 checkpoint, runs one MoE layer +through our CuTeDSL runner, and compares against the reference moe_pipeline +(which uses the same NVFP4 weights but with dynamic gs). + +This catches issues that the small layertest (3 experts, 8 tokens) misses: +- Scale assembly with 48 experts × 8 chunks +- Uneven expert assignment +- Real activation magnitudes +- swiglu_limit clamping +- Variable padded expert offsets at scale +""" +import torch +import sys +import os +import math + +# Add paths +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/../cutedsl') +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/../vllm') + +from cutedsl.moe_pipeline import moe_pipeline +from vllm.nvfp4_cutedsl import CuTeDSLMoERunner + +# ============================================================ +# CONFIG +# ============================================================ +MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +LAYER_IDX = 0 +NUM_EXPERTS = 48 # local experts per EP rank +HIDDEN_SIZE = 7168 +INTERMEDIATE_SIZE = 18432 +NUM_TOKENS = 64 +TOP_K = 6 +SWIGLU_LIMIT = 10.0 +DEVICE = "cuda" + +def load_expert_weights(layer_idx, num_experts): + """Load NVFP4 weights for one layer from the checkpoint.""" + from safetensors import safe_open + + # Find the layer shard file + shard_dir = os.path.join(MODEL_PATH, f"model-0000{layer_idx+1:02d}-of-00010.safetensors") + # Try to find the right shard + import glob + shards = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors"))) + + l1_fp4 = [] + l1_sf = [] + l1_gs = [] + l2_fp4 = [] + l2_sf = [] + l2_gs = [] + + for shard_path in shards: + with safe_open(shard_path, framework="pt", device="cpu") as f: + for e in range(num_experts): + global_e = e # For rank 0, local = global + + # L1 (gate+up) + w13_key = f"model.layers.{layer_idx}.mlp.experts.{global_e}.w13_weight" + sf13_key = f"model.layers.{layer_idx}.mlp.experts.{global_e}.w13_weight_scale" + gs13_key = f"model.layers.{layer_idx}.mlp.experts.{global_e}.w13_weight_scale_2" + + # L2 (down) + w2_key = f"model.layers.{layer_idx}.mlp.experts.{global_e}.w2_weight" + sf2_key = f"model.layers.{layer_idx}.mlp.experts.{global_e}.w2_weight_scale" + gs2_key = f"model.layers.{layer_idx}.mlp.experts.{global_e}.w2_weight_scale_2" + + if w13_key in f.keys(): + l1_fp4.append(f.get_tensor(w13_key).to(DEVICE)) + l1_sf.append(f.get_tensor(sf13_key).to(DEVICE)) + l1_gs.append(f.get_tensor(gs13_key).to(DEVICE)) + l2_fp4.append(f.get_tensor(w2_key).to(DEVICE)) + l2_sf.append(f.get_tensor(sf2_key).to(DEVICE)) + l2_gs.append(f.get_tensor(gs2_key).to(DEVICE)) + + if len(l1_fp4) == num_experts: + break + + if len(l1_fp4) != num_experts: + raise RuntimeError(f"Only loaded {len(l1_fp4)}/{num_experts} experts from checkpoint") + + return { + 'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, + 'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs, + } + + +def main(): + torch.cuda.set_device(0) + torch.manual_seed(42) + + print(f"=== Pipeline Test: Layer {LAYER_IDX}, {NUM_EXPERTS} experts, {NUM_TOKENS} tokens, top_k={TOP_K} ===") + + # Load real weights + print("Loading weights from checkpoint...") + weights = load_expert_weights(LAYER_IDX, NUM_EXPERTS) + print(f"Loaded {NUM_EXPERTS} experts") + + # Create runner + 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, + ) + + # Set weights + runner.l1_fp4 = weights['l1_fp4'] + runner.l1_sf = weights['l1_sf'] + runner.l1_gs = weights['l1_gs'] + runner.l2_fp4 = weights['l2_fp4'] + runner.l2_sf = weights['l2_sf'] + runner.l2_gs = weights['l2_gs'] + runner.set_swiglu_limit(SWIGLU_LIMIT) + + # Create input + hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) + + # Create top-k assignments (realistic: uneven distribution) + topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE) + for i in range(NUM_TOKENS): + # Each token picks TOP_K random experts + experts = torch.randperm(NUM_EXPERTS)[:TOP_K] + topk_ids[i] = experts + topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K + + # ---- Stage 1: Reference pipeline (dynamic gs) ---- + print("\n--- Reference pipeline (dynamic gs) ---") + with torch.no_grad(): + ref_out = moe_pipeline( + hidden_states=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + l1_fp4=weights['l1_fp4'], + l1_sf=weights['l1_sf'], + l1_gs=weights['l1_gs'], + l2_fp4=weights['l2_fp4'], + l2_sf=weights['l2_sf'], + l2_gs=weights['l2_gs'], + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + swiglu_limit=SWIGLU_LIMIT, + ) + print(f"Reference: shape={ref_out.shape} amax={ref_out.amax().item():.4f} mean={ref_out.mean().item():.4f}") + print(f" NaN: {torch.isnan(ref_out).any().item()} Inf: {torch.isinf(ref_out).any().item()}") + + # ---- Stage 2: Runner with warmup gs ---- + print("\n--- Runner (warmup gs) ---") + with torch.no_grad(): + # Compute warmup gs + runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids) + print(f"Warmup gs: L1={runner._l1_activation_global_scale:.6f} L2={runner._l2_activation_global_scale:.6f}") + + # Run + runner_out = runner.run(hidden_states, topk_weights, topk_ids) + print(f"Runner: shape={runner_out.shape} amax={runner_out.amax().item():.4f} mean={runner_out.mean().item():.4f}") + print(f" NaN: {torch.isnan(runner_out).any().item()} Inf: {torch.isinf(runner_out).any().item()}") + + # ---- Comparison ---- + print("\n--- Comparison ---") + # Overall cosine + cos = torch.nn.functional.cosine_similarity( + ref_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0) + ).item() + mse = (ref_out - runner_out).pow(2).mean().item() + print(f"Cosine: {cos:.6f} MSE: {mse:.4f}") + + if cos < 0.90: + print("\n⚠️ LOW COSINE — investigating per-token differences...") + for i in range(min(NUM_TOKENS, 8)): + cos_i = torch.nn.functional.cosine_similarity( + ref_out[i].unsqueeze(0), runner_out[i].unsqueeze(0) + ).item() + print(f" Token {i}: cosine={cos_i:.4f} ref_max={ref_out[i].amax().item():.4f} run_max={runner_out[i].amax().item():.4f}") + + if cos >= 0.98: + print(f"\n✅ PASS: cosine {cos:.6f} >= 0.98") + elif cos >= 0.90: + print(f"\n⚠️ MARGINAL: cosine {cos:.6f} — close but degraded") + else: + print(f"\n❌ FAIL: cosine {cos:.6f} < 0.90 — significant quality loss") + + +if __name__ == "__main__": + main()