add option to disable pipeline partitioning

added pipeline configuration to deepspeed config under pipeline section,
with defaults of
pipe_partitioned = True
grad_partitioned = True
can be used to enable or disable activations and grads partitioning in TP/PP mode.
When not using TP, activation and grads partitioning will be disabled
regardless of this confgiuration.

deepspeed/runtime/config.py

@@ -446,6 +446,8 @@ def get_pipeline_config(param_dict):
         "partition": "best",
         "seed_layers": False,
         "activation_checkpoint_interval": 0,
+        "pipe_partitioned": True,
+        "grad_partitioned": True,
     }
     config = default_pipeline
     for key, val in param_dict.get("pipeline", {}).items():

deepspeed/runtime/pipe/engine.py

@@ -129,8 +129,12 @@ def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs):
 
         # Partition input/output buffers
         # XXX temporarily disable while I revert some partition hacks.
-        self.is_pipe_partitioned = self.is_model_parallel
-        self.is_grad_partitioned = self.is_model_parallel
+        assert isinstance(self._config.pipeline['pipe_partitioned'], bool)
+        assert isinstance(self._config.pipeline['grad_partitioned'], bool)
+        self.is_pipe_partitioned = self.is_model_parallel and self._config.pipeline['pipe_partitioned']
+        self.is_grad_partitioned = self.is_model_parallel and self._config.pipeline['grad_partitioned']
+        logger.info(f'is_pipe_partitioned= {self.is_pipe_partitioned}',
+                    f'is_grad_partitioned= {self.is_grad_partitioned}')
 
         model_parameters = filter(lambda p: p.requires_grad, self.module.parameters())
         num_params = sum([p.numel() for p in model_parameters])