working with interleaving

megatron/arguments.py

@@ -147,7 +147,7 @@ def parse_args(extra_args_provider=None, defaults={},
         #     args.num_layers_per_virtual_pipeline_stage
         # <<<
         args.virtual_pipeline_model_parallel_size = \
-            (args.num_layers // transformer_pipeline_size) // \
+            (args.num_layers // args.transformer_pipeline_model_parallel_size) // \
             args.num_layers_per_virtual_pipeline_stage
         # >>>
         # from lutil import pax

megatron/mpu/initialize.py

@@ -350,13 +350,13 @@ def get_num_layers(args, is_encoder_and_decoder_model):
             #     get_pipeline_model_parallel_world_size()
             # )
             # <<<
-            assert args.num_layers % transformer_pipeline_size == 0, \
-                'num_layers must be divisible by transformer_pipeline_size'
+            assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \
+                'num_layers must be divisible by transformer_pipeline_model_parallel_size'
             num_layers = (
                 0
                 if args.standalone_embed_stage
                 and get_pipeline_model_parallel_rank() == 0 else
-                args.num_layers // transformer_pipeline_size
+                args.num_layers // args.transformer_pipeline_model_parallel_size
             )
     else:
         num_layers = args.num_layers

megatron/schedules.py

@@ -40,9 +40,26 @@ def get_forward_backward_func():
             #     "pipeline size" : args.pipeline_model_parallel_size,
             # })
             # <<<
-            assert get_num_microbatches() % args.pipeline_model_parallel_size == 0, \
-                'number of microbatches is not divisible by pipeline-parallel ' \
-                'size when using interleaved schedule'
+            # >>>
+            # assert get_num_microbatches() % args.pipeline_model_parallel_size == 0, \
+            #     'number of microbatches is not divisible by pipeline-parallel ' \
+            #     'size when using interleaved schedule'
+            # assert get_num_microbatches() % \
+            #     args.transformer_pipeline_model_parallel_size == 0, \
+            #     'number of microbatches (%d) is not divisible by transformer-' \
+            #     'pipeline-model-parallel-size (%d) when using interleaved ' \
+            #     'schedule' % (
+            #         get_num_microbatches(),
+            #         args.transformer_pipeline_model_parallel_size,
+            #     )
+            assert get_num_microbatches() % \
+                args.pipeline_model_parallel_size == 0, \
+                'number of microbatches (%d) is not divisible by pipeline-' \
+                'model-parallel-size (%d) when using interleaved schedule' % (
+                    get_num_microbatches(),
+                    args.pipeline_model_parallel_size,
+                )
+            # <<<
         else:
             forward_backward_func = forward_backward_pipelining_without_interleaving
     else: