From 83c47d07d164226d95c62caac76cefc77846e8c8 Mon Sep 17 00:00:00 2001
From: Wenwen Qu <vinny_qu@163.com>
Date: Thu, 19 Oct 2023 10:59:54 +0800
Subject: [PATCH] fix is_no_pp_or_last_stage logic

---
 internlm/core/context/parallel_context.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/internlm/core/context/parallel_context.py b/internlm/core/context/parallel_context.py
index 997bd46..915905a 100644
--- a/internlm/core/context/parallel_context.py
+++ b/internlm/core/context/parallel_context.py
@@ -329,7 +329,8 @@ class ParallelContext(metaclass=SingletonMeta):
         return self.is_last_rank(ParallelMode.PIPELINE)
 
     def is_no_pp_or_last_stage(self):
-        return not self.is_initialized(ParallelMode.PIPELINE) or self.is_pipeline_last_stage()
+        # NOTICE!!!, this will ignore virutal stage
+        return not self.is_initialized(ParallelMode.PIPELINE) or self.is_last_rank(ParallelMode.PIPELINE)
 
     def get_world_size(self, parallel_mode: ParallelMode):
         """Returns the world size for `parallel_mode`.