From 7d4093be86f9c343f45be866f9ab7eece3b1eddc Mon Sep 17 00:00:00 2001 From: Yang Zhang Date: Thu, 16 May 2024 19:16:14 +0000 Subject: [PATCH] Fix changed behavior of pipe_parallel --- megatron/neox_arguments/arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/neox_arguments/arguments.py b/megatron/neox_arguments/arguments.py index ff4f4bc21..d9a586eb4 100644 --- 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 (