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
97 lines
4.8 KiB
Python
97 lines
4.8 KiB
Python
"""
|
|
Critical check: weight_scale_2 values are ~4.65e-05 (TINY).
|
|
When folded: block_sf * 4.65e-05 → most products near zero → float8 can't represent
|
|
This is likely THE bug: folding a float8 scale by a tiny global scale produces
|
|
subnormal/zero values in float8.
|
|
"""
|
|
from safetensors import safe_open
|
|
import glob
|
|
import os
|
|
import torch
|
|
|
|
MODEL_PATH = "/model"
|
|
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
|
|
|
|
# Get layer 0, expert 0 scales
|
|
for f in ckpt_files:
|
|
with safe_open(f, framework="pt") as st:
|
|
keys = list(st.keys())
|
|
if any("layers.0.mlp.experts.0.gate_proj.weight_scale" in k for k in keys):
|
|
# Gate
|
|
gate_sf = st.get_tensor("model.layers.0.mlp.experts.0.gate_proj.weight_scale")
|
|
gate_gs = st.get_tensor("model.layers.0.mlp.experts.0.gate_proj.weight_scale_2")
|
|
# Up
|
|
up_sf = st.get_tensor("model.layers.0.mlp.experts.0.up_proj.weight_scale")
|
|
up_gs = st.get_tensor("model.layers.0.mlp.experts.0.up_proj.weight_scale_2")
|
|
# Down
|
|
down_sf = st.get_tensor("model.layers.0.mlp.experts.0.down_proj.weight_scale")
|
|
down_gs = st.get_tensor("model.layers.0.mlp.experts.0.down_proj.weight_scale_2")
|
|
|
|
print("=" * 60)
|
|
print("LAYER 0, EXPERT 0 — Scale Analysis")
|
|
print("=" * 60)
|
|
|
|
for name, sf, gs in [("gate", gate_sf, gate_gs), ("up", up_sf, up_gs), ("down", down_sf, down_gs)]:
|
|
sf_f32 = sf.float()
|
|
gs_f32 = gs.float()
|
|
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()
|
|
|
|
n_total = product.numel()
|
|
n_clamped = (product > 448.0).sum().item()
|
|
n_zeroed = (folded_back == 0.0).sum().item()
|
|
n_nonzero_orig = (sf_f32 > 0).sum().item()
|
|
n_nonzero_folded = (folded_back > 0).sum().item()
|
|
|
|
rel_err = (folded_back - product).abs() / product.clamp(min=1e-10)
|
|
|
|
print(f"\n {name}_proj:")
|
|
print(f" block_sf: shape={list(sf.shape)} range=[{sf_f32.min():.4e}, {sf_f32.max():.4e}] unique_u8={torch.unique(sf.view(torch.uint8)).numel()}")
|
|
print(f" global_sf: {gs_f32.item():.6e}")
|
|
print(f" product (sf*gs): range=[{product.min():.4e}, {product.max():.4e}]")
|
|
print(f" folded (float8): range=[{folded_back.min():.4e}, {folded_back.max():.4e}]")
|
|
print(f" Clamped to 448: {n_clamped}/{n_total} ({100*n_clamped/n_total:.1f}%)")
|
|
print(f" Became zero: {n_zeroed}/{n_total} ({100*n_zeroed/n_total:.1f}%)")
|
|
print(f" Was nonzero → became zero: {n_nonzero_orig - n_nonzero_folded}/{n_nonzero_orig}")
|
|
print(f" Rel error: max={rel_err.max():.4f} mean={rel_err.mean():.4f}")
|
|
|
|
# Show the float8 step size at the product magnitude
|
|
if product.max() > 0:
|
|
typical = product.median().item()
|
|
if typical > 0:
|
|
f8_typ = torch.tensor(typical, dtype=torch.float32).to(torch.float8_e4m3fn)
|
|
f8_back = f8_typ.float()
|
|
if f8_back > 0:
|
|
step = (f8_typ.view(torch.uint8) + 1).view(torch.float8_e4m3fn).float() - f8_back
|
|
print(f" Float8 step at median ({typical:.4e}): Δ={step.item():.4e} rel={step.item()/f8_back.item():.2%}")
|
|
|
|
break
|
|
|
|
# Now check: what if we DON'T fold, and instead pass global_scale as GEMM alpha?
|
|
print("\n" + "=" * 60)
|
|
print("ALTERNATIVE: Pass global_scale as GEMM alpha")
|
|
print("=" * 60)
|
|
print("""
|
|
The fold is lossy because float8 can't represent the product range.
|
|
But if we DON'T fold, the CUTLASS GEMM needs a separate global scale mechanism.
|
|
|
|
Option 1: Multiply the GEMM alpha by the weight's global_scale
|
|
- alpha already carries the activation global scale
|
|
- We could fold weight global scale into alpha: alpha_new = alpha * weight_gs
|
|
- BUT: alpha is a single scalar, weight_gs varies per-expert
|
|
- For grouped GEMM, each expert needs its own alpha
|
|
|
|
Option 2: Keep block scales as-is (no fold), multiply output by global_scale
|
|
- After GEMM: output *= weight_global_scale
|
|
- This is exact (float32 multiply on bf16 output)
|
|
- Requires passing global_scale to nvfp4_mega_moe_full
|
|
|
|
Option 3: Fold global_scale into the GEMM alpha per-expert
|
|
- In cutlass_grouped_nvfp4_gemm, each expert gets its own alpha
|
|
- alpha_expert = l1_global_scale * l1_weight_global_scale[expert_id]
|
|
- This is EXACT and doesn't lose precision
|
|
- The block scales stay at their original float8 values (no folding)
|
|
""")
|