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:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user