Bake runtime fixes into model code: SwiGLU clamp, router stability, from_pretrained auto-fix

fuse2_model.py

@@ -66,6 +66,11 @@ class SwiGLUExpert(nn.Module):
     down_proj: (hidden, intermediate)
     """
 
+    # DeepSeek V4 Flash uses swiglu_limit=10.0 to clamp intermediate
+    # activations. Without this, outlier values grow exponentially across
+    # 36 layers and produce NaN by layer 6.
+    SWIGLU_LIMIT = 10.0
+
     def __init__(self, hidden_size: int, intermediate_size: int):
         super().__init__()
         self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
@@ -73,7 +78,9 @@ class SwiGLUExpert(nn.Module):
         self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+        gate_up = F.silu(self.gate_proj(x)) * self.up_proj(x)
+        gate_up = gate_up.clamp(-self.SWIGLU_LIMIT, self.SWIGLU_LIMIT)
+        return self.down_proj(gate_up)
 
 
 class Fuse2Router(nn.Module):
@@ -115,8 +122,11 @@ class Fuse2Router(nn.Module):
             router_logits: (batch*seq, num_experts) — raw logits for load balancing
         """
         # sqrtsoftplus scoring (from DeepSeek V4)
+        # Clamp softplus to min=1e-6 before sqrt to prevent NaN gradients
+        # when logits are very negative (softplus → 0 → sqrt(0) = 0, but
+        # gradient sqrt'(0) = inf).
         logits = self.gate(hidden_states)  # (tokens, num_experts)
-        scores = F.softplus(logits).sqrt()
+        scores = F.softplus(logits).clamp(min=1e-6).sqrt()
 
         # Top-k selection
         topk_weights, topk_indices = scores.topk(self.top_k, dim=-1)
@@ -514,6 +524,69 @@ class Fuse2ForCausalLM(Qwen3ForCausalLM):
             **kwargs,
         )
 
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """Load from HuggingFace Hub with automatic runtime fixes.
+
+        This overrides the default from_pretrained to apply three critical
+        fixes after weight loading:
+
+        1. Initialize coding_gate and coding_norm if they're still on meta
+           device (these params are not in the safetensors checkpoint).
+        2. Cast all RMSNorm/LayerNorm weights from float32 to bfloat16 to
+           enable fused SDPA kernel dispatch (otherwise falls back to slow
+           Python loops).
+        3. Ensure coding_enabled is True (config default).
+
+        With these fixes, from_pretrained produces a working model without
+        any manual post-load patching.
+        """
+        model = super().from_pretrained(*args, **kwargs)
+        model._apply_runtime_fixes()
+        return model
+
+    def _apply_runtime_fixes(self):
+        """Apply runtime fixes after weight loading.
+
+        Called automatically by from_pretrained. Can also be called manually
+        if the model was loaded via a custom path (e.g., init_empty_weights +
+        manual safetensors loading).
+        """
+        device = next(self.parameters()).device
+        fixed_meta = 0
+        fixed_norms = 0
+
+        for layer in self.model.layers:
+            if not isinstance(layer, Fuse2AugmentedLayer):
+                continue
+
+            # Fix 1: coding_gate on meta device → initialize to -2.0
+            if hasattr(layer, 'coding_gate'):
+                if layer.coding_gate.device.type == 'meta':
+                    layer.coding_gate = nn.Parameter(
+                        torch.tensor(-2.0, device=device))
+                    fixed_meta += 1
+
+            # Fix 2: coding_norm on meta device → create fresh RMSNorm
+            if hasattr(layer, 'coding_norm'):
+                if hasattr(layer.coding_norm, 'weight') and \
+                   layer.coding_norm.weight.device.type == 'meta':
+                    layer.coding_norm = nn.RMSNorm(
+                        layer.coding_norm.weight.shape[0], eps=1e-6).to(device)
+                    fixed_meta += 1
+
+        # Fix 3: Cast float32 norm weights to bfloat16 for fused kernels
+        for module in self.modules():
+            if hasattr(module, 'weight') and hasattr(module, 'eps'):
+                if module.weight.dtype == torch.float32:
+                    module.weight.data = module.weight.data.to(torch.bfloat16)
+                    fixed_norms += 1
+
+        # Ensure coding is enabled
+        self.set_coding_enabled(True)
+
+        return {"meta_params_fixed": fixed_meta, "norms_cast_to_bf16": fixed_norms}
+
 
 def load_expert_weights(
     model: Fuse2ForCausalLM,