Add swiglu_limit=10.0 activation clamping (was missing)
DeepSeek-V4 uses SiluAndMulWithClamp(10.0) which clamps: - silu(gate) to max 10.0 - up to [-10.0, 10.0] Our runner was doing plain F.silu(gate) * up without clamping. Large gate values could produce unbounded SiLU output, causing numerical issues in the L2 GEMM. This is likely contributing to garbage model output.
This commit is contained in:
@@ -46,6 +46,7 @@ class CuTeDSLMoERunner:
|
||||
self.top_k = top_k
|
||||
self.device = device
|
||||
self.experts_start_idx = experts_start_idx
|
||||
self._swiglu_limit = None # Set via set_swiglu_limit()
|
||||
|
||||
# Weight storage (set before _ensure_stacked)
|
||||
self.l1_fp4 = None
|
||||
@@ -92,6 +93,10 @@ class CuTeDSLMoERunner:
|
||||
f"pid={os.getpid()}", flush=True)
|
||||
self._buffers_allocated = False
|
||||
|
||||
def set_swiglu_limit(self, limit: float | None):
|
||||
"""Set the swiglu_limit for activation clamping."""
|
||||
self._swiglu_limit = limit
|
||||
|
||||
def _fill_token_indices(self):
|
||||
"""Fill _token_indices with [0,0,..0, 1,1,..1, ...] (each token repeated top_k times).
|
||||
|
||||
@@ -438,10 +443,15 @@ class CuTeDSLMoERunner:
|
||||
# Extract real token outputs from padded GEMM output
|
||||
l1_out_real = l1_out[padded_dst]
|
||||
|
||||
# === SiLU(gate) * up ===
|
||||
# === SiLU(gate) * up (with swiglu_limit clamp) ===
|
||||
gate = l1_out_real[:, :self.intermediate_size]
|
||||
up = l1_out_real[:, self.intermediate_size:]
|
||||
activated = torch.nn.functional.silu(gate) * up
|
||||
gate_silu = torch.nn.functional.silu(gate)
|
||||
# Apply DeepSeek-V4 swiglu_limit: clamp both silu(gate) and up
|
||||
if self._swiglu_limit is not None:
|
||||
gate_silu = gate_silu.clamp(max=self._swiglu_limit)
|
||||
up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
|
||||
activated = gate_silu * up
|
||||
|
||||
# === L2: down ===
|
||||
padded_activated = self._shared_bufs['activated']
|
||||
|
||||
@@ -551,6 +551,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
|
||||
# the actual amax and compute correct gs values.
|
||||
self._warmup_activation_global_scales()
|
||||
|
||||
# Set swiglu_limit for activation clamping in the runner
|
||||
if self.swiglu_limit is not None:
|
||||
self._cutedsl_runner.set_swiglu_limit(float(self.swiglu_limit))
|
||||
|
||||
def _warmup_activation_global_scales(self) -> None:
|
||||
"""Run a warmup forward pass to compute correct activation global scales.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user