[Verification] Add pipeline and 3D parallelism (DP/TP/PP) support

ci/task_unit_test.sh

@@ -39,6 +39,9 @@ echo "Running unit tests..."
 #   -p "no:randomly": disable randomly plugin for sharding tests.
 torchrun --nproc_per_node 2 -r 1:1 -m pytest -rxXs -p "no:randomly" tests
 
+echo "Running DeepSpeed unit tests..."
+deepspeed --num_gpus 4 tests/test_ds_pipeline.py
+
 echo "Downloading test data..."
 bash benchmark/download_benchmark_dataset.sh
 

conftest.py

@@ -14,6 +14,13 @@ def pytest_collection_modifyitems(items):
     are running different tests, then the entire tests will stuck.
     """
     items.sort(key=lambda item: item.name)
+    new_items = []
+    for item in items:
+        # Skip DeepSpeed tests
+        if item.parent.name not in ["test_ds_pipeline.py"]:
+            new_items.append(item)
+    items[:] = new_items
+    return items
 
 
 @pytest.fixture(scope="session")

slapo/build.py

@@ -249,15 +249,26 @@ def build(
     if sch.metadata.primitives["cut_pipeline_stage"]:
         # pipeline stages will be wrapped into PipeStageWrapper
         sch = generate_pipeline_partition(sch)
-        # Re-analyzie tie weights before consolidation.
-        sch.metadata.tie_weights = analyze_tie_weights(
-            sch.mod, is_pipeline_partitioned=True
-        )
+        is_pipeline = True
+    else:
+        is_pipeline = False
+    # Re-analyzie tie weights before consolidation.
+    sch.metadata.tie_weights = analyze_tie_weights(
+        sch.mod, is_pipeline_partitioned=is_pipeline
+    )
 
     # delay initialization
     if init_weights:
         init_weight_fn = init_weights if isinstance(init_weights, Callable) else None
         sch = consolidate_model(sch, target, init_weight_fn, **kwargs)
+        tie_weights = list(sch.metadata.tie_weights.values())
+        if tie_weights and not is_pipeline:
+            if not hasattr(sch.mod, "tie_weights"):
+                raise RuntimeError(
+                    "Model needs to tie weights but does not have `tie_weights` method. Probably because the model has been traced."
+                )
+            # https://github.com/huggingface/transformers/blob/v4.28.1/src/transformers/modeling_utils.py#L1274-L1277
+            sch.mod.tie_weights()
 
     if sch.metadata.primitives["cut_pipeline_stage"] and target is not None:
         # Generate pipeline modules for a particular target.

slapo/framework_dialect/deepspeed/pipeline.py

@@ -26,6 +26,73 @@ class WrappedTypeCode(Enum):
     TUPLE = 4
 
 
+def get_ds_config(
+    batch_size,
+    micro_batch_size_per_gpu,
+    fp16=True,
+    zero_stage=0,
+    desc="",
+    bf16=False,
+    sequence_parallel=False,
+):
+    # https://github.com/microsoft/DeepSpeed/blob/ff42743/tests/unit/model_parallelism/test_configurable_parallel_pp.py#L20
+    logger.info(f"fp16={fp16}, bf16={bf16}")
+    config_dict = {
+        "help": desc,
+        "steps_per_print": 10,
+        "optimizer": {"type": "AdamW", "params": {"lr": 0.0001}},
+        "fp16": {"enabled": fp16, "initial_scale_power": 12},
+        "bf16": {"enabled": bf16},
+        "gradient_clipping": 1.0,
+        "train_batch_size": batch_size,
+        "train_micro_batch_size_per_gpu": micro_batch_size_per_gpu,
+        "pipeline": {
+            "sequence_parallel": sequence_parallel,
+        },
+        "wall_clock_breakdown": False,
+    }
+
+    if zero_stage > 0:
+        zero_config_dict = {
+            "zero_optimization": {
+                "stage": zero_stage,
+                "overlap_comm": True,
+                "reduce_scatter": True,
+                "contiguous_gradients": False,
+                "prefetch_bucket_size": 5e8,
+            },
+            "zero_allow_untested_optimizer": True,
+        }
+        config_dict.update(zero_config_dict)
+
+    return config_dict
+
+
+_groups = []
+
+
+def create_dist_group_for_pipeline(num_pp, num_mp):
+    from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
+
+    world_size = dist.get_world_size()
+    num_dp = world_size // (num_pp * num_mp)
+    topology = PipeModelDataParallelTopology(
+        num_pp=num_pp, num_mp=num_mp, num_dp=num_dp
+    )
+    model_groups = topology.get_axis_comm_lists("model")
+
+    global_rank = dist.get_rank()
+    group = None
+
+    for g in model_groups:
+        proc_group = dist.new_group(ranks=g)
+        _groups.append(proc_group)
+        if global_rank in g:
+            group = proc_group
+
+    return topology, group
+
+
 def get_simple_nested_list_str(data):
     """A helper function that prints a nested structure without printing
     tensor values.
@@ -88,7 +155,7 @@ def flat_and_name_tensor_list(data, name, suffix):
     values = data.values() if isinstance(data, dict) else data
 
     if isinstance(values, (list, tuple)) and any(
-        isinstance(t, torch.Tensor) for t in values
+        isinstance(t, (torch.Tensor, torch.Size)) for t in values
     ):
         for idx, tensor in enumerate(values):
             name_n_value.extend(
@@ -297,7 +364,7 @@ def forward(self, *args, **kwargs):
                 ), f"[{self.name}] Arg {arg_name} not found in liveness list: {liveness}"
                 idx = liveness.index(arg_name)
             ordered_args.append(unordered_args[idx])
-            if i > 0:
+            if i > 0 and not self.last:
                 # https://github.com/microsoft/DeepSpeed/blob/v0.9.2/deepspeed/runtime/pipe/engine.py#L639
                 assert (
                     torch.is_tensor(unordered_args[idx])