github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix is_no_pp_or_last_stage logic

pull/423/head
Wenwen Qu 2023-10-19 10:59:54 +08:00
parent 6480e03949
commit 83c47d07d1
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
internlm/core/context/parallel_context.py
View File

@ -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`.