Fix wrongly all_gather for mixtral finetune (huggingface#965)

Co-authored-by: ccrhx4 <[email protected]>
skavulya · May 8, 2024 · 3a14236 · 3a14236
1 parent aa175ea
commit 3a14236
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py
@@ -367,7 +367,7 @@ def gaudi_mixtral_block_sparse_moe_forward(self, hidden_states: torch.Tensor) ->
     # router_logits: (batch * sequence_length, n_experts)
     router_logits = self.gate(hidden_states)
 
-    if is_deepspeed_available():
+    if is_deepspeed_available() and (not self.training):
         from deepspeed import comm as dist
 
         if dist.is_initialized():