Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -65,6 +65,10 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "GemmoeConfig"
 
+class GemmoeDistributedDataParallel(nn.parallel.DistributedDataParallel):
+    def __init__(self, model, **kwargs):
+        super().__init__(model, find_unused_parameters=True, **kwargs)
+
 def approx_gelu(x):
     return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x**3)))
 
@@ -80,6 +84,7 @@ def _get_unpad_data(attention_mask):
     )
 
 
+
 class GemmoeRMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
@@ -1112,6 +1117,28 @@ 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):
@@ -1126,6 +1153,13 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         # Initialize weights and apply final processing
         self.post_init()
 
+    def parallelize(self, device_map=None):
+        self.model = GemmoeDistributedDataParallel(
+            self.model,
+            device_ids=[torch.cuda.current_device()],
+            output_device=torch.cuda.current_device(),
+    )
+
     def get_input_embeddings(self):
         return self.model.embed_tokens
 
@@ -1215,11 +1249,10 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
 
         # Handle unused parameters
         if self.training:
-            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)
+            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)
 
         loss = None
         if labels is not None:
@@ -1328,15 +1361,6 @@ 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).