Add pipeline test with real model weights, add swiglu_limit to reference moe_pipeline
This commit is contained in:
@@ -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)
|
||||
|
||||
# ════════════════════════════════════════════════════════════════
|
||||
|
||||
191
tests/test_pipeline_real_weights.py
Normal file
191
tests/test_pipeline_real_weights.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user