Fix warmup: pass local expert IDs (not global), remove incorrect _warmup_done guard
compute_activation_global_scales expects local IDs (0..num_experts-1), not global IDs. EP5/EP7 were getting L2 gs=0 because global IDs (240+, 336+) didn't match expert_id_range (0..47), so no tokens matched any expert → L1 GEMM got zero inputs → L2 gs=0 → NaN/crash. Also removed _warmup_done guard since each layer needs its own warmup (different weights, different gs values).
This commit is contained in:
@@ -551,19 +551,14 @@ class DeepseekV4MegaMoEExperts(nn.Module):
|
||||
# the actual amax and compute correct gs values.
|
||||
self._warmup_activation_global_scales()
|
||||
|
||||
_warmup_done: bool = False
|
||||
|
||||
def _warmup_activation_global_scales(self) -> None:
|
||||
"""Run a warmup forward pass to compute correct activation global scales.
|
||||
|
||||
Called once during finalize_weights, before cudagraph capture.
|
||||
Called once per layer during finalize_weights, before cudagraph capture.
|
||||
Uses quantize_to_nvfp4 (which calls .max()) to get the exact gs
|
||||
from real activation magnitudes, then stores them for use by
|
||||
quantize_activation_nvfp4 (no .max(), cudagraph-safe).
|
||||
"""
|
||||
if self._warmup_done:
|
||||
return
|
||||
self._warmup_done = True
|
||||
import torch
|
||||
runner = self._cutedsl_runner
|
||||
device = runner.device
|
||||
@@ -574,11 +569,12 @@ class DeepseekV4MegaMoEExperts(nn.Module):
|
||||
# Sample hidden states: typical BF16 activations have amax ~1-10
|
||||
hidden_states = torch.randn(num_tokens, runner.hidden_size,
|
||||
dtype=torch.bfloat16, device=device)
|
||||
# Assign all tokens to the first few experts evenly
|
||||
# Assign all tokens to local experts (0..num_local_experts-1)
|
||||
# compute_activation_global_scales expects local IDs, not global
|
||||
topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int64, device=device)
|
||||
for i in range(num_tokens):
|
||||
for j in range(top_k):
|
||||
topk_ids[i, j] = (runner.experts_start_idx + j) % (runner.experts_start_idx + runner.num_experts)
|
||||
topk_ids[i, j] = j % runner.num_experts
|
||||
topk_weights = torch.ones(num_tokens, top_k, dtype=torch.float32, device=device) / top_k
|
||||
|
||||
runner.compute_activation_global_scales(
|
||||
|
||||
Reference in New Issue
Block a user