biondizzle
fd59222fc0
The fold block_sf (float8) * global_sf (float32) -> float8 loses ~25% precision. Product of ~56-448 block_sf * ~4.65e-05 global_sf lands in float8 low-precision zone where step size is 25%. This makes model output garbage despite finite values. Fix: keep block scales as original float8, return global scales separately as float32 per-expert vectors. Apply global scale as per-expert GEMM alpha in cutlass_grouped_nvfp4_gemm (already iterates per-expert). For L1 with separate gate/up global scales, use gate_gs as alpha and apply up_correction ratio to the up half post-GEMM. weight_transform.py: no more _fold_global_scale, returns (w, sf, global_sf) nvfp4_mega_moe.py: per-expert alpha = activation_gs * weight_gs kernel.py: per_expert_alpha parameter in grouped GEMM deepseek_v4.py: updated type hints and comments
280 lines
12 KiB
Python
280 lines
12 KiB
Python
"""
|
|
NVFP4 MegaMoE Diagnostic — B200
|
|
Checks:
|
|
1. weight_scale_2 values (are they nonzero / loaded correctly?)
|
|
2. Folded scale ranges (clamp/precision loss)
|
|
3. L2 weight/SF orientation sanity
|
|
4. Dequant reference vs CUTLASS output comparison
|
|
5. Single-expert, single-layer test
|
|
"""
|
|
import torch
|
|
import sys
|
|
import os
|
|
import json
|
|
from pathlib import Path
|
|
|
|
MODEL_PATH = "/model" # inside the container
|
|
|
|
def inspect_checkpoint_scales():
|
|
"""Check raw checkpoint weight_scale_2 values."""
|
|
from safetensors import safe_open
|
|
import glob
|
|
|
|
print("=" * 60)
|
|
print("CHECK 1: Checkpoint weight_scale_2 Values")
|
|
print("=" * 60)
|
|
|
|
# Find checkpoint files
|
|
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
|
|
print(f"Found {len(ckpt_files)} safetensors files")
|
|
|
|
# Look for expert weight_scale_2 params
|
|
w13_gs_found = 0
|
|
w2_gs_found = 0
|
|
w13_gs_values = {}
|
|
w2_gs_values = {}
|
|
|
|
for f in ckpt_files:
|
|
with safe_open(f, framework="pt") as st:
|
|
for key in st.keys():
|
|
if "weight_scale_2" in key and ("experts" in key or "ffn" in key):
|
|
val = st.get_tensor(key)
|
|
if "w13" in key or "gate_up" in key or "w1" in key or "w3" in key:
|
|
w13_gs_found += 1
|
|
if w13_gs_found <= 3:
|
|
w13_gs_values[key] = {"shape": list(val.shape), "dtype": str(val.dtype),
|
|
"min": val.float().min().item(), "max": val.float().max().item(),
|
|
"mean": val.float().mean().item()}
|
|
elif "w2" in key or "down" in key:
|
|
w2_gs_found += 1
|
|
if w2_gs_found <= 3:
|
|
w2_gs_values[key] = {"shape": list(val.shape), "dtype": str(val.dtype),
|
|
"min": val.float().min().item(), "max": val.float().max().item(),
|
|
"mean": val.float().mean().item()}
|
|
|
|
print(f"w13 weight_scale_2 entries: {w13_gs_found}")
|
|
print(f"w2 weight_scale_2 entries: {w2_gs_found}")
|
|
for k, v in w13_gs_values.items():
|
|
print(f" {k}: {v}")
|
|
for k, v in w2_gs_values.items():
|
|
print(f" {k}: {v}")
|
|
|
|
return w13_gs_found > 0 and w2_gs_found > 0
|
|
|
|
|
|
def inspect_loaded_model():
|
|
"""Check the model's weight_scale_2 after loading (before finalize_weights)."""
|
|
print("\n" + "=" * 60)
|
|
print("CHECK 2: Model weight_scale_2 After Loading")
|
|
print("=" * 60)
|
|
|
|
# We need to load the model and inspect before finalize_weights nukes the params
|
|
# The vLLM server is already running, so let's check the live model
|
|
# Actually, let's load a fresh model instance for inspection
|
|
|
|
# Simpler approach: just check the checkpoint directly for scale_2
|
|
# The real check is whether finalize_weights gets called with nonzero scale_2
|
|
print(" (Checkpoint inspection is more reliable — see CHECK 1)")
|
|
print(" The [SF-DEBUG] prints from weight_transform.py should also show this")
|
|
|
|
|
|
def check_fold_precision_real():
|
|
"""Check float8 folding precision with real checkpoint scales."""
|
|
print("\n" + "=" * 60)
|
|
print("CHECK 3: Float8 Folding Precision (Real Scales)")
|
|
print("=" * 60)
|
|
|
|
from safetensors import safe_open
|
|
import glob
|
|
|
|
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
|
|
|
|
# Find one layer's expert scales
|
|
for f in ckpt_files:
|
|
with safe_open(f, framework="pt") as st:
|
|
keys = list(st.keys())
|
|
# Find w2 weight_scale and weight_scale_2 for layer 0
|
|
w2_sf_key = None
|
|
w2_gs_key = None
|
|
w13_sf_key = None
|
|
w13_gs_key = None
|
|
|
|
for k in keys:
|
|
if "layers.0" in k:
|
|
if "w2" in k and k.endswith("weight_scale") and "scale_2" not in k:
|
|
w2_sf_key = k
|
|
elif "w2" in k and "weight_scale_2" in k:
|
|
w2_gs_key = k
|
|
elif ("w13" in k or "gate_up" in k) and k.endswith("weight_scale") and "scale_2" not in k:
|
|
w13_sf_key = k
|
|
elif ("w13" in k or "gate_up" in k) and "weight_scale_2" in k:
|
|
w13_gs_key = k
|
|
|
|
if w2_sf_key and w2_gs_key:
|
|
w2_sf = st.get_tensor(w2_sf_key)
|
|
w2_gs = st.get_tensor(w2_gs_key)
|
|
print(f" L2 block scale: shape={list(w2_sf.shape)} dtype={w2_sf.dtype} "
|
|
f"range=[{w2_sf.float().min():.4e}, {w2_sf.float().max():.4e}]")
|
|
print(f" L2 global scale: shape={list(w2_gs.shape)} dtype={w2_gs.dtype} "
|
|
f"range=[{w2_gs.float().min():.4e}, {w2_gs.float().max():.4e}]")
|
|
|
|
# Fold and check precision
|
|
sf_f32 = w2_sf.float()
|
|
gs_f32 = w2_gs.float()
|
|
|
|
# Reshape gs for broadcast
|
|
while gs_f32.dim() < sf_f32.dim():
|
|
gs_f32 = gs_f32.unsqueeze(-1)
|
|
|
|
product = sf_f32 * gs_f32
|
|
product_clamped = product.clamp(0.0, 448.0)
|
|
folded_f8 = product_clamped.to(torch.float8_e4m3fn)
|
|
folded_back = folded_f8.float()
|
|
|
|
# Stats
|
|
n_clamped = (product > 448.0).sum().item()
|
|
n_total = product.numel()
|
|
n_zeroed = (folded_back == 0.0).sum().item()
|
|
|
|
rel_err = (folded_back - product).abs() / product.clamp(min=1e-10)
|
|
print(f"\n L2 Fold results:")
|
|
print(f" Clamped to 448: {n_clamped}/{n_total} ({100*n_clamped/n_total:.1f}%)")
|
|
print(f" Zeroed (subnormal): {n_zeroed}/{n_total} ({100*n_zeroed/n_total:.1f}%)")
|
|
print(f" Rel error: max={rel_err.max():.4f} mean={rel_err.mean():.4f} p99={rel_err.quantile(0.99):.4f}")
|
|
|
|
# Show distribution of folded values
|
|
fb_hist = torch.histc(folded_back, bins=10, min=0, max=448)
|
|
print(f" Folded value histogram (0-448, 10 bins): {fb_hist.int().tolist()}")
|
|
|
|
# CRITICAL CHECK: is the product range within float8?
|
|
print(f" Product range: [{product.min():.4e}, {product.max():.4e}]")
|
|
if n_clamped > 0:
|
|
print(f" ⚠️ {n_clamped} values clamped — this IS precision loss!")
|
|
|
|
if w13_sf_key and w13_gs_key:
|
|
w13_sf = st.get_tensor(w13_sf_key)
|
|
w13_gs = st.get_tensor(w13_gs_key)
|
|
print(f"\n L1 block scale: shape={list(w13_sf.shape)} dtype={w13_sf.dtype} "
|
|
f"range=[{w13_sf.float().min():.4e}, {w13_sf.float().max():.4e}]")
|
|
print(f" L1 global scale: shape={list(w13_gs.shape)} dtype={w13_gs.dtype} "
|
|
f"range=[{w13_gs.float().min():.4e}, {w13_gs.float().max():.4e}]")
|
|
|
|
break # Just check one file that has layer 0
|
|
|
|
|
|
def check_l2_weight_semantics():
|
|
"""Verify L2 weight layout by dequantizing and checking against reference."""
|
|
print("\n" + "=" * 60)
|
|
print("CHECK 4: L2 Weight Dequantization Sanity")
|
|
print("=" * 60)
|
|
|
|
from safetensors import safe_open
|
|
import glob
|
|
|
|
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
|
|
|
|
for f in ckpt_files:
|
|
with safe_open(f, framework="pt") as st:
|
|
keys = list(st.keys())
|
|
# Find layer 0 w2 weight, weight_scale, weight_scale_2
|
|
w2_w = w2_sf = w2_gs = None
|
|
for k in keys:
|
|
if "layers.0" in k:
|
|
if "w2" in k and k.endswith(".weight") and "scale" not in k:
|
|
w2_w = st.get_tensor(k)
|
|
elif "w2" in k and "weight_scale" == k.split(".")[-1]:
|
|
w2_sf = st.get_tensor(k)
|
|
elif "w2" in k and "weight_scale_2" in k:
|
|
w2_gs = st.get_tensor(k)
|
|
|
|
if w2_w is not None and w2_sf is not None and w2_gs is not None:
|
|
print(f" w2_weight: shape={list(w2_w.shape)} dtype={w2_w.dtype}")
|
|
print(f" w2_weight_scale: shape={list(w2_sf.shape)} dtype={w2_sf.dtype}")
|
|
print(f" w2_weight_scale_2: shape={list(w2_gs.shape)} dtype={w2_gs.dtype}")
|
|
|
|
# Dequantize a small patch
|
|
# w2 is down_proj: (hidden, intermediate) in BF16, or (hidden, inter//2) uint8 for NVFP4
|
|
if w2_w.dtype == torch.uint8:
|
|
# Unpack E2M1
|
|
FP4_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=w2_w.device)
|
|
lower = FP4_LUT[(w2_w[:4, :8] & 0x0F).long()]
|
|
upper = FP4_LUT[((w2_w[:4, :8] >> 4) & 0x0F).long()]
|
|
unpacked = torch.empty(4, 16, dtype=torch.float32)
|
|
unpacked[:, 0::2] = lower
|
|
unpacked[:, 1::2] = upper
|
|
|
|
# Apply scales
|
|
sf_slice = w2_sf[:4, :1].float() # (4, 1)
|
|
gs = w2_gs.float()
|
|
print(f" Dequantized w2[:4, :16] with sf[:4,:1]={sf_slice.flatten().tolist()}")
|
|
print(f" global_scale_2 = {gs.item() if gs.numel() == 1 else gs[:4].flatten().tolist()}")
|
|
dequant = unpacked * sf_slice * gs.float()
|
|
print(f" Dequantized range: [{dequant.min():.4f}, {dequant.max():.4f}]")
|
|
print(f" Dequantized[:2, :8]: {dequant[:2, :8].tolist()}")
|
|
else:
|
|
print(f" w2_weight is {w2_w.dtype}, not uint8 — may be BF16 checkpoint")
|
|
print(f" w2[:4, :8] = {w2_w[:4, :8].tolist()}")
|
|
break
|
|
|
|
|
|
def check_ep_reduce_contract():
|
|
"""Verify the EP all-reduce contract with a synthetic test."""
|
|
print("\n" + "=" * 60)
|
|
print("CHECK 5: EP Reduce Contract (Synthetic)")
|
|
print("=" * 60)
|
|
|
|
# Simulate 2 ranks
|
|
M, HIDDEN = 4, 8
|
|
# Rank 0: experts 0,1 — tokens routed to expert 0 (slot_weight=0.7) and 1 (slot_weight=0.3)
|
|
y0 = torch.zeros(M, HIDDEN, dtype=torch.bfloat16)
|
|
slot_token_0 = torch.tensor([0, 0, 1, 2, 3]) # which tokens
|
|
slot_weight_0 = torch.tensor([0.7, 0.3, 0.5, 0.6, 0.4], dtype=torch.bfloat16)
|
|
l2_slots_0 = torch.randn(5, HIDDEN, dtype=torch.bfloat16)
|
|
y0.index_add_(0, slot_token_0, l2_slots_0 * slot_weight_0.unsqueeze(1))
|
|
|
|
# Rank 1: experts 2,3 — token 0 also routed to expert 2
|
|
y1 = torch.zeros(M, HIDDEN, dtype=torch.bfloat16)
|
|
slot_token_1 = torch.tensor([0, 1])
|
|
slot_weight_1 = torch.tensor([0.2, 0.5], dtype=torch.bfloat16)
|
|
l2_slots_1 = torch.randn(2, HIDDEN, dtype=torch.bfloat16)
|
|
y1.index_add_(0, slot_token_1, l2_slots_1 * slot_weight_1.unsqueeze(1))
|
|
|
|
# All-reduce (sum)
|
|
y_final = y0 + y1 # simulated all-reduce
|
|
|
|
# Verify: token 0 should have contributions from rank0 (experts 0,1) and rank1 (expert 2)
|
|
expected_0 = (0.7 * l2_slots_0[0] + 0.3 * l2_slots_0[1] + 0.2 * l2_slots_1[0]).bfloat16()
|
|
actual_0 = y_final[0].bfloat16()
|
|
diff = (expected_0 - actual_0).abs().max().item()
|
|
print(f" Token 0: expected vs actual diff = {diff:.6f} ✓" if diff < 0.01 else f" Token 0: MISMATCH diff = {diff}")
|
|
print(f" EP reduce contract is correct — sum of partial rank outputs gives full result")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("NVFP4 MegaMoE Diagnostic — B200")
|
|
print(f"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}")
|
|
print(f"GPUs: {torch.cuda.device_count()}")
|
|
print()
|
|
|
|
try:
|
|
inspect_checkpoint_scales()
|
|
except Exception as e:
|
|
print(f"CHECK 1 FAILED: {e}")
|
|
|
|
try:
|
|
check_fold_precision_real()
|
|
except Exception as e:
|
|
print(f"CHECK 3 FAILED: {e}")
|
|
|
|
try:
|
|
check_l2_weight_semantics()
|
|
except Exception as e:
|
|
print(f"CHECK 4 FAILED: {e}")
|
|
|
|
try:
|
|
check_ep_reduce_contract()
|
|
except Exception as e:
|
|
print(f"CHECK 5 FAILED: {e}")
|