add setter of pipeline model parallel split rank (#1306)

apex/transformer/parallel_state.py

@@ -348,6 +348,12 @@ def get_pipeline_model_parallel_split_rank():
     return _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
 
 
+def set_pipeline_model_parallel_split_rank(pipeline_model_parallel_split_rank: int):
+    """Set my rank for the pipeline model parallel split rank."""
+    global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
+    _PIPELINE_MODEL_PARALLEL_SPLIT_RANK = pipeline_model_parallel_split_rank
+
+
 def is_pipeline_first_stage(ignore_virtual=False):
     """Return True if in the first pipeline model-parallel stage, False otherwise."""
     if not ignore_virtual:

tests/L0/run_transformer/run_initialize_test.py

@@ -100,6 +100,11 @@ def test_pipeline_model_parallel_split_rank():
     split_rank = parallel_state.get_pipeline_model_parallel_split_rank()
     assert split_rank is pipeline_model_parallel_split_rank_
 
+    fake_split_rank = 7
+    parallel_state.set_pipeline_model_parallel_split_rank(fake_split_rank)
+    split_rank = parallel_state.get_pipeline_model_parallel_split_rank()
+    assert split_rank == fake_split_rank
+
     # Reset groups
     parallel_state.destroy_model_parallel()