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/slapo/build.py b/slapo/build.py
index 8ecc74e9..143f1e11 100644
--- a/slapo/build.py
+++ b/slapo/build.py
@@ -249,15 +249,26 @@ def build(
     if sch.metadata.primitives["cut_pipeline_stage"]:
         # pipeline stages will be wrapped into PipeStageWrapper
         sch = generate_pipeline_partition(sch)
-        # Re-analyzie tie weights before consolidation.
-        sch.metadata.tie_weights = analyze_tie_weights(
-            sch.mod, is_pipeline_partitioned=True
-        )
+        is_pipeline = True
+    else:
+        is_pipeline = False
+    # Re-analyzie tie weights before consolidation.
+    sch.metadata.tie_weights = analyze_tie_weights(
+        sch.mod, is_pipeline_partitioned=is_pipeline
+    )
 
     # delay initialization
     if init_weights:
         init_weight_fn = init_weights if isinstance(init_weights, Callable) else None
         sch = consolidate_model(sch, target, init_weight_fn, **kwargs)
+        tie_weights = list(sch.metadata.tie_weights.values())
+        if tie_weights and not is_pipeline:
+            if not hasattr(sch.mod, "tie_weights"):
+                raise RuntimeError(
+                    "Model needs to tie weights but does not have `tie_weights` method. Probably because the model has been traced."
+                )
+            # https://github.com/huggingface/transformers/blob/v4.28.1/src/transformers/modeling_utils.py#L1274-L1277
+            sch.mod.tie_weights()
 
     if sch.metadata.primitives["cut_pipeline_stage"] and target is not None:
         # Generate pipeline modules for a particular target.
diff --git a/slapo/framework_dialect/deepspeed/pipeline.py b/slapo/framework_dialect/deepspeed/pipeline.py
index c220821c..1739e205 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.
@@ -88,7 +155,7 @@ def flat_and_name_tensor_list(data, name, suffix):
     values = data.values() if isinstance(data, dict) else data
 
     if isinstance(values, (list, tuple)) and any(
-        isinstance(t, torch.Tensor) for t in values
+        isinstance(t, (torch.Tensor, torch.Size)) for t in values
     ):
         for idx, tensor in enumerate(values):
             name_n_value.extend(
@@ -297,7 +364,7 @@ def forward(self, *args, **kwargs):
                 ), f"[{self.name}] Arg {arg_name} not found in liveness list: {liveness}"
                 idx = liveness.index(arg_name)
             ordered_args.append(unordered_args[idx])
-            if i > 0:
+            if i > 0 and not self.last:
                 # https://github.com/microsoft/DeepSpeed/blob/v0.9.2/deepspeed/runtime/pipe/engine.py#L639
                 assert (
                     torch.is_tensor(unordered_args[idx])
diff --git a/slapo/verify.py b/slapo/verify.py
index 591273b2..6d695d32 100644
--- a/slapo/verify.py
+++ b/slapo/verify.py
@@ -1,6 +1,8 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+# pylint: disable=too-many-branches, too-many-instance-attributes
 
+import re
 import sys
 import copy
 from contextlib import ContextDecorator
@@ -19,16 +21,46 @@
 
 
 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,
+        init_weights=False,
+        **kwargs,
+    ):
         if not isinstance(example_inputs, list):
             example_inputs = [example_inputs]
-        self.example_inputs = example_inputs
+        self.device = device
+        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)
+            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.init_weights = init_weights
+        self.kwargs = kwargs
 
     def __enter__(self):
         self.original_trace = sys.gettrace()
@@ -56,70 +88,219 @@ 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
+        if self.sch.metadata.primitives["cut_pipeline_stage"]:
+            try:
+                # pylint: disable=unused-import
+                import deepspeed
+            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"
+            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"
+            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")
-        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
-        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)
+        #    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.
+        #    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)
+                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.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
+                )
         # 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)
+        original_output = original_mod(*self.original_inputs)
+        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
+            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.
         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=0)
         # 5. Delete the original model to avoid excessive memory usage
         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 that 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():
             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
+        original_param_names = original_state_dict.keys()
 
         def init_weights(mod, path):
             for name, _ in mod.named_parameters(recurse=False):
+                full_name = f"{path}.{name}"
+                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:
+                            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
+                    # 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
+                    else:
+                        raise RuntimeError(
+                            f"Cannot find the original parameter for {full_name}"
+                        )
+                else:
+                    original_name = full_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 is_pipeline:
+            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)
+        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
         set_random_seed(2023)
-        new_output = new_mod(*self.example_inputs)
+        if is_pipeline:
+            from deepspeed.utils import RepeatingLoader
+
+            data_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
+            if self.eval_mode:
+                new_output = new_mod.eval_batch(
+                    data_iter, compute_loss=bool(self.loss_fn)
+                )
+            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)
+        # 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
+            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!")
-        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
new file mode 100644
index 00000000..f8e3957d
--- /dev/null
+++ b/tests/test_ds_pipeline.py
@@ -0,0 +1,457 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Test DeepSpeed Pipeline."""
+import os
+import inspect
+
+import torch
+from torch import nn
+import torch.distributed as dist
+import torch.nn.functional as F
+import deepspeed
+
+import slapo
+from slapo.framework_dialect.deepspeed.pipeline import (
+    get_ds_config,
+    create_dist_group_for_pipeline,
+)
+from slapo.logger import get_logger
+
+logger = get_logger()
+
+
+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, has_relu=False):
+        super().__init__()
+        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:
+            x = layer(x)
+        return x
+
+
+def test_pipeline_2stages_pp_dp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 4:
+        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():
+        model = Model(12)
+    # 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 // 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.5"].cut_pipeline_stage()
+
+
+def test_pipeline_2stages_pp_tp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 4:
+        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():
+        model = Model(2, has_relu=True)
+    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,
+    ):
+        # 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()
+
+
+def test_pipeline_4stages_pp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 4:
+        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():
+        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()
+
+
+def test_pipeline_2stages_pp_tp_dp():
+    deepspeed.init_distributed(dist_backend="nccl")
+    if dist.get_world_size() != 8:
+        logger.info("This test requires 8 GPUs.")
+        return
+    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)
+    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()
+
+
+def test_bert_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 BertLMHeadModel, AutoConfig
+
+    config = AutoConfig.from_pretrained("bert-large-uncased")
+    # config.tie_word_embeddings = False
+    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
+    }
+
+    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
+    )
+    token_type_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,
+        example_inputs=[input_ids, attention_mask, token_type_ids],
+        example_outputs=labels,
+        loss_fn=loss_fn,
+        topology=topology,
+        config=ds_config_dict,
+    ):
+        sch.trace_until(
+            "bert.encoder", tracer="huggingface", concrete_args=concrete_args
+        )
+        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", "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=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.float16, 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,
+        dtype=torch.float16,
+    ):
+        sch.trace_until(
+            "transformer", tracer="huggingface", concrete_args=concrete_args
+        )
+        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=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.float16, 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,
+        dtype=torch.float16,
+    ):
+        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()
+    test_pipeline_4stages_pp()
+    test_pipeline_2stages_pp_tp_dp()
+    test_bert_2stages_pp()
+    test_gpt2_2stages_pp()
+    test_gpt2_4stages_pp()
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"]