Add interleaved pipeline schedule to GPT

examples/nlp/language_modeling/conf/megatron_gpt_config.yaml

@@ -18,6 +18,7 @@ trainer:
   accumulate_grad_batches: 1 # do not modify, grad acc is automatic for training megatron models
   gradient_clip_val: 1.0
   benchmark: False
+  enable_model_summary: False # default PTL callback for this does not support model parallelism, instead we log manually
 
 exp_manager:
   explicit_log_dir: null
@@ -47,7 +48,7 @@ model:
   global_batch_size: 8 # will use more micro batches to reach global batch size
   tensor_model_parallel_size: 1 # intra-layer model parallelism
   pipeline_model_parallel_size: 1 # inter-layer model parallelism
-  resume_from_checkpoint: null # manually set the checkpoint file to load from
+  virtual_pipeline_model_parallel_size: null # interleaved pipeline
 
   # model architecture
   encoder_seq_length: 512
@@ -92,6 +93,7 @@ model:
 
   # miscellaneous
   seed: 1234
+  resume_from_checkpoint: null # manually set the checkpoint file to load from
   use_cpu_initialization: False # Init weights on the CPU (slow for large models)
   onnx_safe: False # Use work-arounds for known problems with Torch ONNX exporter.
   apex_transformer_log_level: 30 # Python logging level displays logs with severity greater than or equal to this

examples/nlp/language_modeling/megatron_gpt_eval.py

@@ -171,6 +171,7 @@ def main(cfg) -> None:
                 app_state.model_parallel_size,
                 app_state.data_parallel_size,
                 app_state.pipeline_model_parallel_split_rank,
+                app_state.virtual_pipeline_model_parallel_rank,
             ) = fake_initialize_model_parallel(
                 world_size=app_state.model_parallel_size,
                 rank=trainer.global_rank,

examples/nlp/language_modeling/megatron_t5_eval.py

@@ -70,6 +70,7 @@ def main():
             app_state.model_parallel_size,
             app_state.data_parallel_size,
             app_state.pipeline_model_parallel_split_rank,
+            app_state.virtual_pipeline_model_parallel_rank,
         ) = fake_initialize_model_parallel(
             world_size=app_state.model_parallel_size,
             rank=trainer.global_rank,

examples/nlp/language_modeling/megatron_t5_prompt_learning_eval.py

@@ -56,6 +56,7 @@ def main(cfg) -> None:
             app_state.model_parallel_size,
             app_state.data_parallel_size,
             app_state.pipeline_model_parallel_split_rank,
+            app_state.virtual_pipeline_model_parallel_rank,
         ) = fake_initialize_model_parallel(
             world_size=app_state.model_parallel_size,
             rank=trainer.global_rank,

examples/nlp/language_modeling/tuning/megatron_t5_adapter_eval.py

@@ -57,6 +57,7 @@ def main(cfg) -> None:
             app_state.model_parallel_size,
             app_state.data_parallel_size,
             app_state.pipeline_model_parallel_split_rank,
+            app_state.virtual_pipeline_model_parallel_rank,
         ) = fake_initialize_model_parallel(
             world_size=app_state.model_parallel_size,
             rank=trainer.global_rank,

examples/nlp/language_modeling/tuning/megatron_t5_ia3_eval.py

@@ -57,6 +57,7 @@ def main(cfg) -> None:
             app_state.model_parallel_size,
             app_state.data_parallel_size,
             app_state.pipeline_model_parallel_split_rank,
+            app_state.virtual_pipeline_model_parallel_rank,
         ) = fake_initialize_model_parallel(
             world_size=app_state.model_parallel_size,
             rank=trainer.global_rank,

examples/nlp/machine_translation/nmt_transformer_infer_megatron.py

@@ -62,6 +62,7 @@ def main(cfg) -> None:
         app_state.model_parallel_size,
         app_state.data_parallel_size,
         app_state.pipeline_model_parallel_split_rank,
+        app_state.virtual_pipeline_model_parallel_rank,
     ) = fake_initialize_model_parallel(
         world_size=app_state.model_parallel_size,
         rank=trainer.global_rank,

nemo/collections/nlp/models/language_modeling/megatron_base_model.py

@@ -33,6 +33,7 @@
 from nemo.collections.nlp.parts.nlp_overrides import GradScaler
 from nemo.core.optim import MainParamsOptimizerWrapper, prepare_lr_scheduler
 from nemo.utils import AppState, logging
+from nemo.utils.get_rank import is_global_rank_zero
 
 try:
     from apex.transformer import parallel_state
@@ -87,6 +88,7 @@ def __init__(self, cfg: DictConfig, trainer: Trainer, no_lm_init=True):
             local_rank=trainer.local_rank,
             tensor_model_parallel_size=cfg.get('tensor_model_parallel_size', 1),
             pipeline_model_parallel_size=cfg.get('pipeline_model_parallel_size', 1),
+            virtual_pipeline_model_parallel_size=cfg.get('virtual_pipeline_model_parallel_size', None),
             pipeline_model_parallel_split_rank=cfg.get('pipeline_model_parallel_split_rank', 0),
             micro_batch_size=cfg.get('micro_batch_size'),
             global_batch_size=cfg.get('global_batch_size'),
@@ -389,3 +391,17 @@ def _validate_config(self):
             logging.info("Gradient accumulation fusion can only be used with megatron amp O2 mixed precision.")
             with open_dict(self.cfg):
                 self.cfg.gradient_accumulation_fusion = False
+
+    def is_data_parallel_rank_zero(self):
+        if is_global_rank_zero():
+            return True
+        else:
+            try:
+                data_parallel_rank = parallel_state.get_data_parallel_rank()
+            except:
+                data_parallel_rank = None
+
+            if data_parallel_rank is not None and data_parallel_rank == 0:
+                return True
+            else:
+                return False