Simplify router gate: dequant NVFP4 → BF16, F.linear (no FP8 middleman)
Same as what worked before. The checkpoint stores NVFP4 weights, so we dequantize once at load time and use cuBLAS F.linear. No FP8 re-quantize step needed — that was just adding noise on top of the NVFP4 dequant.
This commit is contained in:
@@ -1306,36 +1306,28 @@ def main():
|
||||
router.load_weights(hash_lut=all_w[f"{pfx}.gate.tid2eid"].to(dev, torch.int32))
|
||||
else:
|
||||
eb = all_w.get(f"{pfx}.gate.e_score_correction_bias")
|
||||
# FP8_E4M3 router gate — quantize weight to FP8, dequantize to BF16, F.linear
|
||||
# This avoids NVFP4's multi-scale complexity while still using FP8 compression.
|
||||
# BF16 router gate — dequantize NVFP4 to BF16, use F.linear
|
||||
E = cfg["n_routed_experts"]
|
||||
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, pfx, 'gate')
|
||||
if gate_w is not None and gate_ws is not None:
|
||||
# Checkpoint has NVFP4 gate weight — dequantize to BF16 first, then re-quantize to FP8
|
||||
# Checkpoint has NVFP4 gate weight — dequantize to BF16
|
||||
from dsv4.ops.quantize import dequantize_nvfp4
|
||||
ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0
|
||||
isc_v = gate_isc.float().item() if gate_isc is not None else 1.0/(6.0*448.0)
|
||||
gsb = 1.0 * ws2_v # global_scale_b = gs * ws2
|
||||
gsa = torch.tensor([gsb] * gate_w.shape[0], device=dev, dtype=torch.float32)
|
||||
gate_bf16 = dequantize_nvfp4(gate_w.to(dev), gate_ws.to(dev), gsa) # (E_packed*2, H)
|
||||
gate_bf16 = gate_bf16.T.contiguous() # (H, E) for W_gate
|
||||
router.W_gate = gate_bf16.T.contiguous() # (H, E) for F.linear(x, W_gate.T)
|
||||
else:
|
||||
# BF16 gate weight from checkpoint
|
||||
gw = all_w.get(f"{pfx}.gate.weight")
|
||||
gate_bf16 = gw.bfloat16().to(dev)
|
||||
if gate_bf16.shape[0] != H:
|
||||
gate_bf16 = gate_bf16.T.contiguous() # ensure (H, E)
|
||||
# Quantize to FP8_E4M3: scale = amax / 448.0
|
||||
gate_amax = gate_bf16.abs().max().float().item()
|
||||
gate_scale = gate_amax / 448.0
|
||||
gate_fp8 = (gate_bf16.float() / gate_scale).to(torch.float8_e4m3fn)
|
||||
# Dequantize back to BF16 for F.linear (FP8 round-trip ~0.9999 cos)
|
||||
gate_dequant = gate_fp8.to(torch.bfloat16) * gate_scale
|
||||
router.W_gate = gate_dequant.contiguous() # (H, E) for F.linear(x, W_gate.T)
|
||||
router.W_gate = gate_bf16.contiguous()
|
||||
# No gate_lin — force BF16 dispatch path
|
||||
router.gate_lin = None
|
||||
router.load_weights(e_bias=eb.to(dev, torch.float32))
|
||||
if li < 5: print(f" L{li}: FP8_E4M3 router gate (scale={gate_scale:.6f}, amax={gate_amax:.4f})", flush=True)
|
||||
if li < 5: print(f" L{li}: BF16 router gate (dequantized from NVFP4)", flush=True)
|
||||
router.finalize_weights(); routers[li] = router
|
||||
|
||||
moe = Nvfp4MoE(num_experts=cfg["n_routed_experts"], hidden_size=H,
|
||||
|
||||
Reference in New Issue
Block a user