From 16c911cffc504ba700c9562ab9fa8fd7315194e2 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Fri, 9 Jun 2023 20:48:54 +0000
Subject: [PATCH 01/28] Move ds_config and dist_group to framework_dialect

---
 slapo/framework_dialect/deepspeed/pipeline.py | 67 +++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/slapo/framework_dialect/deepspeed/pipeline.py b/slapo/framework_dialect/deepspeed/pipeline.py
index c220821c..31cb0d20 100644
--- a/slapo/framework_dialect/deepspeed/pipeline.py
+++ b/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.

From cd8d85a68907f79dbdac76708337b85c8e7b7b99 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 05:53:27 +0000
Subject: [PATCH 02/28] Add test_ds_pipeline

---
 tests/test_ds_pipeline.py | 92 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)
 create mode 100644 tests/test_ds_pipeline.py

diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
new file mode 100644
index 00000000..19a1bf8c
--- /dev/null
+++ b/tests/test_ds_pipeline.py
@@ -0,0 +1,92 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Test DeepSpeed Pipeline."""
+import pytest
+
+import os
+import torch
+from torch import nn
+import deepspeed
+
+import slapo
+from slapo.framework_dialect.deepspeed.pipeline import (
+    get_ds_config,
+    create_dist_group_for_pipeline,
+)
+from slapo.random import set_random_seed
+import torch.distributed as dist
+import torch.nn.functional as F
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(12)])
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def test_pipeline():
+    deepspeed.init_distributed(dist_backend="nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    with slapo.init_empty_weights():
+        model = Model()
+    topology, group = create_dist_group_for_pipeline(num_pp=2, num_mp=1)
+    sch = slapo.create_schedule(model, group=group)
+    orig_sch = slapo.create_schedule(model)
+    sch.trace_until("")
+    sch["layers.5"].cut_pipeline_stage()
+    bs = 8
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=bs,
+        fp16=False,
+    )
+    original_model, _ = slapo.build(orig_sch)
+    original_model.to(dist.get_rank())
+    set_random_seed(2023)
+    inp = torch.randn(bs, 10, device=dist.get_rank())
+    label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
+    if dist.get_world_size() > 1:
+        dist.broadcast(inp, src=0)
+        dist.broadcast(label, src=0)
+    original_output = original_model(inp)
+    original_output = F.cross_entropy(original_output, label)
+    print("original output: ", original_output)
+    original_state_dict = original_model.state_dict()
+
+    def init_weights(mod, path):
+        for name, _ in mod.named_parameters(recurse=False):
+            old_name = ".".join(path.split(".")[1:]).replace("_", ".") + "." + name
+            setattr(
+                mod,
+                name,
+                nn.Parameter(
+                    original_state_dict[old_name].detach().to(dist.get_rank())
+                ),
+            )
+
+    model, _ = slapo.build(
+        sch,
+        init_weights=init_weights,
+        topology=topology,
+        target="deepspeed",
+        config=ds_config_dict,
+        loss_fn=nn.CrossEntropyLoss(),
+    )
+    model.to(dist.get_rank())
+    train_iter = iter([(inp, label) for _ in range(100)])
+    output = model.train_batch(data_iter=train_iter)
+    print("new ouput", output)
+    if dist.get_rank() == 1:
+        assert torch.allclose(original_output, output)
+
+
+if __name__ == "__main__":
+    # pytest.main([__file__])
+    test_pipeline()

From 767b30a0da1b26a9a8aa1fddff8c62227f385433 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 06:31:40 +0000
Subject: [PATCH 03/28] Update verify API

---
 slapo/verify.py           | 100 +++++++++++++++++++++++++++++++-------
 tests/test_ds_pipeline.py |  42 +++-------------
 2 files changed, 90 insertions(+), 52 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 591273b2..93bee0c7 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -19,16 +19,33 @@
 
 
 class Verify(ContextDecorator):
-    def __init__(self, sch, example_inputs, device="cuda", eval_mode=True, enable=True):
+    def __init__(
+        self,
+        sch,
+        example_inputs,
+        example_outputs=None,
+        loss_fn=None,
+        device="cuda",
+        topology=None,
+        eval_mode=True,
+        enable=True,
+        **kwargs,
+    ):
         if not isinstance(example_inputs, list):
             example_inputs = [example_inputs]
-        self.example_inputs = example_inputs
+        self.device = device
+        self.example_inputs = [x.to(self.device) for x in example_inputs]
+        self.example_outputs = (
+            example_outputs.to(self.device) if example_outputs else None
+        )
+        self.loss_fn = loss_fn
         self.original_trace = None
         self.sch = sch
         self.original_sch = create_schedule(copy.deepcopy(self.sch.mod))
-        self.device = device
+        self.topology = topology
         self.enable = enable
         self.eval_mode = eval_mode
+        self.kwargs = kwargs
 
     def __enter__(self):
         self.original_trace = sys.gettrace()
@@ -59,6 +76,25 @@ def __exit__(self, *exc):
         """
         if not self.enable:
             return
+        if self.sch.metadata.primitives["cut_pipeline_stage"]:
+            try:
+                import deepspeed
+            except ImportError:
+                raise ImportError(
+                    "deepspeed is required when pipeline parallelism is used"
+                )
+            assert (
+                self.example_outputs is not None
+            ), "example_outputs must be provided when pipeline parallelism is used"
+            assert (
+                self.loss_fn is not None
+            ), "loss_fn must be provided when pipeline parallelism is used"
+            assert (
+                self.topology is not None
+            ), "topology must be provided when pipeline parallelism is used"
+            assert (
+                "config" in self.kwargs
+            ), "config must be provided when pipeline parallelism is used"
         # 1. Build the original model with random weights
         named_params = self.original_sch.mod.named_parameters()
         is_initialized = named_params.__next__()[1].device != torch.device("meta")
@@ -66,24 +102,31 @@ def __exit__(self, *exc):
         #    make sure all the buffers are on the right device
         original_mod = original_mod.to(self.device)
         # 2. Get the example inputs
-        self.example_inputs = [x.to(self.device) for x in self.example_inputs]
-        #   Broadcast the example inputs from rank 0 to other ranks
-        if self.sch.world_size > 1:
-            for inp in self.example_inputs:
-                dist.broadcast(inp, src=0, group=self.sch.group)
+        #    Broadcast the example inputs from rank 0 to other ranks
+        group_src_rank = (
+            dist.get_global_rank(self.sch.group, 0) if self.sch.group is not None else 0
+        )
+        for inp in self.example_inputs:
+            dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
         # 3. Run the original model
         #    make sure the random seeds are the same, which may affect the output of dropout
         if self.eval_mode:
             original_mod.eval()
         set_random_seed(2023)
         original_output = original_mod(*self.example_inputs)
+        if self.example_outputs is not None:
+            assert (
+                self.loss_fn is not None
+            ), "loss_fn must be provided when example_outputs is provided"
+            original_output = self.loss_fn(original_output, self.example_outputs)
         # 4. Broadcast the original model from rank 0 to other ranks
         original_state_dict = original_mod.state_dict()
-        if self.sch.world_size > 1:
-            for param_name in original_state_dict:
-                dist.broadcast(
-                    original_state_dict[param_name], src=0, group=self.sch.group
-                )
+        for param_name in original_state_dict:
+            dist.broadcast(
+                original_state_dict[param_name],
+                src=group_src_rank,
+                group=self.sch.group,
+            )
         # 5. Delete the original model to avoid excessive memory usage
         del original_mod
         # 6. Get the transformed model from the schedule
@@ -94,22 +137,41 @@ def __exit__(self, *exc):
         for param_name, param in self.sch.mod.named_parameters():
             if hasattr(param, "orig_shape"):
                 copied_mod.get_parameter(param_name).orig_shape = param.orig_shape
-        new_sch = create_schedule(copied_mod)
+        new_sch = create_schedule(copied_mod, group=self.sch.group)
+        # copy schedule metadata
+        new_sch.metadata = copy.deepcopy(self.sch.metadata)
         # 7. Use original weights to initialize the new model
         #    Notice init_weights is called before actual sharding, so we only need to
         #    assign the original weights to the corresponding modules
 
         def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
+                # TODO: fix submod name
+                if self.sch.metadata.primitives["cut_pipeline_stage"]:
+                    original_name = (
+                        ".".join(path.split(".")[1:]).replace("_", ".") + "." + name
+                    )
+                else:
+                    original_name = f"{path}.{name}"
                 setattr(
                     mod,
                     name,
                     nn.Parameter(
-                        original_state_dict[f"{path}.{name}"].detach().to(self.device)
+                        original_state_dict[original_name].detach().to(self.device)
                     ),
                 )
 
-        new_mod, _ = build(new_sch, init_weights=init_weights)
+        if self.sch.metadata.primitives["cut_pipeline_stage"]:
+            new_mod, _ = build(
+                new_sch,
+                init_weights=init_weights,
+                target="deepspeed",
+                topology=self.topology,
+                config=self.kwargs["config"],
+                loss_fn=self.loss_fn,
+            )
+        else:
+            new_mod, _ = build(new_sch, init_weights=init_weights)
         # 8. Run the new model
         #    make sure all the buffers are on the right device
         new_mod.to(self.device)
@@ -117,7 +179,11 @@ def init_weights(mod, path):
             new_mod.eval()
         #    make sure the random seeds are the same, which may affect the output of dropout
         set_random_seed(2023)
-        new_output = new_mod(*self.example_inputs)
+        if self.sch.metadata.primitives["cut_pipeline_stage"]:
+            train_iter = iter([tuple(self.example_inputs + [self.example_outputs])])
+            new_output = new_mod.train_batch(train_iter)
+        else:
+            new_output = new_mod(*self.example_inputs)
         # 9. Compare the outputs
         torch.testing.assert_close(original_output, new_output)
         logger.info("Passed verification!")
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 19a1bf8c..3f2cc635 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -38,53 +38,25 @@ def test_pipeline():
         model = Model()
     topology, group = create_dist_group_for_pipeline(num_pp=2, num_mp=1)
     sch = slapo.create_schedule(model, group=group)
-    orig_sch = slapo.create_schedule(model)
     sch.trace_until("")
-    sch["layers.5"].cut_pipeline_stage()
+
     bs = 8
     ds_config_dict = get_ds_config(
         batch_size=bs,
         micro_batch_size_per_gpu=bs,
         fp16=False,
     )
-    original_model, _ = slapo.build(orig_sch)
-    original_model.to(dist.get_rank())
-    set_random_seed(2023)
     inp = torch.randn(bs, 10, device=dist.get_rank())
     label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
-    if dist.get_world_size() > 1:
-        dist.broadcast(inp, src=0)
-        dist.broadcast(label, src=0)
-    original_output = original_model(inp)
-    original_output = F.cross_entropy(original_output, label)
-    print("original output: ", original_output)
-    original_state_dict = original_model.state_dict()
-
-    def init_weights(mod, path):
-        for name, _ in mod.named_parameters(recurse=False):
-            old_name = ".".join(path.split(".")[1:]).replace("_", ".") + "." + name
-            setattr(
-                mod,
-                name,
-                nn.Parameter(
-                    original_state_dict[old_name].detach().to(dist.get_rank())
-                ),
-            )
-
-    model, _ = slapo.build(
+    with slapo.Verify(
         sch,
-        init_weights=init_weights,
+        example_inputs=[inp],
+        example_outputs=label,
+        loss_fn=F.cross_entropy,
         topology=topology,
-        target="deepspeed",
         config=ds_config_dict,
-        loss_fn=nn.CrossEntropyLoss(),
-    )
-    model.to(dist.get_rank())
-    train_iter = iter([(inp, label) for _ in range(100)])
-    output = model.train_batch(data_iter=train_iter)
-    print("new ouput", output)
-    if dist.get_rank() == 1:
-        assert torch.allclose(original_output, output)
+    ):
+        sch["layers.5"].cut_pipeline_stage()
 
 
 if __name__ == "__main__":

From cfb34c81b37ec6e74ec51c988ead980819eb1e7d Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 15:39:19 +0000
Subject: [PATCH 04/28] Fix name

---
 slapo/verify.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 93bee0c7..53a6d4a2 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -81,7 +81,7 @@ def __exit__(self, *exc):
                 import deepspeed
             except ImportError:
                 raise ImportError(
-                    "deepspeed is required when pipeline parallelism is used"
+                    "DeepSpeed is required when pipeline parallelism is used"
                 )
             assert (
                 self.example_outputs is not None
@@ -146,13 +146,14 @@ def __exit__(self, *exc):
 
         def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
-                # TODO: fix submod name
                 if self.sch.metadata.primitives["cut_pipeline_stage"]:
-                    original_name = (
-                        ".".join(path.split(".")[1:]).replace("_", ".") + "." + name
-                    )
-                else:
-                    original_name = f"{path}.{name}"
+                    path = path.split(".")
+                    for idx, subpath in enumerate(path):
+                        if "submod_" not in subpath:
+                            break
+                    # Fix ModuleList name
+                    path = ".".join(path[idx:]).replace("_", ".")
+                original_name = f"{path}.{name}"
                 setattr(
                     mod,
                     name,

From 7365f897ed202cf25b0acae71e0bae7b734a1817 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 16:12:21 +0000
Subject: [PATCH 05/28] Verify correctness by setting random seed

---
 slapo/verify.py           | 14 ++++++++++++--
 tests/test_ds_pipeline.py |  1 +
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 53a6d4a2..41d88317 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -36,7 +36,9 @@ def __init__(
         self.device = device
         self.example_inputs = [x.to(self.device) for x in example_inputs]
         self.example_outputs = (
-            example_outputs.to(self.device) if example_outputs else None
+            example_outputs.to(self.device)
+            if isinstance(example_outputs, torch.Tensor)
+            else None
         )
         self.loss_fn = loss_fn
         self.original_trace = None
@@ -73,6 +75,10 @@ def trace_calls(frame, event, arg):
     def __exit__(self, *exc):
         """Verify the correctness of the schedule.
         TODO: Support backward verification
+
+        TP: Takes the same model and same inputs, and expects the same outputs across devices.
+        DP: Takes the same model and different inputs, and expects the same outputs on the same device.
+        PP: Only need to verify the correctness in the last stage.
         """
         if not self.enable:
             return
@@ -101,13 +107,17 @@ def __exit__(self, *exc):
         original_mod, _ = build(self.original_sch, init_weights=not is_initialized)
         #    make sure all the buffers are on the right device
         original_mod = original_mod.to(self.device)
-        # 2. Get the example inputs
+        # 2. Get the example inputs and outputs
         #    Broadcast the example inputs from rank 0 to other ranks
         group_src_rank = (
             dist.get_global_rank(self.sch.group, 0) if self.sch.group is not None else 0
         )
         for inp in self.example_inputs:
             dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
+        if self.example_outputs is not None:
+            dist.broadcast(
+                self.example_outputs, src=group_src_rank, group=self.sch.group
+            )
         # 3. Run the original model
         #    make sure the random seeds are the same, which may affect the output of dropout
         if self.eval_mode:
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 3f2cc635..122f3d2e 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -46,6 +46,7 @@ def test_pipeline():
         micro_batch_size_per_gpu=bs,
         fp16=False,
     )
+    set_random_seed(2013)
     inp = torch.randn(bs, 10, device=dist.get_rank())
     label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
     with slapo.Verify(

From aac338ba6105732690de24dd1e0c8ee221872ae6 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 17:43:24 +0000
Subject: [PATCH 06/28] Correct PP

---
 slapo/verify.py           | 50 +++++++++++++++++++++++++--------------
 tests/test_ds_pipeline.py |  2 --
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 41d88317..e5fbe0d3 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -101,6 +101,9 @@ def __exit__(self, *exc):
             assert (
                 "config" in self.kwargs
             ), "config must be provided when pipeline parallelism is used"
+            is_pipeline = True
+        else:
+            is_pipeline = False
         # 1. Build the original model with random weights
         named_params = self.original_sch.mod.named_parameters()
         is_initialized = named_params.__next__()[1].device != torch.device("meta")
@@ -108,16 +111,29 @@ def __exit__(self, *exc):
         #    make sure all the buffers are on the right device
         original_mod = original_mod.to(self.device)
         # 2. Get the example inputs and outputs
-        #    Broadcast the example inputs from rank 0 to other ranks
-        group_src_rank = (
-            dist.get_global_rank(self.sch.group, 0) if self.sch.group is not None else 0
-        )
-        for inp in self.example_inputs:
-            dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
-        if self.example_outputs is not None:
-            dist.broadcast(
-                self.example_outputs, src=group_src_rank, group=self.sch.group
-            )
+        #    Broadcast the example inputs from rank 0 in each TP/PP group
+        #    to other ranks in the same group.
+        #    Only the first stage of PP needs the example inputs & outputs,
+        #    but for verification, each device holds an entire copy of the original model,
+        #    so we need to broadcast the example inputs & outputs to all the TP&PP devices.
+        #    Notice for each device in the DP group, they should take different inputs.
+        if is_pipeline:
+            for i in range(self.topology.get_dim("data")):
+                tp_pp_group = dist.new_group(ranks=self.topology.filter_match(data=i))
+                group_src_rank = dist.get_global_rank(tp_pp_group, 0)
+                for inp in self.example_inputs:
+                    dist.broadcast(inp, src=group_src_rank, group=tp_pp_group)
+                dist.broadcast(
+                    self.example_outputs, src=group_src_rank, group=tp_pp_group
+                )
+        else:
+            group_src_rank = 0
+            for inp in self.example_inputs:
+                dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
+            if self.example_outputs is not None:
+                dist.broadcast(
+                    self.example_outputs, src=group_src_rank, group=self.sch.group
+                )
         # 3. Run the original model
         #    make sure the random seeds are the same, which may affect the output of dropout
         if self.eval_mode:
@@ -130,13 +146,11 @@ def __exit__(self, *exc):
             ), "loss_fn must be provided when example_outputs is provided"
             original_output = self.loss_fn(original_output, self.example_outputs)
         # 4. Broadcast the original model from rank 0 to other ranks
+        #    Since for verification, each device holds an entire copy of the original
+        #    model, here we directly broadcast the model to all the devices.
         original_state_dict = original_mod.state_dict()
         for param_name in original_state_dict:
-            dist.broadcast(
-                original_state_dict[param_name],
-                src=group_src_rank,
-                group=self.sch.group,
-            )
+            dist.broadcast(original_state_dict[param_name], src=0)
         # 5. Delete the original model to avoid excessive memory usage
         del original_mod
         # 6. Get the transformed model from the schedule
@@ -156,7 +170,7 @@ def __exit__(self, *exc):
 
         def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
-                if self.sch.metadata.primitives["cut_pipeline_stage"]:
+                if is_pipeline:
                     path = path.split(".")
                     for idx, subpath in enumerate(path):
                         if "submod_" not in subpath:
@@ -172,7 +186,7 @@ def init_weights(mod, path):
                     ),
                 )
 
-        if self.sch.metadata.primitives["cut_pipeline_stage"]:
+        if is_pipeline:
             new_mod, _ = build(
                 new_sch,
                 init_weights=init_weights,
@@ -190,7 +204,7 @@ def init_weights(mod, path):
             new_mod.eval()
         #    make sure the random seeds are the same, which may affect the output of dropout
         set_random_seed(2023)
-        if self.sch.metadata.primitives["cut_pipeline_stage"]:
+        if is_pipeline:
             train_iter = iter([tuple(self.example_inputs + [self.example_outputs])])
             new_output = new_mod.train_batch(train_iter)
         else:
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 122f3d2e..822bfeeb 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -14,7 +14,6 @@
     get_ds_config,
     create_dist_group_for_pipeline,
 )
-from slapo.random import set_random_seed
 import torch.distributed as dist
 import torch.nn.functional as F
 
@@ -46,7 +45,6 @@ def test_pipeline():
         micro_batch_size_per_gpu=bs,
         fp16=False,
     )
-    set_random_seed(2013)
     inp = torch.randn(bs, 10, device=dist.get_rank())
     label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
     with slapo.Verify(

From 6b3aa7c0a6620351174a5b6e41a3b0796bf0481d Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 17:47:35 +0000
Subject: [PATCH 07/28] Add 4stages

---
 tests/test_ds_pipeline.py | 41 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 822bfeeb..414a0bb9 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -19,9 +19,9 @@
 
 
 class Model(nn.Module):
-    def __init__(self):
+    def __init__(self, num_layers=12):
         super().__init__()
-        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(12)])
+        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(num_layers)])
 
     def forward(self, x):
         for layer in self.layers:
@@ -29,12 +29,12 @@ def forward(self, x):
         return x
 
 
-def test_pipeline():
+def test_pipeline_2stages():
     deepspeed.init_distributed(dist_backend="nccl")
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     with slapo.init_empty_weights():
-        model = Model()
+        model = Model(12)
     topology, group = create_dist_group_for_pipeline(num_pp=2, num_mp=1)
     sch = slapo.create_schedule(model, group=group)
     sch.trace_until("")
@@ -58,6 +58,37 @@ def test_pipeline():
         sch["layers.5"].cut_pipeline_stage()
 
 
+def test_pipeline_4stages():
+    deepspeed.init_distributed(dist_backend="nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    with slapo.init_empty_weights():
+        model = Model(12)
+    topology, group = create_dist_group_for_pipeline(num_pp=4, num_mp=1)
+    sch = slapo.create_schedule(model, group=group)
+    sch.trace_until("")
+
+    bs = 8
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=bs,
+        fp16=False,
+    )
+    inp = torch.randn(bs, 10, device=dist.get_rank())
+    label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
+    with slapo.Verify(
+        sch,
+        example_inputs=[inp],
+        example_outputs=label,
+        loss_fn=F.cross_entropy,
+        topology=topology,
+        config=ds_config_dict,
+    ):
+        sch["layers.2"].cut_pipeline_stage()
+        sch["layers.5"].cut_pipeline_stage()
+        sch["layers.8"].cut_pipeline_stage()
+
+
 if __name__ == "__main__":
     # pytest.main([__file__])
-    test_pipeline()
+    test_pipeline_4stages()

From 4aedbb3d1878f25a1725c96f109eee090eddb708 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 17:54:57 +0000
Subject: [PATCH 08/28] Add DP

---
 tests/test_ds_pipeline.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 414a0bb9..325d052d 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -29,20 +29,24 @@ def forward(self, x):
         return x
 
 
-def test_pipeline_2stages():
+def test_pipeline_2stages_pp_dp():
     deepspeed.init_distributed(dist_backend="nccl")
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     with slapo.init_empty_weights():
         model = Model(12)
-    topology, group = create_dist_group_for_pipeline(num_pp=2, num_mp=1)
+    # If total number of devices is 4, then num_dp = 2
+    num_pp = 2
+    num_mp = 1
+    num_dp = dist.get_world_size() // (num_pp * num_mp)
+    topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
     sch = slapo.create_schedule(model, group=group)
     sch.trace_until("")
 
     bs = 8
     ds_config_dict = get_ds_config(
         batch_size=bs,
-        micro_batch_size_per_gpu=bs,
+        micro_batch_size_per_gpu=bs // num_dp,
         fp16=False,
     )
     inp = torch.randn(bs, 10, device=dist.get_rank())
@@ -58,7 +62,7 @@ def test_pipeline_2stages():
         sch["layers.5"].cut_pipeline_stage()
 
 
-def test_pipeline_4stages():
+def test_pipeline_4stages_pp():
     deepspeed.init_distributed(dist_backend="nccl")
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
@@ -91,4 +95,4 @@ def test_pipeline_4stages():
 
 if __name__ == "__main__":
     # pytest.main([__file__])
-    test_pipeline_4stages()
+    test_pipeline_2stages_pp_dp()

From be504b1333fac8c6992adc62da026e3db383bc56 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 18:18:01 +0000
Subject: [PATCH 09/28] Add TP+PP

---
 slapo/verify.py           | 15 +++++++--
 tests/test_ds_pipeline.py | 66 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index e5fbe0d3..e0c886a6 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -155,7 +155,17 @@ def __exit__(self, *exc):
         del original_mod
         # 6. Get the transformed model from the schedule
         #    Copy it and build a new schedule to prevent the original schedule from being modified
-        copied_mod = copy.deepcopy(self.sch.mod)
+        try:
+            copied_mod = copy.deepcopy(self.sch.mod)
+            is_copy_failed = False
+        except TypeError:
+            # One example is ProcessGroup cannot be copied:
+            # https://github.com/pytorch/pytorch/issues/73825
+            is_copy_failed = True
+            logger.warning(
+                "Failed to copy the model, using the original model to verify"
+            )
+            copied_mod = self.sch.mod
         # copy original attributes
         # TODO: find a better way to copy attributes
         for param_name, param in self.sch.mod.named_parameters():
@@ -212,5 +222,6 @@ def init_weights(mod, path):
         # 9. Compare the outputs
         torch.testing.assert_close(original_output, new_output)
         logger.info("Passed verification!")
-        del new_mod
+        if not is_copy_failed:
+            del new_mod
         sys.settrace(self.original_trace)
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 325d052d..4e77aaa6 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -31,6 +31,8 @@ def forward(self, x):
 
 def test_pipeline_2stages_pp_dp():
     deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 4:
+        pytest.skip("This test requires 4 GPUs.")
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     with slapo.init_empty_weights():
@@ -62,8 +64,70 @@ def test_pipeline_2stages_pp_dp():
         sch["layers.5"].cut_pipeline_stage()
 
 
+def test_pipeline_2stages_pp_tp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 4:
+        pytest.skip("This test requires 4 GPUs.")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    class LinearReLU(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = nn.Linear(10, 10)
+
+        def forward(self, x):
+            x = self.linear(x)
+            x = F.relu(x)
+            return x
+
+    class Model(nn.Module):
+        def __init__(self, num_layers=2):
+            super().__init__()
+            self.layers = nn.ModuleList([LinearReLU() for _ in range(num_layers)])
+
+        def forward(self, x):
+            for layer in self.layers:
+                x = layer(x)
+            return x
+
+    with slapo.init_empty_weights():
+        model = Model(2)
+    num_pp = 2
+    num_mp = 2
+    num_dp = dist.get_world_size() // (num_pp * num_mp)
+    topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
+    sch = slapo.create_schedule(model, group=group)
+    sch.trace_until("")
+
+    bs = 8
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=bs // num_dp,
+        fp16=False,
+    )
+    inp = torch.randn(bs, 10, device=dist.get_rank())
+    label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
+    with slapo.Verify(
+        sch,
+        example_inputs=[inp],
+        example_outputs=label,
+        loss_fn=F.cross_entropy,
+        topology=topology,
+        config=ds_config_dict,
+    ):
+        sch["layers.0.linear"].shard("weight", axis=0)
+        sch["layers.0.linear"].shard("bias", axis=0)
+        sch["layers.1.linear"].shard("weight", axis=1)
+        sch["layers.1.linear"].sync("fwd_post", sync_op_or_fn="all_reduce")
+        sch["layers.0.linear"].sync("bwd_post", sync_op_or_fn="all_reduce")
+        sch["layers.0"].cut_pipeline_stage()
+
+
 def test_pipeline_4stages_pp():
     deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 4:
+        pytest.skip("This test requires 4 GPUs.")
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     with slapo.init_empty_weights():
@@ -95,4 +159,4 @@ def test_pipeline_4stages_pp():
 
 if __name__ == "__main__":
     # pytest.main([__file__])
-    test_pipeline_2stages_pp_dp()
+    test_pipeline_2stages_pp_tp()

From 95948ae2e4bb2198dcf76edabf07362c27475583 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 19:03:23 +0000
Subject: [PATCH 10/28] Fix DP

---
 slapo/verify.py           | 8 ++++++++
 tests/test_ds_pipeline.py | 3 ++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index e0c886a6..d8462c60 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -145,6 +145,13 @@ def __exit__(self, *exc):
                 self.loss_fn is not None
             ), "loss_fn must be provided when example_outputs is provided"
             original_output = self.loss_fn(original_output, self.example_outputs)
+        if is_pipeline:
+            # average the loss across all the DP devices
+            for ranks in self.topology.get_axis_comm_lists("data"):
+                if dist.get_rank() in ranks:
+                    dp_group = dist.new_group(ranks=ranks)
+                    dist.all_reduce(original_output, group=dp_group)
+                    original_output /= len(ranks)
         # 4. Broadcast the original model from rank 0 to other ranks
         #    Since for verification, each device holds an entire copy of the original
         #    model, here we directly broadcast the model to all the devices.
@@ -216,6 +223,7 @@ def init_weights(mod, path):
         set_random_seed(2023)
         if is_pipeline:
             train_iter = iter([tuple(self.example_inputs + [self.example_outputs])])
+            # DeepSpeed will automatically broadcast the output to each device
             new_output = new_mod.train_batch(train_iter)
         else:
             new_output = new_mod(*self.example_inputs)
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 4e77aaa6..25383ad9 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -116,6 +116,7 @@ def forward(self, x):
         topology=topology,
         config=ds_config_dict,
     ):
+        # TODO: Fix layers.0/1 sharding bug
         sch["layers.0.linear"].shard("weight", axis=0)
         sch["layers.0.linear"].shard("bias", axis=0)
         sch["layers.1.linear"].shard("weight", axis=1)
@@ -159,4 +160,4 @@ def test_pipeline_4stages_pp():
 
 if __name__ == "__main__":
     # pytest.main([__file__])
-    test_pipeline_2stages_pp_tp()
+    test_pipeline_2stages_pp_dp()

From 53887fab47bc2159190ce9ccf22bc0665d3ff0ca Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 19:34:30 +0000
Subject: [PATCH 11/28] Add DS testing

---
 ci/task_unit_test.sh      |  3 ++
 conftest.py               |  7 ++++
 tests/test_ds_pipeline.py | 85 ++++++++++++++++++++++++++++-----------
 3 files changed, 71 insertions(+), 24 deletions(-)

diff --git a/ci/task_unit_test.sh b/ci/task_unit_test.sh
index 6de80663..2bf4ef6c 100644
--- a/ci/task_unit_test.sh
+++ b/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
 
diff --git a/conftest.py b/conftest.py
index d0dcbfa7..22108dc6 100644
--- a/conftest.py
+++ b/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")
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 25383ad9..e010b91e 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -18,10 +18,24 @@
 import torch.nn.functional as F
 
 
+class LinearReLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 10)
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = F.relu(x)
+        return x
+
+
 class Model(nn.Module):
-    def __init__(self, num_layers=12):
+    def __init__(self, num_layers=12, has_relu=False):
         super().__init__()
-        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(num_layers)])
+        if not has_relu:
+            self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(num_layers)])
+        else:
+            self.layers = nn.ModuleList([LinearReLU() for _ in range(num_layers)])
 
     def forward(self, x):
         for layer in self.layers:
@@ -71,28 +85,8 @@ def test_pipeline_2stages_pp_tp():
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
 
-    class LinearReLU(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.linear = nn.Linear(10, 10)
-
-        def forward(self, x):
-            x = self.linear(x)
-            x = F.relu(x)
-            return x
-
-    class Model(nn.Module):
-        def __init__(self, num_layers=2):
-            super().__init__()
-            self.layers = nn.ModuleList([LinearReLU() for _ in range(num_layers)])
-
-        def forward(self, x):
-            for layer in self.layers:
-                x = layer(x)
-            return x
-
     with slapo.init_empty_weights():
-        model = Model(2)
+        model = Model(2, has_relu=True)
     num_pp = 2
     num_mp = 2
     num_dp = dist.get_world_size() // (num_pp * num_mp)
@@ -158,6 +152,49 @@ def test_pipeline_4stages_pp():
         sch["layers.8"].cut_pipeline_stage()
 
 
+# def test_pipeline_2stages_pp_tp_dp():
+#     deepspeed.init_distributed(dist_backend="nccl")
+#     if dist.get_world_size() != 8:
+#         pytest.skip("This test requires 8 GPUs.")
+#     local_rank = int(os.environ["LOCAL_RANK"])
+#     torch.cuda.set_device(local_rank)
+
+#     with slapo.init_empty_weights():
+#         model = Model(2, has_relu=True)
+#     num_pp = 2
+#     num_mp = 2
+#     num_dp = dist.get_world_size() // (num_pp * num_mp)
+#     print("num_dp:", num_dp, "num_pp:", num_pp, "num_mp:", num_mp)
+#     topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
+#     sch = slapo.create_schedule(model, group=group)
+#     sch.trace_until("")
+
+#     bs = 8
+#     ds_config_dict = get_ds_config(
+#         batch_size=bs,
+#         micro_batch_size_per_gpu=bs // num_dp,
+#         fp16=False,
+#     )
+#     inp = torch.randn(bs, 10, device=dist.get_rank())
+#     label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
+#     with slapo.Verify(
+#         sch,
+#         example_inputs=[inp],
+#         example_outputs=label,
+#         loss_fn=F.cross_entropy,
+#         topology=topology,
+#         config=ds_config_dict,
+#     ):
+#         # TODO: Fix layers.0/1 sharding bug
+#         sch["layers.0.linear"].shard("weight", axis=0)
+#         sch["layers.0.linear"].shard("bias", axis=0)
+#         sch["layers.1.linear"].shard("weight", axis=1)
+#         sch["layers.1.linear"].sync("fwd_post", sync_op_or_fn="all_reduce")
+#         sch["layers.0.linear"].sync("bwd_post", sync_op_or_fn="all_reduce")
+#         sch["layers.0"].cut_pipeline_stage()
+
+
 if __name__ == "__main__":
-    # pytest.main([__file__])
     test_pipeline_2stages_pp_dp()
+    test_pipeline_2stages_pp_tp()
+    test_pipeline_4stages_pp()

From 322f0301aa1cdef9ec20b5d434bf80029aec475b Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 19:48:04 +0000
Subject: [PATCH 12/28] Fix 3D parallelism

---
 slapo/verify.py           |  4 +-
 tests/test_ds_pipeline.py | 83 ++++++++++++++++++++-------------------
 2 files changed, 44 insertions(+), 43 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index d8462c60..dfc3debf 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -148,9 +148,9 @@ def __exit__(self, *exc):
         if is_pipeline:
             # average the loss across all the DP devices
             for ranks in self.topology.get_axis_comm_lists("data"):
+                dp_group = dist.new_group(ranks=ranks)
+                dist.all_reduce(original_output, group=dp_group)
                 if dist.get_rank() in ranks:
-                    dp_group = dist.new_group(ranks=ranks)
-                    dist.all_reduce(original_output, group=dp_group)
                     original_output /= len(ranks)
         # 4. Broadcast the original model from rank 0 to other ranks
         #    Since for verification, each device holds an entire copy of the original
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index e010b91e..0b17a288 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -110,7 +110,7 @@ def test_pipeline_2stages_pp_tp():
         topology=topology,
         config=ds_config_dict,
     ):
-        # TODO: Fix layers.0/1 sharding bug
+        # TODO: Fix layers.0/1 indexing bug
         sch["layers.0.linear"].shard("weight", axis=0)
         sch["layers.0.linear"].shard("bias", axis=0)
         sch["layers.1.linear"].shard("weight", axis=1)
@@ -152,49 +152,50 @@ def test_pipeline_4stages_pp():
         sch["layers.8"].cut_pipeline_stage()
 
 
-# def test_pipeline_2stages_pp_tp_dp():
-#     deepspeed.init_distributed(dist_backend="nccl")
-#     if dist.get_world_size() != 8:
-#         pytest.skip("This test requires 8 GPUs.")
-#     local_rank = int(os.environ["LOCAL_RANK"])
-#     torch.cuda.set_device(local_rank)
-
-#     with slapo.init_empty_weights():
-#         model = Model(2, has_relu=True)
-#     num_pp = 2
-#     num_mp = 2
-#     num_dp = dist.get_world_size() // (num_pp * num_mp)
-#     print("num_dp:", num_dp, "num_pp:", num_pp, "num_mp:", num_mp)
-#     topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
-#     sch = slapo.create_schedule(model, group=group)
-#     sch.trace_until("")
-
-#     bs = 8
-#     ds_config_dict = get_ds_config(
-#         batch_size=bs,
-#         micro_batch_size_per_gpu=bs // num_dp,
-#         fp16=False,
-#     )
-#     inp = torch.randn(bs, 10, device=dist.get_rank())
-#     label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
-#     with slapo.Verify(
-#         sch,
-#         example_inputs=[inp],
-#         example_outputs=label,
-#         loss_fn=F.cross_entropy,
-#         topology=topology,
-#         config=ds_config_dict,
-#     ):
-#         # TODO: Fix layers.0/1 sharding bug
-#         sch["layers.0.linear"].shard("weight", axis=0)
-#         sch["layers.0.linear"].shard("bias", axis=0)
-#         sch["layers.1.linear"].shard("weight", axis=1)
-#         sch["layers.1.linear"].sync("fwd_post", sync_op_or_fn="all_reduce")
-#         sch["layers.0.linear"].sync("bwd_post", sync_op_or_fn="all_reduce")
-#         sch["layers.0"].cut_pipeline_stage()
+def test_pipeline_2stages_pp_tp_dp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 8:
+        pytest.skip("This test requires 8 GPUs.")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    with slapo.init_empty_weights():
+        model = Model(2, has_relu=True)
+    num_pp = 2
+    num_mp = 2
+    num_dp = dist.get_world_size() // (num_pp * num_mp)
+    print("num_dp:", num_dp, "num_pp:", num_pp, "num_mp:", num_mp)
+    topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
+    sch = slapo.create_schedule(model, group=group)
+    sch.trace_until("")
+
+    bs = 8
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=bs // num_dp,
+        fp16=False,
+    )
+    inp = torch.randn(bs, 10, device=dist.get_rank())
+    label = torch.randint(0, 10, (bs,), dtype=torch.long, device=dist.get_rank())
+    with slapo.Verify(
+        sch,
+        example_inputs=[inp],
+        example_outputs=label,
+        loss_fn=F.cross_entropy,
+        topology=topology,
+        config=ds_config_dict,
+    ):
+        # TODO: Fix layers.0/1 indexing bug
+        sch["layers.0.linear"].shard("weight", axis=0)
+        sch["layers.0.linear"].shard("bias", axis=0)
+        sch["layers.1.linear"].shard("weight", axis=1)
+        sch["layers.1.linear"].sync("fwd_post", sync_op_or_fn="all_reduce")
+        sch["layers.0.linear"].sync("bwd_post", sync_op_or_fn="all_reduce")
+        sch["layers.0"].cut_pipeline_stage()
 
 
 if __name__ == "__main__":
     test_pipeline_2stages_pp_dp()
     test_pipeline_2stages_pp_tp()
     test_pipeline_4stages_pp()
+    test_pipeline_2stages_pp_tp_dp()

From 3ba52099bd84466d75db876eaacdbff4977dcce1 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sat, 10 Jun 2023 19:57:03 +0000
Subject: [PATCH 13/28] Fix pylint

---
 slapo/verify.py           | 9 ++++++---
 tests/test_ds_pipeline.py | 6 +++---
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index dfc3debf..7070e996 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -84,11 +84,12 @@ def __exit__(self, *exc):
             return
         if self.sch.metadata.primitives["cut_pipeline_stage"]:
             try:
+                # pylint: disable=unused-import
                 import deepspeed
-            except ImportError:
+            except ImportError as exc:
                 raise ImportError(
                     "DeepSpeed is required when pipeline parallelism is used"
-                )
+                ) from exc
             assert (
                 self.example_outputs is not None
             ), "example_outputs must be provided when pipeline parallelism is used"
@@ -189,8 +190,10 @@ def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
                 if is_pipeline:
                     path = path.split(".")
-                    for idx, subpath in enumerate(path):
+                    idx = 0
+                    for i, subpath in enumerate(path):
                         if "submod_" not in subpath:
+                            idx = i
                             break
                     # Fix ModuleList name
                     path = ".".join(path[idx:]).replace("_", ".")
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 0b17a288..92ba9b61 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -2,11 +2,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """Test DeepSpeed Pipeline."""
+import os
 import pytest
 
-import os
 import torch
 from torch import nn
+import torch.distributed as dist
+import torch.nn.functional as F
 import deepspeed
 
 import slapo
@@ -14,8 +16,6 @@
     get_ds_config,
     create_dist_group_for_pipeline,
 )
-import torch.distributed as dist
-import torch.nn.functional as F
 
 
 class LinearReLU(nn.Module):

From ccda82fd690c88857738a354b34d3fe356e8438f Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 14:54:33 +0000
Subject: [PATCH 14/28] Fix naming

---
 slapo/verify.py | 31 +++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 7070e996..f554318f 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -185,19 +185,34 @@ def __exit__(self, *exc):
         # 7. Use original weights to initialize the new model
         #    Notice init_weights is called before actual sharding, so we only need to
         #    assign the original weights to the corresponding modules
+        original_param_names = original_state_dict.keys()
 
         def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
-                if is_pipeline:
-                    path = path.split(".")
-                    idx = 0
-                    for i, subpath in enumerate(path):
+                full_name = f"{path}.{name}"
+                # FIXME: this is a workaround for ModuleList
+                full_name = full_name.replace("layer_", "layer.")
+                if full_name not in original_param_names:
+                    # Remove all the leading submod_ in the full_name
+                    subpaths = full_name.split(".")
+                    new_subpaths = []
+                    for subpath in subpaths:
                         if "submod_" not in subpath:
-                            idx = i
+                            new_subpaths.append(subpath)
+                    full_name = ".".join(new_subpaths)
+                    # We only match the last part of the full_name
+                    # e.g., submod_1.submod_1.submod_1.layer.12.attention.self.query.weight
+                    # should match bert.encoder.layer.12.attention.self.query.weight
+                    for param_name in original_param_names:
+                        if param_name.endswith(full_name):
+                            original_name = param_name
                             break
-                    # Fix ModuleList name
-                    path = ".".join(path[idx:]).replace("_", ".")
-                original_name = f"{path}.{name}"
+                    else:
+                        raise RuntimeError(
+                            f"Cannot find the original parameter for {full_name}"
+                        )
+                else:
+                    original_name = full_name
                 setattr(
                     mod,
                     name,

From 563e1f4b02db63763cb9efec73698abc2ebf39f6 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 14:54:52 +0000
Subject: [PATCH 15/28] Add BERT PP test

---
 tests/test_ds_pipeline.py | 74 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 69 insertions(+), 5 deletions(-)

diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 92ba9b61..3ccea31a 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -4,6 +4,7 @@
 """Test DeepSpeed Pipeline."""
 import os
 import pytest
+import inspect
 
 import torch
 from torch import nn
@@ -16,6 +17,7 @@
     get_ds_config,
     create_dist_group_for_pipeline,
 )
+from slapo.op.cross_entropy import ParallelCrossEntropy
 
 
 class LinearReLU(nn.Module):
@@ -164,7 +166,6 @@ def test_pipeline_2stages_pp_tp_dp():
     num_pp = 2
     num_mp = 2
     num_dp = dist.get_world_size() // (num_pp * num_mp)
-    print("num_dp:", num_dp, "num_pp:", num_pp, "num_mp:", num_mp)
     topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
     sch = slapo.create_schedule(model, group=group)
     sch.trace_until("")
@@ -194,8 +195,71 @@ def test_pipeline_2stages_pp_tp_dp():
         sch["layers.0"].cut_pipeline_stage()
 
 
+def test_bert_2stages_pp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 2:
+        pytest.skip("This test requires 2 GPUs.")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    from transformers import BertLMHeadModel, AutoConfig
+
+    config = AutoConfig.from_pretrained("bert-large-uncased")
+    with slapo.init_empty_weights():
+        model = BertLMHeadModel(config)
+
+    num_pp = 2
+    num_mp = 1
+    num_dp = dist.get_world_size() // (num_pp * num_mp)
+    topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
+    sch = slapo.create_schedule(model, group=group)
+    input_names = ["input_ids", "attention_mask", "token_type_ids"]
+    sig = inspect.signature(sch.mod.forward)
+    concrete_args = {
+        p.name: p.default for p in sig.parameters.values() if p.name not in input_names
+    }
+    sch.trace_until(f"bert.encoder", tracer="huggingface", concrete_args=concrete_args)
+
+    loss_fct = ParallelCrossEntropy(group=group)
+
+    def loss_fn(outputs, labels):
+        # (bs, seq, vocab)
+        prediction_scores = outputs["logits"]
+        shifted_prediction_scores = prediction_scores[..., :-1, :].contiguous()
+        shifted_prediction_scores = shifted_prediction_scores
+        labels = labels[..., 1:].contiguous()
+        lm_loss = loss_fct(shifted_prediction_scores, labels)
+        lm_loss = lm_loss.contiguous().mean()
+        return lm_loss
+
+    bs = 2
+    seq_len = 512
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=bs // num_dp,
+        fp16=False,
+    )
+    device = "cuda"
+    input_ids = torch.ones(bs, seq_len, dtype=torch.long, device=device)
+    attention_mask = torch.ones(bs, seq_len, dtype=torch.float32, device=device)
+    token_type_ids = torch.ones(
+        bs, seq_len, dtype=torch.long, requires_grad=False, device=device
+    )
+    labels = torch.randint(0, 10, (bs, seq_len), dtype=torch.long, device=sch.rank)
+    with slapo.Verify(
+        sch,
+        example_inputs=[input_ids, attention_mask, token_type_ids],
+        example_outputs=labels,
+        loss_fn=loss_fn,
+        topology=topology,
+        config=ds_config_dict,
+    ):
+        sch[f"bert.encoder.layer.11"].cut_pipeline_stage()
+
+
 if __name__ == "__main__":
-    test_pipeline_2stages_pp_dp()
-    test_pipeline_2stages_pp_tp()
-    test_pipeline_4stages_pp()
-    test_pipeline_2stages_pp_tp_dp()
+    # test_pipeline_2stages_pp_dp()
+    # test_pipeline_2stages_pp_tp()
+    # test_pipeline_4stages_pp()
+    # test_pipeline_2stages_pp_tp_dp()
+    test_bert_2stages_pp()

From e912dae948738260232ff0ae3930d7b6dbff78e3 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 15:22:05 +0000
Subject: [PATCH 16/28] Fix dataloader

---
 slapo/verify.py           | 14 +++++++++++++-
 tests/test_ds_pipeline.py | 21 +++++++++++++++------
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index f554318f..5275ecbc 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -240,7 +240,19 @@ def init_weights(mod, path):
         #    make sure the random seeds are the same, which may affect the output of dropout
         set_random_seed(2023)
         if is_pipeline:
-            train_iter = iter([tuple(self.example_inputs + [self.example_outputs])])
+            from deepspeed.utils import RepeatingLoader
+
+            train_iter = RepeatingLoader(
+                [
+                    # First batch: (inputs, labels)
+                    (
+                        tuple(self.example_inputs),  # inputs
+                        self.example_outputs,  # labels
+                    ),
+                    # Rest of the batches
+                    # ...
+                ]
+            )
             # DeepSpeed will automatically broadcast the output to each device
             new_output = new_mod.train_batch(train_iter)
         else:
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 3ccea31a..88cf16d3 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -224,7 +224,11 @@ def test_bert_2stages_pp():
 
     def loss_fn(outputs, labels):
         # (bs, seq, vocab)
-        prediction_scores = outputs["logits"]
+        if isinstance(outputs, torch.Tensor):
+            # DS PP output has already removed the data structure
+            prediction_scores = outputs
+        else:
+            prediction_scores = outputs["logits"]
         shifted_prediction_scores = prediction_scores[..., :-1, :].contiguous()
         shifted_prediction_scores = shifted_prediction_scores
         labels = labels[..., 1:].contiguous()
@@ -234,18 +238,23 @@ def loss_fn(outputs, labels):
 
     bs = 2
     seq_len = 512
+    micro_bs = bs // num_dp
     ds_config_dict = get_ds_config(
         batch_size=bs,
-        micro_batch_size_per_gpu=bs // num_dp,
+        micro_batch_size_per_gpu=micro_bs,
         fp16=False,
     )
     device = "cuda"
-    input_ids = torch.ones(bs, seq_len, dtype=torch.long, device=device)
-    attention_mask = torch.ones(bs, seq_len, dtype=torch.float32, device=device)
+    input_ids = torch.ones(micro_bs, seq_len, dtype=torch.long, device=device)
+    attention_mask = torch.ones(
+        micro_bs, seq_len, dtype=torch.float32, requires_grad=False, device=device
+    )
     token_type_ids = torch.ones(
-        bs, seq_len, dtype=torch.long, requires_grad=False, device=device
+        micro_bs, seq_len, dtype=torch.long, requires_grad=False, device=device
+    )
+    labels = torch.randint(
+        0, 10, (micro_bs, seq_len), dtype=torch.long, device=sch.rank
     )
-    labels = torch.randint(0, 10, (bs, seq_len), dtype=torch.long, device=sch.rank)
     with slapo.Verify(
         sch,
         example_inputs=[input_ids, attention_mask, token_type_ids],

From d08b2b4b31d95cfb81abb35d05dcf6c98f9329c3 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 16:48:39 +0000
Subject: [PATCH 17/28] Remove pytest

---
 slapo/verify.py           |  1 +
 tests/test_ds_pipeline.py | 29 ++++++++++++++++++-----------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 5275ecbc..056a3572 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -192,6 +192,7 @@ def init_weights(mod, path):
                 full_name = f"{path}.{name}"
                 # FIXME: this is a workaround for ModuleList
                 full_name = full_name.replace("layer_", "layer.")
+                full_name = full_name.replace("layers_", "layers.")
                 if full_name not in original_param_names:
                     # Remove all the leading submod_ in the full_name
                     subpaths = full_name.split(".")
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 88cf16d3..c01dfa6d 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -3,7 +3,6 @@
 
 """Test DeepSpeed Pipeline."""
 import os
-import pytest
 import inspect
 
 import torch
@@ -18,6 +17,9 @@
     create_dist_group_for_pipeline,
 )
 from slapo.op.cross_entropy import ParallelCrossEntropy
+from slapo.logger import get_logger
+
+logger = get_logger()
 
 
 class LinearReLU(nn.Module):
@@ -48,7 +50,8 @@ def forward(self, x):
 def test_pipeline_2stages_pp_dp():
     deepspeed.init_distributed(dist_backend="nccl")
     if dist.get_world_size() != 4:
-        pytest.skip("This test requires 4 GPUs.")
+        logger.info("This test requires 4 GPUs.")
+        return
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     with slapo.init_empty_weights():
@@ -83,7 +86,8 @@ def test_pipeline_2stages_pp_dp():
 def test_pipeline_2stages_pp_tp():
     deepspeed.init_distributed(dist_backend="nccl")
     if dist.get_world_size() != 4:
-        pytest.skip("This test requires 4 GPUs.")
+        logger.info("This test requires 4 GPUs.")
+        return
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
 
@@ -124,7 +128,8 @@ def test_pipeline_2stages_pp_tp():
 def test_pipeline_4stages_pp():
     deepspeed.init_distributed(dist_backend="nccl")
     if dist.get_world_size() != 4:
-        pytest.skip("This test requires 4 GPUs.")
+        logger.info("This test requires 4 GPUs.")
+        return
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
     with slapo.init_empty_weights():
@@ -157,7 +162,8 @@ def test_pipeline_4stages_pp():
 def test_pipeline_2stages_pp_tp_dp():
     deepspeed.init_distributed(dist_backend="nccl")
     if dist.get_world_size() != 8:
-        pytest.skip("This test requires 8 GPUs.")
+        logger.info("This test requires 8 GPUs.")
+        return
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
 
@@ -198,7 +204,8 @@ def test_pipeline_2stages_pp_tp_dp():
 def test_bert_2stages_pp():
     deepspeed.init_distributed(dist_backend="nccl")
     if dist.get_world_size() != 2:
-        pytest.skip("This test requires 2 GPUs.")
+        logger.info("This test requires 2 GPUs.")
+        return
     local_rank = int(os.environ["LOCAL_RANK"])
     torch.cuda.set_device(local_rank)
 
@@ -267,8 +274,8 @@ def loss_fn(outputs, labels):
 
 
 if __name__ == "__main__":
-    # test_pipeline_2stages_pp_dp()
-    # test_pipeline_2stages_pp_tp()
-    # test_pipeline_4stages_pp()
-    # test_pipeline_2stages_pp_tp_dp()
-    test_bert_2stages_pp()
+    test_pipeline_2stages_pp_dp()
+    test_pipeline_2stages_pp_tp()
+    test_pipeline_4stages_pp()
+    test_pipeline_2stages_pp_tp_dp()
+    # test_bert_2stages_pp()

From e09d64433171f9e2cc91dd3c6e5d1015eafb6159 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 20:24:53 +0000
Subject: [PATCH 18/28] Support tie_weights

---
 slapo/build.py            | 18 ++++++++++++++----
 tests/test_ds_pipeline.py | 19 +++++++++++--------
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/slapo/build.py b/slapo/build.py
index 8ecc74e9..80e6fb6d 100644
--- a/slapo/build.py
+++ b/slapo/build.py
@@ -249,15 +249,25 @@ 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."
+                )
+            sch.mod.tie_weights()
 
     if sch.metadata.primitives["cut_pipeline_stage"] and target is not None:
         # Generate pipeline modules for a particular target.
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index c01dfa6d..8f438691 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -212,6 +212,7 @@ def test_bert_2stages_pp():
     from transformers import BertLMHeadModel, AutoConfig
 
     config = AutoConfig.from_pretrained("bert-large-uncased")
+    # config.tie_word_embeddings = False
     with slapo.init_empty_weights():
         model = BertLMHeadModel(config)
 
@@ -225,9 +226,8 @@ def test_bert_2stages_pp():
     concrete_args = {
         p.name: p.default for p in sig.parameters.values() if p.name not in input_names
     }
-    sch.trace_until(f"bert.encoder", tracer="huggingface", concrete_args=concrete_args)
 
-    loss_fct = ParallelCrossEntropy(group=group)
+    loss_fct = nn.CrossEntropyLoss()
 
     def loss_fn(outputs, labels):
         # (bs, seq, vocab)
@@ -236,11 +236,11 @@ def loss_fn(outputs, labels):
             prediction_scores = outputs
         else:
             prediction_scores = outputs["logits"]
-        shifted_prediction_scores = prediction_scores[..., :-1, :].contiguous()
-        shifted_prediction_scores = shifted_prediction_scores
-        labels = labels[..., 1:].contiguous()
-        lm_loss = loss_fct(shifted_prediction_scores, labels)
-        lm_loss = lm_loss.contiguous().mean()
+        shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+        lm_loss = loss_fct(
+            shifted_prediction_scores.view(-1, config.vocab_size), labels.view(-1)
+        )
         return lm_loss
 
     bs = 2
@@ -270,6 +270,9 @@ def loss_fn(outputs, labels):
         topology=topology,
         config=ds_config_dict,
     ):
+        sch.trace_until(
+            f"bert.encoder", tracer="huggingface", concrete_args=concrete_args
+        )
         sch[f"bert.encoder.layer.11"].cut_pipeline_stage()
 
 
@@ -278,4 +281,4 @@ def loss_fn(outputs, labels):
     test_pipeline_2stages_pp_tp()
     test_pipeline_4stages_pp()
     test_pipeline_2stages_pp_tp_dp()
-    # test_bert_2stages_pp()
+    test_bert_2stages_pp()

From 8023b0ebc97e3a35080b63688607e8c9ce8b2c60 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 21:00:00 +0000
Subject: [PATCH 19/28] Fix eval_batch

---
 slapo/verify.py | 50 +++++++++++++++++++++++++++++--------------------
 1 file changed, 30 insertions(+), 20 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 056a3572..18e953d9 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -93,9 +93,6 @@ def __exit__(self, *exc):
             assert (
                 self.example_outputs is not None
             ), "example_outputs must be provided when pipeline parallelism is used"
-            assert (
-                self.loss_fn is not None
-            ), "loss_fn must be provided when pipeline parallelism is used"
             assert (
                 self.topology is not None
             ), "topology must be provided when pipeline parallelism is used"
@@ -124,14 +121,15 @@ def __exit__(self, *exc):
                 group_src_rank = dist.get_global_rank(tp_pp_group, 0)
                 for inp in self.example_inputs:
                     dist.broadcast(inp, src=group_src_rank, group=tp_pp_group)
-                dist.broadcast(
-                    self.example_outputs, src=group_src_rank, group=tp_pp_group
-                )
+                if isinstance(self.example_outputs, torch.Tensor):
+                    dist.broadcast(
+                        self.example_outputs, src=group_src_rank, group=tp_pp_group
+                    )
         else:
             group_src_rank = 0
             for inp in self.example_inputs:
                 dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
-            if self.example_outputs is not None:
+            if isinstance(self.example_outputs, torch.Tensor):
                 dist.broadcast(
                     self.example_outputs, src=group_src_rank, group=self.sch.group
                 )
@@ -141,18 +139,19 @@ def __exit__(self, *exc):
             original_mod.eval()
         set_random_seed(2023)
         original_output = original_mod(*self.example_inputs)
-        if self.example_outputs is not None:
-            assert (
-                self.loss_fn is not None
-            ), "loss_fn must be provided when example_outputs is provided"
+        if self.loss_fn is not None:
+            assert isinstance(
+                self.example_outputs, torch.Tensor
+            ), "example_outputs must be provided when loss_fn is provided"
             original_output = self.loss_fn(original_output, self.example_outputs)
         if is_pipeline:
             # average the loss across all the DP devices
-            for ranks in self.topology.get_axis_comm_lists("data"):
-                dp_group = dist.new_group(ranks=ranks)
-                dist.all_reduce(original_output, group=dp_group)
-                if dist.get_rank() in ranks:
-                    original_output /= len(ranks)
+            if self.topology.get_dim("data") > 1:
+                for ranks in self.topology.get_axis_comm_lists("data"):
+                    dp_group = dist.new_group(ranks=ranks)
+                    dist.all_reduce(original_output, group=dp_group)
+                    if dist.get_rank() in ranks:
+                        original_output /= len(ranks)
         # 4. Broadcast the original model from rank 0 to other ranks
         #    Since for verification, each device holds an entire copy of the original
         #    model, here we directly broadcast the model to all the devices.
@@ -167,7 +166,7 @@ def __exit__(self, *exc):
             copied_mod = copy.deepcopy(self.sch.mod)
             is_copy_failed = False
         except TypeError:
-            # One example is ProcessGroup cannot be copied:
+            # One example is ProcessGroup that cannot be copied:
             # https://github.com/pytorch/pytorch/issues/73825
             is_copy_failed = True
             logger.warning(
@@ -243,7 +242,7 @@ def init_weights(mod, path):
         if is_pipeline:
             from deepspeed.utils import RepeatingLoader
 
-            train_iter = RepeatingLoader(
+            data_iter = RepeatingLoader(
                 [
                     # First batch: (inputs, labels)
                     (
@@ -255,11 +254,22 @@ def init_weights(mod, path):
                 ]
             )
             # DeepSpeed will automatically broadcast the output to each device
-            new_output = new_mod.train_batch(train_iter)
+            if self.eval_mode:
+                new_output = new_mod.eval_batch(
+                    data_iter, compute_loss=True if self.loss_fn else False
+                )
+            else:
+                new_output = new_mod.train_batch(data_iter)
         else:
             new_output = new_mod(*self.example_inputs)
         # 9. Compare the outputs
-        torch.testing.assert_close(original_output, new_output)
+        if is_pipeline:
+            if isinstance(original_output, dict):
+                original_output = original_output["logits"]
+        if new_output is not None:
+            if self.loss_fn is not None and new_output.shape != original_output.shape:
+                new_output = new_output.view(original_output.shape)
+            torch.testing.assert_close(original_output, new_output)
         logger.info("Passed verification!")
         if not is_copy_failed:
             del new_mod

From 24d9b1df39db302c86a8966246da1297171f42ab Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 21:07:37 +0000
Subject: [PATCH 20/28] Fix pylint

---
 slapo/verify.py           | 5 ++++-
 tests/test_ds_pipeline.py | 5 ++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 18e953d9..faeb30c7 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -1,5 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+# pylint: disable=too-many-branches
 
 import sys
 import copy
@@ -256,7 +257,7 @@ def init_weights(mod, path):
             # DeepSpeed will automatically broadcast the output to each device
             if self.eval_mode:
                 new_output = new_mod.eval_batch(
-                    data_iter, compute_loss=True if self.loss_fn else False
+                    data_iter, compute_loss=bool(self.loss_fn)
                 )
             else:
                 new_output = new_mod.train_batch(data_iter)
@@ -264,6 +265,8 @@ def init_weights(mod, path):
             new_output = new_mod(*self.example_inputs)
         # 9. Compare the outputs
         if is_pipeline:
+            # DS only outputs a single tensor, while the original
+            # HF model may output a dictionary
             if isinstance(original_output, dict):
                 original_output = original_output["logits"]
         if new_output is not None:
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 8f438691..82f2bc84 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -16,7 +16,6 @@
     get_ds_config,
     create_dist_group_for_pipeline,
 )
-from slapo.op.cross_entropy import ParallelCrossEntropy
 from slapo.logger import get_logger
 
 logger = get_logger()
@@ -271,9 +270,9 @@ def loss_fn(outputs, labels):
         config=ds_config_dict,
     ):
         sch.trace_until(
-            f"bert.encoder", tracer="huggingface", concrete_args=concrete_args
+            "bert.encoder", tracer="huggingface", concrete_args=concrete_args
         )
-        sch[f"bert.encoder.layer.11"].cut_pipeline_stage()
+        sch["bert.encoder.layer.11"].cut_pipeline_stage()
 
 
 if __name__ == "__main__":

From e77dfc91b9f0f078a7a2cf728f18c8d83bf53b05 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 21:40:54 +0000
Subject: [PATCH 21/28] Fix init_weights

---
 slapo/verify.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index faeb30c7..4db9dce7 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -30,6 +30,7 @@ def __init__(
         topology=None,
         eval_mode=True,
         enable=True,
+        init_weights=False,
         **kwargs,
     ):
         if not isinstance(example_inputs, list):
@@ -48,6 +49,7 @@ def __init__(
         self.topology = topology
         self.enable = enable
         self.eval_mode = eval_mode
+        self.init_weights = init_weights
         self.kwargs = kwargs
 
     def __enter__(self):
@@ -106,7 +108,12 @@ def __exit__(self, *exc):
         # 1. Build the original model with random weights
         named_params = self.original_sch.mod.named_parameters()
         is_initialized = named_params.__next__()[1].device != torch.device("meta")
-        original_mod, _ = build(self.original_sch, init_weights=not is_initialized)
+        original_mod, _ = build(
+            self.original_sch,
+            init_weights=self.init_weights
+            if self.init_weights
+            else (not is_initialized),
+        )
         #    make sure all the buffers are on the right device
         original_mod = original_mod.to(self.device)
         # 2. Get the example inputs and outputs

From 2af61a4e353bde97787447ae75978897d7079bfe Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 21:58:49 +0000
Subject: [PATCH 22/28] Fix DS inputs

---
 slapo/verify.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index 4db9dce7..e136a524 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=too-many-branches
 
+import re
 import sys
 import copy
 from contextlib import ContextDecorator
@@ -36,7 +37,16 @@ def __init__(
         if not isinstance(example_inputs, list):
             example_inputs = [example_inputs]
         self.device = device
-        self.example_inputs = [x.to(self.device) for x in example_inputs]
+        self.original_inputs = []
+        self.example_inputs = []
+        for x in example_inputs:
+            if isinstance(x, torch.Tensor):
+                x = x.to(self.device)
+                self.example_inputs.append(x)
+                self.original_inputs.append(x)
+            else:
+                # DS pipeline does not accept non-tensor inputs
+                self.original_inputs.append(x)
         self.example_outputs = (
             example_outputs.to(self.device)
             if isinstance(example_outputs, torch.Tensor)
@@ -135,8 +145,9 @@ def __exit__(self, *exc):
                     )
         else:
             group_src_rank = 0
-            for inp in self.example_inputs:
-                dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
+            for inp in self.original_inputs:
+                if isinstance(inp, torch.Tensor):
+                    dist.broadcast(inp, src=group_src_rank, group=self.sch.group)
             if isinstance(self.example_outputs, torch.Tensor):
                 dist.broadcast(
                     self.example_outputs, src=group_src_rank, group=self.sch.group
@@ -146,7 +157,7 @@ def __exit__(self, *exc):
         if self.eval_mode:
             original_mod.eval()
         set_random_seed(2023)
-        original_output = original_mod(*self.example_inputs)
+        original_output = original_mod(*self.original_inputs)
         if self.loss_fn is not None:
             assert isinstance(
                 self.example_outputs, torch.Tensor
@@ -197,9 +208,6 @@ def __exit__(self, *exc):
         def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
                 full_name = f"{path}.{name}"
-                # FIXME: this is a workaround for ModuleList
-                full_name = full_name.replace("layer_", "layer.")
-                full_name = full_name.replace("layers_", "layers.")
                 if full_name not in original_param_names:
                     # Remove all the leading submod_ in the full_name
                     subpaths = full_name.split(".")
@@ -207,6 +215,9 @@ def init_weights(mod, path):
                     for subpath in subpaths:
                         if "submod_" not in subpath:
                             new_subpaths.append(subpath)
+                    # FIXME: this is a workaround for ModuleList
+                    if re.match(r".*_[0-9]+", new_subpaths[0]):
+                        new_subpaths[0] = new_subpaths[0].replace("_", ".")
                     full_name = ".".join(new_subpaths)
                     # We only match the last part of the full_name
                     # e.g., submod_1.submod_1.submod_1.layer.12.attention.self.query.weight

From bb2ba5f3d6b28c90b82e0f5fe2822f6e847d6d27 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Sun, 11 Jun 2023 22:04:59 +0000
Subject: [PATCH 23/28] Add GPT2 PP test (not done yet)

---
 tests/test_ds_pipeline.py | 78 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 82f2bc84..d729cf3a 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -275,9 +275,87 @@ def loss_fn(outputs, labels):
         sch["bert.encoder.layer.11"].cut_pipeline_stage()
 
 
+def test_gpt2_2stages_pp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 2:
+        logger.info("This test requires 2 GPUs.")
+        return
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    from transformers import GPT2LMHeadModel, AutoConfig
+
+    config = AutoConfig.from_pretrained("gpt2-xl")
+    config.use_cache = False
+    with slapo.init_empty_weights():
+        model = GPT2LMHeadModel(config)
+
+    num_pp = 2
+    num_mp = 1
+    num_dp = dist.get_world_size() // (num_pp * num_mp)
+    topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
+    sch = slapo.create_schedule(model, group=group)
+    input_names = ["input_ids", "attention_mask", "token_type_ids"]
+    sig = inspect.signature(sch.mod.forward)
+    concrete_args = {
+        p.name: p.default for p in sig.parameters.values() if p.name not in input_names
+    }
+
+    loss_fct = nn.CrossEntropyLoss()
+
+    def loss_fn(outputs, labels):
+        # (bs, seq, vocab)
+        if isinstance(outputs, torch.Tensor):
+            # DS PP output has already removed the data structure
+            prediction_scores = outputs
+        else:
+            prediction_scores = outputs["logits"]
+        shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+        lm_loss = loss_fct(
+            shifted_prediction_scores.view(-1, config.vocab_size), labels.view(-1)
+        )
+        return lm_loss
+
+    bs = 2
+    seq_len = 512
+    micro_bs = bs // num_dp
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=micro_bs,
+        fp16=False,
+    )
+    device = "cuda"
+    input_ids = torch.ones(micro_bs, seq_len, dtype=torch.long, device=device)
+    attention_mask = torch.ones(
+        micro_bs, seq_len, dtype=torch.float32, requires_grad=False, device=device
+    )
+    position_ids = torch.ones(
+        micro_bs, seq_len, dtype=torch.long, requires_grad=False, device=device
+    )
+    labels = torch.randint(
+        0, 10, (micro_bs, seq_len), dtype=torch.long, device=sch.rank
+    )
+    with slapo.Verify(
+        sch,
+        # (input_ids, past_key_values, attention_mask, token_type_ids, position_ids)
+        example_inputs=[input_ids, None, attention_mask, None, position_ids],
+        example_outputs=labels,
+        loss_fn=loss_fn,
+        topology=topology,
+        config=ds_config_dict,
+        init_weights=model._init_weights,
+    ):
+        sch.trace_until(
+            "transformer", tracer="huggingface", concrete_args=concrete_args
+        )
+        sch["transformer.h.23"].cut_pipeline_stage()
+
+
 if __name__ == "__main__":
     test_pipeline_2stages_pp_dp()
     test_pipeline_2stages_pp_tp()
     test_pipeline_4stages_pp()
     test_pipeline_2stages_pp_tp_dp()
     test_bert_2stages_pp()
+    test_gpt2_2stages_pp()

From 961750632680018ab7afb0dc0d509aa4bfcec895 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Mon, 12 Jun 2023 05:23:41 +0000
Subject: [PATCH 24/28] Fix GPT2 test

---
 slapo/verify.py           | 2 ++
 tests/test_ds_pipeline.py | 3 ++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index e136a524..ca2444d5 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -288,6 +288,8 @@ def init_weights(mod, path):
             if isinstance(original_output, dict):
                 original_output = original_output["logits"]
         if new_output is not None:
+            # DS sometimes output shape-1 tensors, while the original
+            # HF model may output shape-0 tensors for loss
             if self.loss_fn is not None and new_output.shape != original_output.shape:
                 new_output = new_output.view(original_output.shape)
             torch.testing.assert_close(original_output, new_output)
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index d729cf3a..22213f2b 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -295,7 +295,7 @@ def test_gpt2_2stages_pp():
     num_dp = dist.get_world_size() // (num_pp * num_mp)
     topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
     sch = slapo.create_schedule(model, group=group)
-    input_names = ["input_ids", "attention_mask", "token_type_ids"]
+    input_names = ["input_ids", "attention_mask", "position_ids"]
     sig = inspect.signature(sch.mod.forward)
     concrete_args = {
         p.name: p.default for p in sig.parameters.values() if p.name not in input_names
@@ -317,6 +317,7 @@ def loss_fn(outputs, labels):
         )
         return lm_loss
 
+    # avoid OOM
     bs = 2
     seq_len = 512
     micro_bs = bs // num_dp

From 3400275f61476a4fcdbc40d923488e926d8f070e Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Mon, 12 Jun 2023 05:38:32 +0000
Subject: [PATCH 25/28] Add GPT2 4-stage test

---
 slapo/framework_dialect/deepspeed/pipeline.py |  2 +-
 slapo/verify.py                               |  2 +-
 tests/test_ds_pipeline.py                     | 93 +++++++++++++++++++
 3 files changed, 95 insertions(+), 2 deletions(-)

diff --git a/slapo/framework_dialect/deepspeed/pipeline.py b/slapo/framework_dialect/deepspeed/pipeline.py
index 31cb0d20..500c070d 100644
--- a/slapo/framework_dialect/deepspeed/pipeline.py
+++ b/slapo/framework_dialect/deepspeed/pipeline.py
@@ -364,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])
diff --git a/slapo/verify.py b/slapo/verify.py
index ca2444d5..df8563e2 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -1,6 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-# pylint: disable=too-many-branches
+# pylint: disable=too-many-branches, too-many-instance-attributes
 
 import re
 import sys
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index 22213f2b..c3707fd9 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -353,6 +353,98 @@ def loss_fn(outputs, labels):
         sch["transformer.h.23"].cut_pipeline_stage()
 
 
+def test_gpt2_4stages_pp():
+    """
+    def forward(self, input_ids : typing_Union[torch.LongTensor,NoneType] = None, past_key_values = None, attention_mask : typing_Union[torch.FloatTensor,NoneType] = None, token_type_ids = None, position_ids : typing_Union[torch.LongTensor,NoneType] = None, head_mask = None, inputs_embeds = None, encoder_hidden_states = None, encoder_attention_mask = None, labels = None, use_cache = None, output_attentions = None, output_hidden_states = None, return_dict = None):
+        submod_0 = self.submod_0(input_ids, position_ids, attention_mask);  input_ids = position_ids = attention_mask = None
+        submod_1 = self.submod_1(submod_0);  submod_0 = None
+        getitem = submod_1[0]
+        getitem_1 = submod_1[1]
+        getitem_2 = submod_1[2];  submod_1 = None
+        submod_2 = self.submod_2(getitem, getitem_1);  getitem = None
+        submod_3 = self.submod_3(submod_2, getitem_1, getitem_2);  submod_2 = getitem_1 = getitem_2 = None
+        return {'logits': submod_3, 'past_key_values': None, 'hidden_states': None, 'attentions': None, 'cross_attentions': None}
+    """
+    deepspeed.init_distributed(dist_backend="nccl")
+    mem = torch.cuda.get_device_properties(0).total_memory / 1024 / 1024 / 1024
+    if dist.get_world_size() != 4 or mem < 30:
+        logger.info("This test requires 4 GPUs with large memory (~32GB).")
+        return
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    from transformers import GPT2LMHeadModel, AutoConfig
+
+    config = AutoConfig.from_pretrained("gpt2-xl")
+    config.use_cache = False
+    with slapo.init_empty_weights():
+        model = GPT2LMHeadModel(config)
+
+    num_pp = 4
+    num_mp = 1
+    num_dp = dist.get_world_size() // (num_pp * num_mp)
+    topology, group = create_dist_group_for_pipeline(num_pp=num_pp, num_mp=num_mp)
+    sch = slapo.create_schedule(model, group=group)
+    input_names = ["input_ids", "attention_mask", "position_ids"]
+    sig = inspect.signature(sch.mod.forward)
+    concrete_args = {
+        p.name: p.default for p in sig.parameters.values() if p.name not in input_names
+    }
+
+    loss_fct = nn.CrossEntropyLoss()
+
+    def loss_fn(outputs, labels):
+        # (bs, seq, vocab)
+        if isinstance(outputs, torch.Tensor):
+            # DS PP output has already removed the data structure
+            prediction_scores = outputs
+        else:
+            prediction_scores = outputs["logits"]
+        shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+        lm_loss = loss_fct(
+            shifted_prediction_scores.view(-1, config.vocab_size), labels.view(-1)
+        )
+        return lm_loss
+
+    # avoid OOM
+    bs = 2
+    seq_len = 512
+    micro_bs = bs // num_dp
+    ds_config_dict = get_ds_config(
+        batch_size=bs,
+        micro_batch_size_per_gpu=micro_bs,
+        fp16=False,
+    )
+    device = "cuda"
+    input_ids = torch.ones(micro_bs, seq_len, dtype=torch.long, device=device)
+    attention_mask = torch.ones(
+        micro_bs, seq_len, dtype=torch.float32, requires_grad=False, device=device
+    )
+    position_ids = torch.ones(
+        micro_bs, seq_len, dtype=torch.long, requires_grad=False, device=device
+    )
+    labels = torch.randint(
+        0, 10, (micro_bs, seq_len), dtype=torch.long, device=sch.rank
+    )
+    with slapo.Verify(
+        sch,
+        # (input_ids, past_key_values, attention_mask, token_type_ids, position_ids)
+        example_inputs=[input_ids, None, attention_mask, None, position_ids],
+        example_outputs=labels,
+        loss_fn=loss_fn,
+        topology=topology,
+        config=ds_config_dict,
+        init_weights=model._init_weights,
+    ):
+        sch.trace_until(
+            "transformer", tracer="huggingface", concrete_args=concrete_args
+        )
+        sch["transformer.h.11"].cut_pipeline_stage()
+        sch["transformer.h.23"].cut_pipeline_stage()
+        sch["transformer.h.35"].cut_pipeline_stage()
+
+
 if __name__ == "__main__":
     test_pipeline_2stages_pp_dp()
     test_pipeline_2stages_pp_tp()
@@ -360,3 +452,4 @@ def loss_fn(outputs, labels):
     test_pipeline_2stages_pp_tp_dp()
     test_bert_2stages_pp()
     test_gpt2_2stages_pp()
+    test_gpt2_4stages_pp()

From 867b85df6b077d669f23c0b69ff2fcbf00301230 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Mon, 12 Jun 2023 05:55:23 +0000
Subject: [PATCH 26/28] Add fp16 support

---
 slapo/verify.py           |  9 ++++++++-
 tests/test_ds_pipeline.py | 10 ++++++----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/slapo/verify.py b/slapo/verify.py
index df8563e2..ac887ec8 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -115,6 +115,8 @@ def __exit__(self, *exc):
             is_pipeline = True
         else:
             is_pipeline = False
+        if "dtype" in self.kwargs:
+            logger.info("Using %s data type", self.kwargs["dtype"], ranks=0)
         # 1. Build the original model with random weights
         named_params = self.original_sch.mod.named_parameters()
         is_initialized = named_params.__next__()[1].device != torch.device("meta")
@@ -126,6 +128,8 @@ def __exit__(self, *exc):
         )
         #    make sure all the buffers are on the right device
         original_mod = original_mod.to(self.device)
+        #    with the correct data type
+        original_mod = original_mod.to(self.kwargs.get("dtype", torch.float32))
         # 2. Get the example inputs and outputs
         #    Broadcast the example inputs from rank 0 in each TP/PP group
         #    to other ranks in the same group.
@@ -253,7 +257,9 @@ def init_weights(mod, path):
             new_mod, _ = build(new_sch, init_weights=init_weights)
         # 8. Run the new model
         #    make sure all the buffers are on the right device
-        new_mod.to(self.device)
+        new_mod = new_mod.to(self.device)
+        #    with the correct data type
+        new_mod = new_mod.to(self.kwargs.get("dtype", torch.float32))
         if self.eval_mode:
             new_mod.eval()
         #    make sure the random seeds are the same, which may affect the output of dropout
@@ -292,6 +298,7 @@ def init_weights(mod, path):
             # HF model may output shape-0 tensors for loss
             if self.loss_fn is not None and new_output.shape != original_output.shape:
                 new_output = new_output.view(original_output.shape)
+            new_output = new_output.to(original_output.dtype)
             torch.testing.assert_close(original_output, new_output)
         logger.info("Passed verification!")
         if not is_copy_failed:
diff --git a/tests/test_ds_pipeline.py b/tests/test_ds_pipeline.py
index c3707fd9..f8e3957d 100644
--- a/tests/test_ds_pipeline.py
+++ b/tests/test_ds_pipeline.py
@@ -324,12 +324,12 @@ def loss_fn(outputs, labels):
     ds_config_dict = get_ds_config(
         batch_size=bs,
         micro_batch_size_per_gpu=micro_bs,
-        fp16=False,
+        fp16=True,
     )
     device = "cuda"
     input_ids = torch.ones(micro_bs, seq_len, dtype=torch.long, device=device)
     attention_mask = torch.ones(
-        micro_bs, seq_len, dtype=torch.float32, requires_grad=False, device=device
+        micro_bs, seq_len, dtype=torch.float16, requires_grad=False, device=device
     )
     position_ids = torch.ones(
         micro_bs, seq_len, dtype=torch.long, requires_grad=False, device=device
@@ -346,6 +346,7 @@ def loss_fn(outputs, labels):
         topology=topology,
         config=ds_config_dict,
         init_weights=model._init_weights,
+        dtype=torch.float16,
     ):
         sch.trace_until(
             "transformer", tracer="huggingface", concrete_args=concrete_args
@@ -414,12 +415,12 @@ def loss_fn(outputs, labels):
     ds_config_dict = get_ds_config(
         batch_size=bs,
         micro_batch_size_per_gpu=micro_bs,
-        fp16=False,
+        fp16=True,
     )
     device = "cuda"
     input_ids = torch.ones(micro_bs, seq_len, dtype=torch.long, device=device)
     attention_mask = torch.ones(
-        micro_bs, seq_len, dtype=torch.float32, requires_grad=False, device=device
+        micro_bs, seq_len, dtype=torch.float16, requires_grad=False, device=device
     )
     position_ids = torch.ones(
         micro_bs, seq_len, dtype=torch.long, requires_grad=False, device=device
@@ -436,6 +437,7 @@ def loss_fn(outputs, labels):
         topology=topology,
         config=ds_config_dict,
         init_weights=model._init_weights,
+        dtype=torch.float16,
     ):
         sch.trace_until(
             "transformer", tracer="huggingface", concrete_args=concrete_args

From 4c6ac2aeb6f95a6a61fc2a860b9b41f916f0f86b Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Mon, 12 Jun 2023 19:37:33 +0000
Subject: [PATCH 27/28] Fix verify

---
 slapo/build.py       |  1 +
 slapo/verify.py      | 10 +++++-----
 tests/test_verify.py |  4 +++-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/slapo/build.py b/slapo/build.py
index 80e6fb6d..143f1e11 100644
--- a/slapo/build.py
+++ b/slapo/build.py
@@ -267,6 +267,7 @@ def build(
                 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:
diff --git a/slapo/verify.py b/slapo/verify.py
index ac887ec8..6d695d32 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -288,11 +288,11 @@ def init_weights(mod, path):
         else:
             new_output = new_mod(*self.example_inputs)
         # 9. Compare the outputs
-        if is_pipeline:
-            # DS only outputs a single tensor, while the original
-            # HF model may output a dictionary
-            if isinstance(original_output, dict):
-                original_output = original_output["logits"]
+        # Original HF model may output a dictionary
+        if isinstance(original_output, dict):
+            original_output = original_output["logits"]
+        if isinstance(new_output, dict):
+            new_output = new_output["logits"]
         if new_output is not None:
             # DS sometimes output shape-1 tensors, while the original
             # HF model may output shape-0 tensors for loss
diff --git a/tests/test_verify.py b/tests/test_verify.py
index 5e8850f2..e0f414d2 100644
--- a/tests/test_verify.py
+++ b/tests/test_verify.py
@@ -175,7 +175,9 @@ def new_view(tensor, args):
     bs = 2
     seq = 512
     input_ids = torch.ones(bs, seq, dtype=torch.long, device=sch.rank)
-    with slapo.Verify(sch, [input_ids], eval_mode=True):
+    with slapo.Verify(
+        sch, [input_ids], eval_mode=True, init_weights=model._init_weights
+    ):
         for i in range(config.num_hidden_layers):
             # shard attention
             subsch = sch[f"bert.encoder.layer.{i}.attention.self"]

From d3b3240fd5591115bac12612a099b889a69347e0 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Wed, 14 Jun 2023 17:53:19 +0000
Subject: [PATCH 28/28] Fix gptneo

---
 slapo/framework_dialect/deepspeed/pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/slapo/framework_dialect/deepspeed/pipeline.py b/slapo/framework_dialect/deepspeed/pipeline.py
index 500c070d..1739e205 100644
--- a/slapo/framework_dialect/deepspeed/pipeline.py
+++ b/slapo/framework_dialect/deepspeed/pipeline.py
@@ -155,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(