Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -1234,6 +1234,16 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
         logits = logits.float()
+
+            # Handle unused parameters
+    if self.training:
+        for expert in self.model.block_sparse_moe.experts:
+            for param in expert.parameters():
+                if param.requires_grad and param.grad is None:
+                    param.grad = torch.zeros_like(param)
+
+    loss = None
+    if labels is not None:
         
         loss = None
         if labels is not None: