fix: use checkpoint input_scale for activation global scale (not hardcoded 1/2688)

The checkpoint stores input_scale per projection — the pre-computed
activation normalization factor. Using 1/2688 was wrong for most layers
(e.g. down_proj input_scale=0.031 vs 1/2688=0.000372 — 83x off).
This caused under-quantized activations and garbage output.
This commit is contained in:
2026-05-16 21:46:00 +00:00
parent af087e655e
commit 152648789d
2 changed files with 11 additions and 2 deletions

View File

@@ -60,8 +60,8 @@ class CuTeDSLMoERunner:
self._l1_gsb = None
self._l2_gsb = None
self._l1_activation_global_scale = 1.0 / 2688.0
self._l2_activation_global_scale = 1.0 / 2688.0
self._l1_activation_global_scale = None # set from checkpoint input_scale
self._l2_activation_global_scale = None
# Pre-allocated cudagraph buffers (set in _allocate_buffers)
self._token_indices = None

View File

@@ -512,6 +512,15 @@ class DeepseekV4MegaMoEExperts(nn.Module):
self._cutedsl_runner.l2_sf = l2_sf
self._cutedsl_runner.l2_gs = l2_gs
# Set activation global scales from checkpoint input_scale
# The input_scale is the pre-computed activation normalization factor
# for each projection. Using hardcoded 1/2688 gives wrong scale for most layers.
# input_scale shape: (num_experts, 2) for w13 (gate, up), (num_experts, 1) for w2
l1_igs = self._w13_input_scale[:, 0] if self._w13_input_scale is not None else None # gate input_scale
l2_igs = self._w2_input_scale[:, 0] if self._w2_input_scale is not None else None # down input_scale
self._cutedsl_runner.l1_activation_global_scale = l1_igs.mean().item() if l1_igs is not None else 1.0 / 2688.0
self._cutedsl_runner.l2_activation_global_scale = l2_igs.mean().item() if l2_igs is not None else 1.0 / 2688.0
# Drop the original loader-side parameters
self._w13_input_scale = self.w13_input_scale.data.clone()
self._w2_input_scale = self.w2_input_scale.data.clone()