Fix changed behavior of pipe_parallel

EleutherAI · May 17, 2024 · 7d4093b · 7d4093b
1 parent 49cd41f
commit 7d4093b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/neox_arguments/arguments.py b/megatron/neox_arguments/arguments.py
@@ -1035,7 +1035,7 @@ def calculate_derived(self):
         # the sequential model without the PipelineModule wrapper to avoid the overhead it incurs
         self.update_value(
             "is_pipe_parallel",
-            self.pipe_parallel_size > 1 and self.moe_num_experts == 1,
+            self.pipe_parallel_size >= 1 and self.moe_num_experts == 1,
         )
         if self.moe_num_experts > 1:
             assert not (