Throw exception if ring_exchange is not available when pipeline_model_parallel_size > 1

megatron/arguments.py

@@ -59,6 +59,10 @@ def parse_args(extra_args_provider=None, defaults={},
     args.pipeline_model_parallel_size = min(
         args.pipeline_model_parallel_size,
         (args.world_size // args.tensor_model_parallel_size))
+    if args.pipeline_model_parallel_size > 1:
+        if "ring_exchange" not in dir(torch.distributed):
+            raise Exception('PyTorch with torch.distributed.ring_exchange needed '
+                            'to run pipeline MP!')
     if args.num_microbatches_in_minibatch is None:
         args.num_microbatches_in_minibatch = 1
     if args.rank == 0: