Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -1112,28 +1112,6 @@ class GemmoeModel(GemmoePreTrainedModel):
 		return causal_mask
 
 class GemmoeForCausalLM(GemmoePreTrainedModel):
-    r"""
-    The Gemmoe Model transformer with a language modeling head on top for causal language modeling (CLM).
-
-    Args:
-        config (GemmoeConfig): The configuration object for the Gemmoe model.
-
-    Example usage:
-    ```python
-    >>> from transformers import AutoTokenizer, GemmoeForCausalLM
-
-    >>> model = GemmoeForCausalLM.from_pretrained("google/gemmoe-7b")
-    >>> tokenizer = AutoTokenizer.from_pretrained("google/gemmoe-7b")
-
-    >>> prompt = "What is your favorite condiment?"
-    >>> inputs = tokenizer(prompt, return_tensors="pt")
-
-    >>> # Generate
-    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    "What is your favorite condiment?"
-    ```
-    """
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
@@ -1237,10 +1215,11 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
 
         # Handle unused parameters
         if self.training:
-            for expert in self.model.layers[-1].block_sparse_moe.experts:
-                for param in expert.parameters():
-                    if param.requires_grad and param.grad is None:
-                        param.grad = torch.zeros_like(param)
+            for layer in self.model.layers:
+                for expert in layer.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:
@@ -1349,6 +1328,15 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
                 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
             )
         return reordered_past
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
 @add_start_docstrings(
     """
     The Gemmoe Model transformer with a sequence classification head on top (linear layer).