diff --git a/bmtrain/__init__.py b/bmtrain/__init__.py
index 3e025846..25c72052 100644
--- a/bmtrain/__init__.py
+++ b/bmtrain/__init__.py
@@ -10,7 +10,7 @@
 from .block_layer import CheckpointBlock, TransformerBlockList
 from .backward import optim_step
 from .wrapper import BMTrainModelWrapper
-
+from .pipe_layer import PipelineTransformerBlockList
 from . import debug
 from .store import save, load
 
diff --git a/bmtrain/benchmark/__init__.py b/bmtrain/benchmark/__init__.py
index 237cfca4..571d621f 100644
--- a/bmtrain/benchmark/__init__.py
+++ b/bmtrain/benchmark/__init__.py
@@ -1,2 +1,3 @@
 from .all_gather import all_gather
-from .reduce_scatter import reduce_scatter
\ No newline at end of file
+from .reduce_scatter import reduce_scatter
+from .send_recv import send_recv
\ No newline at end of file
diff --git a/bmtrain/benchmark/send_recv.py b/bmtrain/benchmark/send_recv.py
new file mode 100644
index 00000000..e3c971e4
--- /dev/null
+++ b/bmtrain/benchmark/send_recv.py
@@ -0,0 +1,31 @@
+from .. import nccl
+from .shape import SHAPES
+from ..global_var import config
+from ..utils import print_rank
+from .utils import format_size
+import torch
+def send_recv():
+    current_stream = torch.cuda.current_stream()
+    for shape in SHAPES:
+        send_size = shape
+
+        send_buffer = torch.empty( send_size // 2, dtype=torch.half, device="cuda" )
+        recv_buffer = torch.empty( send_size // 2, dtype=torch.half, device="cuda" )
+        
+        start_evt = torch.cuda.Event(enable_timing=True)
+        end_evt = torch.cuda.Event(enable_timing=True)
+
+        current_stream.record_event(start_evt)
+        nccl.groupStart()
+        if config['rank'] in [0,2,4,6]:
+            nccl.send(send_buffer.storage(), config['rank']+1, config['comm'])
+        else:
+            nccl.recv(recv_buffer.storage(), config['rank']-1, config['comm'])
+        nccl.groupEnd()
+        current_stream.record_event(end_evt)
+        current_stream.synchronize()
+        time_usage = start_evt.elapsed_time(end_evt)
+
+        bw = shape / 1024 / 1024 / 1024 * 1000 / time_usage
+        print_rank("Send Recv:\tsize {}\ttime: {:4.3f}\tbw: {:2.6f} GB/s".format(format_size(send_size), time_usage, bw))
+
diff --git a/bmtrain/block_layer.py b/bmtrain/block_layer.py
index 46c47572..138e5d08 100644
--- a/bmtrain/block_layer.py
+++ b/bmtrain/block_layer.py
@@ -1,6 +1,6 @@
-from typing import Dict, Iterable, Iterator, Union
-
+from typing import Dict, Iterable, Iterator, Union, List
 
+from .utils import round_up
 from .global_var import config
 import torch
 from . import nccl
@@ -10,8 +10,6 @@
 from . import debug
 import copy
 
-def round_up(x, d):
-    return (x + d - 1) // d * d
 
 # the flag is used to control the zero level , 0 means normal zero3 , 1 means forward without release parameter ,2 means backward without gather parameter
 class OpCheckpointBlock(torch.autograd.Function):
@@ -120,16 +118,19 @@ def backward(ctx, *grad_outputs):
         return (None, None, None, None) + tuple(grads)
 
 class CheckpointBlockContext:
-    def __init__(self ,block : 'CheckpointBlock',ctx_dict : dict = None, flag : int = 0) -> None:
+    def __init__(self, block : 'CheckpointBlock', ctx_dict : dict = None, flag : int = 0, pipe = False) -> None:
         self.block = block
-        self.ctx_dict=ctx_dict
+        self.ctx_dict = ctx_dict
         self._param_buffer = {}
         self._grad_buffer = {}
         self._param_tensor = {}
         self._grad_tensor = {}
         self.flag = flag
         self._need_release = False
-    
+        if pipe:
+            self.comm = config["zero_comm"] 
+        else:
+            self.comm = config["comm"]
     def enter(self):
         """
         gather parameters
@@ -146,15 +147,15 @@ def enter(self):
                 assert self.block._storage_params[kw].is_cuda
                 assert kw not in self._grad_buffer
                 assert kw not in self._param_buffer
-                local_param = self.block._storage_params[kw]
-
+                local_param = self.block._storage_params[kw]    
+           
                 storage_type = local_param.storage_type()
                 if self.flag != 2:
-                    self._param_buffer[kw] = storage_type(val["partition_size"] * config["world_size"])
+                    self._param_buffer[kw] = storage_type(val["partition_size"] * val["world_size"])
                     self._param_tensor[kw] = torch.tensor([], dtype=self._param_buffer[kw].dtype, device=self._param_buffer[kw].device).set_(self._param_buffer[kw])
 
                 if requires_grad and local_param.requires_grad:
-                    self._grad_buffer[kw] = storage_type(val["partition_size"] * config["world_size"])
+                    self._grad_buffer[kw] = storage_type(val["partition_size"] * val["world_size"])
                     self._grad_tensor[kw] = torch.tensor([], dtype=self._grad_buffer[kw].dtype, device=self._grad_buffer[kw].device).set_(self._grad_buffer[kw]).zero_()
             if self.flag != 2:
                 nccl.groupStart()
@@ -162,7 +163,7 @@ def enter(self):
                     nccl.allGather(
                         self.block._storage_params[kw].storage(),
                         self._param_buffer[kw],
-                        config["comm"]
+                        self.comm
                     )
                 nccl.groupEnd()
 
@@ -177,25 +178,25 @@ def enter(self):
                 self._grad_tensor[kw].record_stream(current_stream)
 
         # update parameters in block
-            for param in self.block._param_info:
-                kw_name = param["kw_name"]
-                offset = param["offset"]
-                shape = param["shape"]
+        for param in self.block._param_info:
+            kw_name = param["kw_name"]
+            offset = param["offset"]
+            shape = param["shape"]
 
-                if self.flag != 2:
-                    dtype = self._param_buffer[kw_name].dtype
-                    device = self._param_buffer[kw_name].device
-                    param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape)                
-                else:
-                    dtype = param["parameter"].data.dtype
-                    device = param["parameter"].data.device
-                    param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self.ctx_dict[kw_name], offset, shape)
+            if self.flag != 2:
+                dtype = self._param_buffer[kw_name].dtype
+                device = self._param_buffer[kw_name].device
+                param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape)                
+            else:
+                dtype = param["parameter"].data.dtype
+                device = param["parameter"].data.device
+                param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self.ctx_dict[kw_name], offset, shape)
 
-                if requires_grad and kw_name in self._grad_buffer:
-                    param["parameter"].requires_grad_(True)
-                    param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape)
-                else:
-                    param["parameter"].requires_grad_(False)
+            if requires_grad and kw_name in self._grad_buffer:
+                param["parameter"].requires_grad_(True)
+                param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape)
+            else:
+                param["parameter"].requires_grad_(False)
 
     def __enter__(self):
         self.enter()
@@ -209,7 +210,6 @@ def exit(self):
             return
         self._need_release = False
         self.block._ready = False
-
         requires_grad = torch.is_grad_enabled()
         if requires_grad:
             for kw, val in self.block._storage_info.items():
@@ -237,7 +237,7 @@ def exit(self):
                             self._grad_buffer[kw],
                             local_param.grad.storage(),
                             "sum",
-                            config["comm"]
+                            self.comm
                         )
                 nccl.groupEnd()
 
@@ -355,7 +355,8 @@ def __init__(self, inner_module : torch.nn.Module):
         offsets = {}
         # intialize storage buffers
         for kw, val in self._storage_info.items():
-            partition_size = round_up(val["total"], config['world_size']) // config['world_size']
+            val["world_size"] = config["world_size"]
+            partition_size = round_up(val["total"], val["world_size"]) // val["world_size"]
             val["partition_size"] = partition_size
             val["begin"] = config['rank'] * partition_size
             val["end"] = (config['rank'] + 1) * partition_size
@@ -538,7 +539,6 @@ def init_parameters(self):
                 # initialzie here
                 tmp_tensor = torch.empty(it["shape"], device=param.device, dtype=param.dtype)
                 param._init_method(tmp_tensor)
-                
                 param_st = it["offset"]
                 param_end = it["offset"] + it["size"]
                 kw_name = it["kw_name"]
@@ -604,7 +604,6 @@ def named_modules(self, memo = None, prefix: str = '', remove_duplicate: bool =
             1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
 
         """
-        # print("here in named_modules hahaha")
 
         if memo is None:
             memo = set()
@@ -654,7 +653,7 @@ def forward(ctx, placeholder, self : 'TransformerBlockList', save_list, hidden_s
         with torch.no_grad():
             for i in range(len(self)):
                 if save_list[i][0] == i:
-                    layer_inputs.append(hidden_state)
+                    layer_inputs.append(hidden_state.detach())
                 cuda_rng_state.append( torch.cuda.get_rng_state() )
                 if config['zero_level']==2:
                     flag = 1
@@ -674,13 +673,20 @@ def forward(ctx, placeholder, self : 'TransformerBlockList', save_list, hidden_s
         ctx.layer_inspector = layer_inspector
         ctx.cuda_rng_state = cuda_rng_state
 
-        
         ctx.save_for_backward(*layer_inputs, *tensors)
-        return hidden_state
+
+        if self.return_hidden_states:
+            middle_hiddens = layer_inputs 
+            for mid in middle_hiddens:
+                mid.requires_grad_()
+            middle_hiddens = torch.stack(middle_hiddens, dim=0)
+            return hidden_state, middle_hiddens
+        else:
+            return hidden_state, None
 
 
     @staticmethod
-    def backward(ctx, grad_hidden_state : torch.Tensor):
+    def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle: List[torch.Tensor]):
         def exit_prev(prev_ctx, prev_grad):
             if prev_ctx is not None:
                 if prev_grad:
@@ -738,7 +744,7 @@ def exit_prev(prev_ctx, prev_grad):
                                 ctx.save_list[j+1][0] = j+1
                 
                     torch.cuda.set_rng_state(ctx.cuda_rng_state[i])
-                    ipt = layer_inputs[ctx.save_list[i][1]].detach().requires_grad_()
+                    ipt = layer_inputs[ctx.save_list[i][1]].requires_grad_()
                     if config['zero_level'] == 2:
                         flag = 2
                     else:
@@ -759,12 +765,15 @@ def exit_prev(prev_ctx, prev_grad):
                         assert it["group"] == ctx.layer_inspector[i][j]["group"], "Backward step changed"
                         
                         # change the tensor in placeholder
+                        ctx.layer_inspector[i][j]["requires_grad"] = it["requires_grad"]
                         ctx.layer_inspector[i][j]["tensor"] = it["tensor"]
                     torch.autograd.backward(
                         [output],
                         [grad_hidden_state]
                     )
                     grad_hidden_state = ipt.grad
+                    if grad_middle is not None:
+                        grad_hidden_state = grad_hidden_state + grad_middle[i]
                 
                 exit_prev(prev_ctx, prev_grad)
 
@@ -836,6 +845,11 @@ def __iter__(self) -> Iterator[CheckpointBlock]:
     def __getitem__(self, index: Union[int, str]) -> CheckpointBlock:
         return self._modules[str(index)]
 
-    def forward(self, hidden_state, *args):
+    def forward(self, hidden_state, *args, return_hidden_states = False):
+        self.return_hidden_states = return_hidden_states
         placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled())
-        return OpTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args)
+        last_hidden, middle_hiddens = OpTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args)
+        if return_hidden_states:
+            return last_hidden, middle_hiddens
+        else:
+            return last_hidden
diff --git a/bmtrain/distributed/__init__.py b/bmtrain/distributed/__init__.py
index 0cb1fd5a..9dc64bb8 100644
--- a/bmtrain/distributed/__init__.py
+++ b/bmtrain/distributed/__init__.py
@@ -1 +1 @@
-from .ops import all_gather, all_reduce
+from .ops import all_gather, all_reduce, broadcast, recv_activations, send_activations
diff --git a/bmtrain/distributed/ops.py b/bmtrain/distributed/ops.py
index 9cfff1b9..6124622d 100644
--- a/bmtrain/distributed/ops.py
+++ b/bmtrain/distributed/ops.py
@@ -1,29 +1,82 @@
 import torch
 from ..global_var import config
-from ..nccl import allGather as ncclAllGather
+from ..nccl import allGather as ncclAllGather, recv
 from ..nccl import allReduce as ncclAllReduce
-
+from ..nccl import broadcast as ncclBroadcast
+from ..nccl import send as ncclSend
+from ..nccl import recv as ncclRecv
+from ..nccl import commCount,commRank,NCCLCommunicator
+DTYPE_LIST = [
+    torch.float64,
+    torch.float32,
+    torch.float16,
+    torch.int64,
+    torch.int32,
+    torch.int16,
+    torch.int8,
+    torch.bfloat16,
+    torch.bool
+]
+def send_activations(hidden_state, next_rank, comm):
+    send_meta(hidden_state, next_rank, comm)
+    ncclSend(hidden_state.storage(), next_rank, comm)
+def recv_activations(prev_rank, comm):
+    dtype, shape = recv_meta(prev_rank, comm)
+    hidden_state = torch.empty(shape, dtype=dtype, device="cuda")
+    ncclRecv(hidden_state.storage(), prev_rank, comm)
+    return hidden_state
+def send_meta(x, next_rank, comm):
+    meta = [len(x.size()), DTYPE_LIST.index(x.dtype)] + list(x.size())
+    meta_data = torch.tensor(data=meta, device=x.device, dtype=torch.long)
+    ncclSend(meta_data.storage(), next_rank, comm)
+def recv_meta(prev_rank, comm):
+    meta_data = torch.tensor(data=[0]*50, device="cuda", dtype=torch.long)
+    ncclRecv(meta_data.storage(), prev_rank, comm)
+    n_dims = meta_data[0].item()
+    dtype = DTYPE_LIST[meta_data[1].item()]
+    shape = meta_data[2:n_dims+2].tolist()
+    return dtype,shape
+class OpBroadcast(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, src, root, comm = None):
+        if comm is None:
+            comm = config["comm"]
+        ctx.comm = comm
+        outputs = torch.empty_like(src, dtype = src.dtype, device = src.device)
+        ncclBroadcast(src.storage(), outputs.storage(), root, comm)
+        return outputs
+    @staticmethod
+    def backward(ctx, grad_output):
+        res = all_reduce(grad_output, "sum", ctx.comm)
+        return res, None, None
+def broadcast(src, root, comm=None):
+    if not config["initialized"]:
+        raise RuntimeError("BMTrain is not initialized")
+    return OpBroadcast.apply(src, root, comm)
 class OpAllGather(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, input : torch.Tensor):
+    def forward(ctx, input : torch.Tensor, comm = None):
+        if comm is None:
+            comm = config["comm"]
+        world_size = commCount(comm)
         if not input.is_contiguous():
             input = input.contiguous()
         if input.storage_offset() != 0 or input.storage().size() != input.numel():
             input = input.clone()
-        output = torch.empty( (config['world_size'],) + input.size(), dtype=input.dtype, device=input.device)
-        
+        output = torch.empty( (world_size,) + input.size(), dtype=input.dtype, device=input.device)
+        ctx.comm = comm
         ncclAllGather(
             input.storage(),
             output.storage(),
-            config['comm']
+            comm
         )
         return output
 
     @staticmethod
     def backward(ctx, grad_output):
-        return grad_output[config['rank']]
+        return grad_output[commRank(ctx.comm)], None
 
-def all_gather(x : torch.Tensor):
+def all_gather(x : torch.Tensor, comm = None):
     """Gathers the input tensor from all processes.
 
     Args:
@@ -36,11 +89,14 @@ def all_gather(x : torch.Tensor):
         raise RuntimeError("BMTrain is not initialized")
     
     assert x.is_cuda
-    return OpAllGather.apply(x)
+    return OpAllGather.apply(x, comm)
 
 class OpAllReduce(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, input : torch.Tensor, op : str):
+    def forward(ctx, input : torch.Tensor, op : str, comm : NCCLCommunicator = None):
+        if comm is None:
+            comm = config["comm"]
+        ctx.comm = comm
         if not input.is_contiguous():
             input = input.contiguous()
         if input.storage_offset() != 0 or input.storage().size() != input.numel():
@@ -51,10 +107,10 @@ def forward(ctx, input : torch.Tensor, op : str):
             input.storage(),
             output.storage(),
             op,
-            config['comm']
+            comm
         )
         ctx.op = op
-
+        
         if op in ["sum", "avg"]:
             pass
         elif op in ["max", "min"]:
@@ -66,15 +122,15 @@ def forward(ctx, input : torch.Tensor, op : str):
     @staticmethod
     def backward(ctx, grad_output):
         if ctx.op == "sum":
-            return grad_output, None
+            return grad_output, None, None
         elif ctx.op == "avg":
-            return grad_output / config['world_size'], None
+            return grad_output / commCount(ctx.comm), None, None
         elif ctx.op in ["max", "min"]:
-            return torch.masked_fill(grad_output, ctx.saved_tensors[0], 0), None
+            return torch.masked_fill(grad_output, ctx.saved_tensors[0], 0), None, None
         else:
-            return grad_output * ctx.saved_tensors[0], None
+            return grad_output * ctx.saved_tensors[0], None, None
 
-def all_reduce(x : torch.Tensor, op : str = "sum"):
+def all_reduce(x : torch.Tensor, op : str = "sum", comm = None):
     """Reduces the input tensor from all processes.
 
     Args:
@@ -89,7 +145,7 @@ def all_reduce(x : torch.Tensor, op : str = "sum"):
         raise RuntimeError("BMTrain is not initialized")
 
     assert x.is_cuda
-    return OpAllReduce.apply(x, op)
+    return OpAllReduce.apply(x, op, comm)
 
 
             
diff --git a/bmtrain/global_var.py b/bmtrain/global_var.py
index 67a54cf7..137fa9cd 100644
--- a/bmtrain/global_var.py
+++ b/bmtrain/global_var.py
@@ -6,13 +6,15 @@ class ConfigMap(TypedDict):
     world_size : int
     local_size : int
     zero_level : int
+    pipe_size : int
+    num_micro_batches : int
     calc_stream : torch.cuda.Stream
     load_stream : torch.cuda.Stream
     load_event : torch.cuda.Event
     barrier_stream : torch.cuda.Stream
     loss_scale_factor : float
     loss_scale_steps : int
-
+    topology : 'topology'
     gradient_inspect : bool
     initialized : bool
 
diff --git a/bmtrain/init.py b/bmtrain/init.py
index 42186923..dfc87e14 100644
--- a/bmtrain/init.py
+++ b/bmtrain/init.py
@@ -13,6 +13,8 @@ def init_distributed(
         loss_scale_factor : float = 2,
         loss_scale_steps : int = 1024,
         zero_level: int = 3,
+        pipe_size: int = -1,
+        num_micro_batches: int = None,
     ):
     """Initialize distributed training.
     This function will initialize the distributed training, set the random seed and global configurations.
@@ -40,7 +42,7 @@ def init_distributed(
 
     """
     torch.backends.cudnn.enabled = False
-    
+
     local_rank = int(os.environ["LOCAL_RANK"])
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
@@ -52,11 +54,11 @@ def init_distributed(
     )
     store, rank, world_size = next(rendezvous_iterator)
     store.set_timeout(timeout)
-
     store = dist.PrefixStore("bmtrain", store)
     torch.cuda.set_device(local_rank)
-
     config["initialized"] = True
+    config["pipe_size"] = pipe_size if pipe_size > 0 else 1
+    config["pipe_enabled"] = pipe_size > 0
     config["local_rank"] = local_rank
     config["local_size"] = local_size
     config["rank"] = rank
@@ -68,7 +70,8 @@ def init_distributed(
     config["zero_level"] = zero_level
     config["loss_scale_factor"] = loss_scale_factor if loss_scale_factor > 1 else 1 / loss_scale_factor
     config["loss_scale_steps"] = loss_scale_steps
-
+    config["topology"] = topology(config)
+    config["zero_rank"] = config["topology"].get_group_rank("zero") if pipe_size > 1 else config['rank']
     cpus_this_worker = None
     
     all_available_cpus = sorted(list(os.sched_getaffinity(0)))
@@ -97,7 +100,21 @@ def init_distributed(
     
     unique_id = bytes.fromhex(store.get("BMTRAIN_UNIQUE_ID").decode())
     config['comm'] = nccl.commInitRank(unique_id, world_size, rank)
-    
+    if config['pipe_enabled']:
+        config["micros"] = num_micro_batches if num_micro_batches else config["pipe_size"]
+        topo = config['topology']
+        if topo.stage_id == 0:
+            unique_id = nccl.getUniqueId()
+            store.set(f"PIPE_UNIQUE_ID{topo.pipe_idx}", unique_id.hex())
+        unique_id = bytes.fromhex(store.get(f"PIPE_UNIQUE_ID{topo.pipe_idx}").decode())
+        config ['pipe_comm'] = nccl.commInitRank(unique_id, pipe_size, topo.stage_id)
+        if topo.zero_id == 0:
+            unique_id = nccl.getUniqueId()
+            store.set(f"ZERO_UNIQUE_ID{topo.zero_idx}", unique_id.hex() )
+        unique_id = bytes.fromhex(store.get(f"ZERO_UNIQUE_ID{topo.zero_idx}").decode())
+        config ['zero_comm'] = nccl.commInitRank(unique_id, world_size//pipe_size, topo.zero_id)
+    else:
+        config['zero_comm'] = config['comm']
     for i in range(world_size):
         if i == rank:
             print_dict("Initialization", {
@@ -110,6 +127,38 @@ def init_distributed(
                 "cpus": cpus_this_worker 
             })
         synchronize()
+class topology:
+    def __init__(self,config):
+        # pipe_idx is the idx of the pipeline in the group
+        self.rank = config['rank']
+        pp_size = config["pipe_size"]
+        world_size = config["world_size"]
+        assert world_size % pp_size == 0, "The nums of GPUs must be divisible by the pipeline parallel size"
+
+        dp_size = world_size // pp_size
+        topo=torch.tensor(range(dp_size*pp_size),dtype=torch.int,device='cuda')
+        topo=topo.view(pp_size,dp_size)
+        self.pp_group=topo.transpose(0,1).reshape(-1,pp_size)
+        self.dp_group=topo
+        self.stage_id = (self.pp_group == self.rank).nonzero()[0,-1].item()
+        self.stages = config['pipe_size']
+        self.pipe_idx = (self.pp_group == self.rank).nonzero()[0, 0].item() # x axes
+        self.zero_id = self.pipe_idx
+        self.zero_idx = self.stage_id
+        self.next_rank = self.stage_id+1 if self.stage_id < config['pipe_size'] - 1 else -1
+        self.prev_rank = self.stage_id-1 if self.stage_id > 0 else -1
+        self.tails = self.pp_group[self.pipe_idx, self.stage_id:].tolist()
+        self.heads = self.pp_group[self.pipe_idx, :self.stage_id + 1].tolist()
+    def get_group_id(self,group_name):
+        if group_name == "pipe":
+            return self.pipe_idx
+        elif group_name == "zero":
+            return self.zero_idx
+    def get_group_rank(self,group_name):
+        if group_name == "pipe":
+            return self.stage_id
+        elif group_name == "zero":
+            return self.zero_id
 
 def is_initialized() -> bool:
     return config["initialized"]
diff --git a/bmtrain/inspect/model.py b/bmtrain/inspect/model.py
index 6bbe57c1..158c2de4 100644
--- a/bmtrain/inspect/model.py
+++ b/bmtrain/inspect/model.py
@@ -1,4 +1,6 @@
 import torch
+from ..store import broadcast_object
+from ..pipe_layer import PipelineTransformerBlockList
 from ..block_layer import CheckpointBlock
 from ..parameter import DistributedParameter
 from .. import nccl
@@ -30,6 +32,91 @@ def _gather_value(value : torch.Tensor, partition_size, origin_size):
 
     return output_tensor
 
+def inspect_pipeline_transformer_block_list(pipe_model: PipelineTransformerBlockList, param_name : str, _prefix : str = ''):
+    ret = []
+    for name, model in pipe_model._modules.items():
+        idx = int(name)
+        prefix = _prefix + name + '.'
+
+        # fast check
+        pass_fast_check = False
+        for param in model._param_info:
+            abs_name = prefix + param["name"]
+            if fnmatch.fnmatch(abs_name, param_name):
+                pass_fast_check = True
+                break
+        if not pass_fast_check:
+            continue
+
+        if idx in pipe_model.layer_ids:
+            _param_buffer = {}
+            _grad_buffer = {}
+            for kw, val in model._storage_info.items():
+                storage_type = model._storage_params[kw].storage_type()
+
+                _param_buffer[kw] = storage_type(val["partition_size"] * val['world_size'])
+                if model._storage_params[kw].grad is not None:
+                    _grad_buffer[kw] = storage_type(val["partition_size"] * val['world_size'])
+            
+            nccl.groupStart()
+            for kw, val in model._storage_info.items():
+                nccl.allGather(
+                    model._storage_params[kw].storage(),
+                    _param_buffer[kw],
+                    config["zero_comm"]
+                )
+                if model._storage_params[kw].grad is not None:
+                    nccl.allGather(
+                        model._storage_params[kw].grad.storage(),
+                        _grad_buffer[kw],
+                        config["zero_comm"]
+                    )
+
+            nccl.groupEnd()
+            for param in model._param_info:
+                abs_name = prefix + param["name"]
+                if fnmatch.fnmatch(abs_name, param_name):
+                    kw_name = param["kw_name"]
+                    dtype = _param_buffer[kw_name].dtype
+                    device = _param_buffer[kw_name].device
+                    offset = param["offset"]
+                    shape = param["shape"]
+                    p = torch.tensor([], dtype=dtype, device=device).set_(_param_buffer[kw_name], offset, shape)
+                    if kw_name in _grad_buffer:
+                        g = torch.tensor([], dtype=dtype, device=device).set_(_grad_buffer[kw_name], offset, shape)
+                        info = {
+                            "name": abs_name,
+                            "shape": tuple(shape),
+                            "std": p.std().cpu().item(),
+                            "mean": p.mean().cpu().item(),
+                            "grad_std": g.std().cpu().item(),
+                            "grad_mean": g.mean().cpu().item(),
+                            "max": p.max().cpu().item(),
+                            "min": p.min().cpu().item(),
+                        }
+                    else:
+                        info = {
+                            "name": abs_name,
+                            "shape": tuple(shape),
+                            "std": p.std().cpu().item(),
+                            "mean": p.mean().cpu().item(),
+                            "grad_std": None,
+                            "grad_mean": None,
+                            "max": p.max().cpu().item(),
+                            "min": p.min().cpu().item(),
+                        }
+                    broadcast_object(info, config["pipe_comm"], pipe_model.get_stage_by_layer_id(idx))
+                    ret.append(info)
+        else:
+            for param in model._param_info:
+                abs_name = prefix + param["name"]
+                if fnmatch.fnmatch(abs_name, param_name):
+                    info = broadcast_object({}, config["pipe_comm"], pipe_model.get_stage_by_layer_id(idx))
+                    ret.append(info)
+
+    return ret
+
+
 def inspect_checkpoint_block(model : CheckpointBlock, param_name : str, prefix : str = ''):
     # fast check
     pass_fast_check = False
@@ -121,7 +208,9 @@ def inspect_model(model : torch.nn.Module, param_name : str, prefix : str = ''):
         ...
 
     """
-    if isinstance(model, CheckpointBlock):
+    if isinstance(model, PipelineTransformerBlockList):
+        return inspect_pipeline_transformer_block_list(model, param_name, prefix)
+    elif isinstance(model, CheckpointBlock):
         return inspect_checkpoint_block(model, param_name, prefix)
     else:
         ret = []
diff --git a/bmtrain/inspect/tensor.py b/bmtrain/inspect/tensor.py
index 5d7030f3..7d56ad7f 100644
--- a/bmtrain/inspect/tensor.py
+++ b/bmtrain/inspect/tensor.py
@@ -3,6 +3,7 @@
 from .. import debug
 from .. import nccl
 from ..global_var import config
+from ..store import broadcast_object
 import math
 
 
@@ -16,62 +17,88 @@ def __init__(self):
         self._summary = []
     
     def _set_summary(self, summary):
-        self._summary = []
-        group_cnt = {}
-        for item in summary:
-            group = item["group"]
-            name = item["name"]
-            if group not in group_cnt:
-                group_cnt[group] = {}
-            if name not in group_cnt[group]:
-                group_cnt[group][name] = 0
-            group_cnt[group][name] += 1
-        
-        group_idx = {}
         for item in summary:
-            group = item["group"]
-            name = item["name"]
-            if group not in group_idx:
-                group_idx[group] = {}
-            if name not in group_idx[group]:
-                group_idx[group][name] = 0
-
-            group_name_prefix = f"{group}." if group is not None else ""
-            if group_cnt[group][name] > 1:
-                item["name"] = f"{group_name_prefix}{group_idx[group][name]}.{name}"
+            item['prefix'] = "" if item["group"] is None else f'{item["group"]}.'
+            
+        self._summary = []
+
+        kw_cnt = {}
+        i = 0
+        while i < len(summary):
+            item = summary[i]
+            if item["inside_pipe"] is not None:
+                assert item["inside_pipe"]["st"]
+                pipe_cnt = {}
+                j = i
+                while j < len(summary):
+                    item = summary[j]
+                    kw = f'{item["prefix"]}{item["name"]}'
+
+                    assert item["inside_pipe"] is not None
+                    stage_id = item["inside_pipe"]["stage_id"]
+                    stages = item["inside_pipe"]["stages"]
+                    st = item["inside_pipe"]["st"]
+                    ed = item["inside_pipe"]["ed"]
+
+                    if kw not in pipe_cnt:
+                        pipe_cnt[kw] = 0
+                    pipe_cnt[kw] += 1
+
+                    j += 1
+                    if ed:
+                        break
+
+                for stage in range(stages):
+                    if stage_id == stage:
+                        broadcast_object(pipe_cnt, config["pipe_comm"], src = stage)
+                        for k in range(i, j):
+                            item = summary[k]
+                            kw = f'{item["prefix"]}{item["name"]}'
+                            if kw not in kw_cnt:
+                                kw_cnt[kw] = 0
+                            item["name"] = f'{item["prefix"]}{kw_cnt[kw]}.{item["name"]}'
+                            item["shape"] = ((item["shape"][0] * config['micros'],) + item["shape"][1:])
+                            if item['tensor'].grad is not None:
+                                grad = torch.cat([summary[k+m*(j-i)]['tensor'].grad for m in range(config['micros'])], dim=0)
+                            else:
+                                grad = None
+                            item["tensor"] = torch.cat([summary[k+m*(j-i)]['tensor'] for m in range(config['micros'])], dim=0)
+                            item["tensor"].grad = grad
+                            self._summary.append(item)
+                            kw_cnt[kw] += 1
+                    else:
+                        cnt = broadcast_object({}, config["pipe_comm"], src = stage)
+                        for kw, val in cnt.items():
+                            if kw not in kw_cnt:
+                                kw_cnt[kw] = 0
+                            for _ in range(val):
+                                self._summary.append({
+                                    "name": f'{item["prefix"]}{kw_cnt[kw]}.{item["name"]}',
+                                    "group": None,
+                                    "requires_grad": None,
+                                    "min": None,
+                                    "max": None,
+                                    "mean": None,
+                                    "std": None,
+                                    "shape": None,
+                                    "grad_mean" : None,
+                                    "grad_std" : None,
+                                    "tensor": None,
+                                    "inside_pipe": {"stage_id": stage},
+                                })
+                                kw_cnt[kw] += 1
+
+                i = i + config['micros'] * (j - i)
             else:
-                item["name"] = f"{group_name_prefix}{name}"
-            if not item["requires_grad"]:
-                x = item["tensor"]
-                info = torch.empty(2, dtype=x.dtype, device=x.device)
-                info[0] = x.mean()
-                info[1] = x.var()
-                nccl.allReduce(
-                    info.storage(),
-                    info.storage(),
-                    "avg",
-                    config['comm']
-                )
-                x_mean = info[0].cpu().item()
-                x_std = math.sqrt(info[1].cpu().item())
-                item["mean"] = x_mean
-                item["std"] = x_std
+                kw = f'{item["prefix"]}{item["name"]}'
+                if kw not in kw_cnt:
+                    kw_cnt[kw] = 0
+                item["name"] = f'{item["prefix"]}{kw_cnt[kw]}.{item["name"]}'
+                self._summary.append(item)
+                kw_cnt[kw] += 1
+                i = i + 1
 
-                info[0] = x.max()
-                info[1] = -x.min()
-                nccl.allReduce(
-                    info.storage(),
-                    info.storage(),
-                    'max',
-                    config['comm']
-                )
-                x_max = info[0].cpu().item()
-                x_min = - info[1].cpu().item()
-                item["max"] = x_max
-                item["min"] = x_min
 
-            self._summary.append(item)
-            group_idx[group][name] += 1
     
     def get_summary(self):
         """Get the summary of the tensors recorded by `record_tensor`.
@@ -90,78 +117,71 @@ def get_summary(self):
         **Note:** This method must be called outside of the `with` block.
 
         """
-        nw_summary = []
+        ret = []
         for item in self._summary:
-            if item["requires_grad"] and item["tensor"].grad is not None:
-                x = item["tensor"]
-                info = torch.empty(4, dtype=x.dtype, device=x.device)
-                info[0] = x.mean()
-                info[1] = x.var()
-                info[2] = x.grad.mean()
-                info[3] = x.grad.var()
-                nccl.allReduce(
-                    info.storage(),
-                    info.storage(),
-                    "avg",
-                    config['comm']
-                )
-                x_mean = info[0].cpu().item()
-                x_std = math.sqrt(info[1].cpu().item())
-                grad_mean = info[2].cpu().item()
-                grad_std = math.sqrt(info[3].cpu().item())
-                
+            comm = config["zero_comm"] if item["inside_pipe"] else config["comm"]
+            if item["tensor"] is not None:
+                if not item["requires_grad"] or item["tensor"].grad is None:
+                    x = item["tensor"]
+                    info = torch.empty(2, dtype=x.dtype, device=x.device)
+                    info[0] = x.mean()
+                    info[1] = x.var()
+                    nccl.allReduce(
+                        info.storage(),
+                        info.storage(),
+                        "avg",
+                        comm
+                    )
+                    x_mean = info[0].cpu().item()
+                    x_std = math.sqrt(info[1].cpu().item())
+                    grad_mean = None
+                    grad_std = None
+                else:
+                    x = item["tensor"]
+                    info = torch.empty(4, dtype=x.dtype, device=x.device)
+                    info[0] = x.mean()
+                    info[1] = x.var()
+                    info[2] = x.grad.mean()
+                    info[3] = x.grad.var()
+                    nccl.allReduce(
+                        info.storage(),
+                        info.storage(),
+                        "avg",
+                        comm
+                    )
+                    x_mean = info[0].cpu().item()
+                    x_std = math.sqrt(info[1].cpu().item())
+                    grad_mean = info[2].cpu().item()
+                    grad_std = math.sqrt(info[3].cpu().item())
+
                 info[0] = x.max()
                 info[1] = -x.min()
                 nccl.allReduce(
                     info.storage(),
                     info.storage(),
                     'max',
-                    config['comm']
+                    comm
                 )
                 x_max = info[0].cpu().item()
                 x_min = - info[1].cpu().item()
 
-                nw_summary.append({
+                summary = {
                     "name": item["name"],
-                    "group": item["group"],
-                    "requires_grad" : False,
                     "min": x_min,
                     "max": x_max,
                     "mean": x_mean,
                     "std": x_std,
                     "shape": item["shape"],
                     "grad_mean" : grad_mean,
-                    "grad_std" : grad_std,
-                    "tensor": item["tensor"]
-                })
-            else:
-                nw_summary.append(item)
-        
-        ret = []
-        self._summary = nw_summary
-        for item in self._summary:
-            if item["requires_grad"]:
-                ret.append({
-                    "name": item["name"],
-                    "min": item["min"],
-                    "max": item["max"],
-                    "mean": item["mean"],
-                    "std": item["std"],
-                    "shape": item["shape"],
-                    "grad_mean" : None,
-                    "grad_std" : None
-                })
+                    "grad_std" : grad_std
+                }
+
+                if item["inside_pipe"] is not None:
+                    broadcast_object(summary, config["pipe_comm"], item["inside_pipe"]["stage_id"])
             else:
-                ret.append({
-                    "name": item["name"],
-                    "min": item["min"],
-                    "max": item["max"],
-                    "mean": item["mean"],
-                    "std": item["std"],
-                    "shape": item["shape"],
-                    "grad_mean" : item["grad_mean"],
-                    "grad_std" : item["grad_std"]
-                })
+                summary = broadcast_object({}, config["pipe_comm"], item["inside_pipe"]["stage_id"])
+
+            ret.append(summary)
         return ret
     
     def get_tensor(self, name : str, group : Optional[str] = None, index : Optional[int] = None) -> torch.Tensor:
@@ -268,5 +288,6 @@ def record_tensor(x : torch.Tensor, name : str, group = None, requires_grad = Tr
         "shape": x_shape,
         "grad_mean" : None,
         "grad_std" : None,
-        "tensor": x
+        "tensor": x,
+        "inside_pipe": None,
     })
diff --git a/bmtrain/nccl/__init__.py b/bmtrain/nccl/__init__.py
index c3f0ae78..bccae34b 100644
--- a/bmtrain/nccl/__init__.py
+++ b/bmtrain/nccl/__init__.py
@@ -38,7 +38,8 @@ def dtype2nccl(dtype : torch.dtype) -> int:
         torch.float32 : ncclFloat32,
         torch.float : ncclFloat,
         torch.float64 : ncclFloat64,
-        torch.double : ncclDouble
+        torch.double : ncclDouble,
+        torch.bool : ncclBool
     }
     if dtype not in MAP:
         raise TypeError("Unsupport dtype %s" % dtype)
@@ -83,9 +84,21 @@ def commDestroy(comm : NCCLCommunicator):
     """
     C.ncclCommDestroy(comm.ptr)
     comm._destroy_ptr()
+def commCount(comm : NCCLCommunicator):
+    """NCCL API: `ncclCommCount <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/comms.html#ncclcommcount>`_
 
+    Args:
+        comm (NCCLCommunicator): NCCL communicator.
+    """
+    return C.ncclCommCount(comm.ptr)
 ### collective
+def commRank(comm : NCCLCommunicator):
+    """NCCL API: `ncclCommUserRank <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/comms.html#ncclCommUserRank>`_
 
+    Args:
+        comm (NCCLCommunicator): NCCL communicator.
+    """
+    return C.ncclCommUserRank(comm.ptr)
 def allReduce(
         src : torch.storage._StorageBase,
         dst : torch.storage._StorageBase,
@@ -124,7 +137,45 @@ def allReduce(
         comm.ptr,
         torch.cuda.current_stream().cuda_stream
     )
+def send(src : torch.storage._StorageBase,
+         peer : int,
+         comm : NCCLCommunicator
+    ):
+    """NCCL API: `ncclsend <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclsend>`_
 
+        Args:
+            src (torch.storage._StorageBase): Source buffer.
+            peer (int): rank peer needs to call ncclRecv
+            comm (NCCLCommunicator): NCCL communicator.
+    """
+
+    sendbuff = src.data_ptr()
+    count = src.size()
+    datatype = dtype2nccl(src.dtype)
+    C.ncclSend(
+        sendbuff,
+        count,
+        datatype,
+        peer,
+        comm.ptr,
+        torch.cuda.current_stream().cuda_stream
+    )
+def recv(dst : torch.storage._StorageBase,
+         peer : int,
+         comm : NCCLCommunicator
+        ):
+    recvbuff = dst.data_ptr()
+    count = dst.size()
+    datatype = dtype2nccl(dst.dtype)
+    C.ncclRecv(
+        recvbuff,
+        count,
+        datatype,
+        peer,
+        comm.ptr,
+        torch.cuda.current_stream().cuda_stream
+    )
+    
 def broadcast(
         src : torch.storage._StorageBase,
         dst : torch.storage._StorageBase,
@@ -221,7 +272,6 @@ def allGather(
     recvbuff = dst.data_ptr()
     sendcount = src.size()
     datatype = dtype2nccl(src.dtype)
-
     assert dst.size() % sendcount == 0, "Buffer size not aligned"
     C.ncclAllGather(
         sendbuff, 
diff --git a/bmtrain/nccl/enums.py b/bmtrain/nccl/enums.py
index 5e388728..dc2adbd4 100644
--- a/bmtrain/nccl/enums.py
+++ b/bmtrain/nccl/enums.py
@@ -3,6 +3,7 @@
 
 ncclInt8       = 0
 ncclChar       = 0
+ncclBool       = 0
 ncclUint8      = 1
 ncclInt32      = 2
 ncclInt        = 2
@@ -16,7 +17,6 @@
 ncclFloat64    = 8
 ncclDouble     = 8
 
-
 ### ncclRedOp_t
 
 ncclSum        = 0
diff --git a/bmtrain/optim/adam.py b/bmtrain/optim/adam.py
index 06ee195e..713d3008 100644
--- a/bmtrain/optim/adam.py
+++ b/bmtrain/optim/adam.py
@@ -152,7 +152,7 @@ def loss_scale(self, loss : torch.Tensor) -> torch.Tensor:
         """
         Backward with loss scale.
         """
-        return loss * (self.scale / config['world_size'])
+        return loss * (self.scale * config['pipe_size'] / config['world_size'])
 
     def load_state_dict(self, state_dict: dict) -> None:
         r"""Loads the optimizer state.
diff --git a/bmtrain/optim/adam_offload.py b/bmtrain/optim/adam_offload.py
index 4f106b07..7ae2f706 100644
--- a/bmtrain/optim/adam_offload.py
+++ b/bmtrain/optim/adam_offload.py
@@ -182,7 +182,7 @@ def loss_scale(self, loss : torch.Tensor) -> torch.Tensor:
         """
         Backward with loss scale.
         """
-        return loss * (self.scale / config['world_size'])
+        return loss * (self.scale * config['pipe_size'] / config['world_size'])
 
     def load_state_dict(self, state_dict: dict) -> None:
         r"""Loads the optimizer state.
diff --git a/bmtrain/optim/clip_grad.py b/bmtrain/optim/clip_grad.py
index dcf5f739..92d9a48e 100644
--- a/bmtrain/optim/clip_grad.py
+++ b/bmtrain/optim/clip_grad.py
@@ -23,18 +23,23 @@ def clip_grad_norm(param_groups, max_norm, scale, norm_type=2, eps=1e-6):
         >>> bmt.optim.clip_grad_norm(optimizer.param_groups, max_norm=1.0, scale=optimizer.scale, norm_type=2)
 
     """
-    scale = scale / config['world_size']
-    parameters = [p for group in param_groups for p in group['params'] if p.grad is not None]
-
+    scale = scale * config['pipe_size'] / config['world_size']
+    grads = []
+    parameters = [p for group in param_groups for p in group['params']]
+    for p in parameters:
+        if p.grad is not None:
+            grads.append(p.grad.data)
+        else:
+            grads.append(torch.zeros_like(p.data))
     if norm_type == 'inf':
-        total_norm_cuda = max(p.grad.data.abs().max() for p in parameters).detach()
+        total_norm_cuda = max(g.data.abs().max() for g in grads).detach()
         nccl.allReduce(total_norm_cuda.storage(), total_norm_cuda.storage(), "max", config["comm"])
         total_norm = total_norm_cuda
     else:
         norm_type = float(norm_type)
         total_norm_cuda = torch.cuda.FloatTensor([0])
-        for index, p in enumerate(parameters):
-            param_norm = p.grad.data.float().norm(norm_type)
+        for index, g in enumerate(grads):
+            param_norm = g.data.float().norm(norm_type)
             total_norm_cuda += param_norm ** norm_type
         nccl.allReduce(total_norm_cuda.storage(), total_norm_cuda.storage(), "sum", config["comm"])
         total_norm = total_norm_cuda[0] ** (1. / norm_type)
@@ -43,5 +48,6 @@ def clip_grad_norm(param_groups, max_norm, scale, norm_type=2, eps=1e-6):
     clip_coef = float(max_norm * scale) / (total_norm + eps)
     if clip_coef < 1:
         for p in parameters:
-            p.grad.data.mul_(clip_coef)
+            if p.grad is not None:
+                p.grad.data.mul_(clip_coef)
     return total_norm / scale
\ No newline at end of file
diff --git a/bmtrain/param_init.py b/bmtrain/param_init.py
index 49f00346..c95f90f8 100644
--- a/bmtrain/param_init.py
+++ b/bmtrain/param_init.py
@@ -1,6 +1,5 @@
 from typing import Generator, Iterable, List, Tuple
 import torch
-
 from .block_layer import CheckpointBlock
 from .parameter import DistributedParameter
 from .global_var import config
@@ -23,6 +22,7 @@ def init_distributed_parameter(params : Iterable[torch.nn.Parameter]):
             param._init_method(tmp_tensor)
 
             # Pytorch 1.11 changed the API of storage.__getitem__
+            # use zero_rank to support pipeline
             torch.tensor([], dtype=param.dtype, device=param.device).set_(param.storage())[:] = \
                 torch.tensor([], dtype=param.dtype, device=param.device).set_(tmp_storage)[partition_size * config['rank'] : partition_size * (config['rank'] + 1)]
             # param.storage().copy_(tmp_storage[partition_size * config['rank'] : partition_size * (config['rank'] + 1)])
diff --git a/bmtrain/pipe_layer.py b/bmtrain/pipe_layer.py
new file mode 100644
index 00000000..a408c954
--- /dev/null
+++ b/bmtrain/pipe_layer.py
@@ -0,0 +1,484 @@
+from collections import OrderedDict
+import copy
+import torch
+import copy
+from typing import Dict, Iterable, Iterator, Tuple, Union, List
+import torch
+
+from .distributed import all_gather, broadcast, all_reduce, send_activations, recv_activations
+from .global_var import config
+from . import nccl
+from .checkpointing import ScopedTensorInspectorContext
+from . import debug
+from .block_layer import CheckpointBlockContext, CheckpointBlock, round_up, _get_param_kw
+
+class OpMicroForward(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, placeholder, self : 'PipelineTransformerBlockList', layers_dict, save_list, hidden_state, *args):
+        with PipeContext(self, hidden_state) as pipe_input:
+            hidden_state = pipe_input[0].detach()
+            tensors = [arg if torch.is_tensor(arg) else None for arg in args]
+            others = [arg if not torch.is_tensor(arg) else None for arg in args]
+            ctx.nontensor_inputs = others
+            ctx.self = self
+            ctx.save_list = copy.deepcopy(save_list)
+            ctx.num_save_needed = save_list[-1][1]+1
+            ctx.layers_dict = layers_dict
+            layer_inputs = []
+            layer_inspector = []
+            cuda_rng_state = []
+            with torch.no_grad():
+                for idx,layer_id in enumerate(self.layer_ids):
+                    if save_list[idx][0] == idx:
+                        layer_inputs.append(hidden_state.detach())
+                    cuda_rng_state.append( torch.cuda.get_rng_state() )
+                    if not ctx.layers_dict[idx]:
+                        block_ctx = CheckpointBlockContext(self._modules[str(layer_id)], ctx.layers_dict[idx], 1, pipe=True)
+                    else:
+                        block_ctx = CheckpointBlockContext(self._modules[str(layer_id)], ctx.layers_dict[idx], 2, pipe=True)
+                    # gather parameter on load stream
+                    block_ctx.enter()
+                    # call inner module directly
+                    with ScopedTensorInspectorContext() as inspector:
+                        hidden_state = self._modules[str(layer_id)]._module._call_impl(hidden_state, *args)
+                    for ith, it in enumerate(inspector.hidden_states):
+                        it["shape"] = ((it["shape"][0] // config['pipe_size'],) + it["shape"][1:])
+                        it["inside_pipe"] = {
+                            "stage_id": self.stage_id,
+                            "stages": self.stages,
+                            "st": (layer_id==self.layer_ids[0] and ith==0),
+                            "ed": (layer_id==self.layer_ids[-1] and ith==len(inspector.hidden_states)-1),
+                        }
+                        debug.append("_inspect_hidden_states", it)
+                    layer_inspector.append(inspector.hidden_states)
+                    block_ctx.exit()
+            
+            ctx.layer_inspector = layer_inspector
+            ctx.cuda_rng_state = cuda_rng_state
+
+            ctx.save_for_backward(*layer_inputs, *tensors)
+            pipe_input[0] = hidden_state
+        if self.return_hidden_states:
+            middle_hiddens = layer_inputs 
+            for mid in middle_hiddens:
+                mid.requires_grad_()
+            middle_hiddens = torch.stack(middle_hiddens, dim=0)
+            return pipe_input[0], middle_hiddens
+        else:
+            return pipe_input[0], None
+
+    @staticmethod
+    def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle : torch.Tensor):
+        def exit_prev(prev_ctx, prev_grad):
+            if prev_ctx is not None:
+                if prev_grad:
+                    with torch.enable_grad():
+                        prev_ctx.exit()
+                        config["load_stream"].record_event(config["load_event"])
+                else:
+                    with torch.no_grad():
+                        prev_ctx.exit()
+                        config["load_stream"].record_event(config["load_event"])
+        if not torch.autograd._is_checkpoint_valid():
+            raise RuntimeError(
+                "Checkpointing is not compatible with .grad() or when an `inputs` parameter"
+                " is passed to .backward(). Please use .backward() and do not pass its `inputs`"
+                " argument.")
+        all_inputs = []
+        input_requires_grad = []
+        
+        layer_inputs = ctx.saved_tensors[:ctx.num_save_needed]
+        save_args = ctx.saved_tensors[ctx.num_save_needed:]
+        for tensor, other in zip(save_args, ctx.nontensor_inputs):
+            if tensor is None:
+                all_inputs.append(other)
+                input_requires_grad.append(False)
+            else:
+                # detach for tensor inputs
+                input_requires_grad.append( tensor.requires_grad )
+                nw_tensor = tensor.detach()
+                nw_tensor.requires_grad = tensor.requires_grad
+                all_inputs.append(nw_tensor)
+        with PipeContext(ctx.self, grad_hidden_state, backward=True) as pipe_input:
+            grad_hidden_state = pipe_input[0]
+            with torch.random.fork_rng(devices=[torch.cuda.current_device()], enabled=True):
+                with torch.enable_grad():
+                    # overlap load and scatter here
+                    prev_ctx = None
+                    prev_grad = False
+                    for idx,layer_id in list(enumerate(ctx.self.layer_ids))[::-1]:
+                        torch.cuda.set_rng_state(ctx.cuda_rng_state[idx])
+                        ipt = layer_inputs[ctx.save_list[idx][1]].requires_grad_()
+                        block_ctx = CheckpointBlockContext(ctx.self._modules[str(layer_id)], ctx.layers_dict[idx], 2, pipe=True)
+                        block_ctx.enter()
+                        exit_prev(prev_ctx, prev_grad)
+                        prev_ctx = block_ctx
+                        prev_grad = True
+
+                        with ScopedTensorInspectorContext() as inspector:
+                            output = ctx.self._modules[str(layer_id)]._module._call_impl(ipt, *all_inputs)
+
+                        assert len(ctx.layer_inspector[idx]) == len(inspector.hidden_states), "Backward step changed"
+                        for j, it in enumerate(inspector.hidden_states):
+                            it["shape"] = ((it["shape"][0] // config['pipe_size'],) + it["shape"][1:])
+                            assert it["name"] == ctx.layer_inspector[idx][j]["name"], "Backward step changed"
+                            assert it["shape"] == ctx.layer_inspector[idx][j]["shape"], "Backward step changed"
+                            assert it["group"] == ctx.layer_inspector[idx][j]["group"], "Backward step changed"
+                            
+                            # change the tensor in placeholder
+                            ctx.layer_inspector[idx][j]["requires_grad"] = it["requires_grad"]
+                            ctx.layer_inspector[idx][j]["tensor"] = it["tensor"]
+                        torch.autograd.backward(
+                            [output],
+                            [grad_hidden_state]
+                        )
+                        grad_hidden_state = ipt.grad
+                        if grad_middle is not None:
+                            grad_hidden_state = grad_hidden_state + grad_middle[idx]
+                    exit_prev(prev_ctx, prev_grad)
+
+            pipe_input[0] = grad_hidden_state
+        grads = []
+        for inp, requires_grad in zip(all_inputs, input_requires_grad):
+            if requires_grad:
+                grads.append(inp.grad)
+            else:
+                grads.append(None)
+        return (None, None, None, None, pipe_input[0]) + tuple(grads)
+
+class OpPipeTransformerBlockList(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, placeholder, self : 'PipelineTransformerBlockList', save_list, hidden_state, *args):
+        num_micros = config["micros"]
+        ctx.self = self
+        ctx.num_micros = num_micros
+        layers_dict = [{} for _ in range(len(self))]
+        args_list = [[] for _ in range(num_micros)]
+        batch_related = args[-1]
+        batch_related_origin = [True if i in args[-1] else False for i in range(len(args[:-1]))]
+        batch_related_rule = []
+        args = args[:-1]
+        batch_size = hidden_state.shape[0]
+        assert (batch_size * config["pipe_size"]) % num_micros == 0, f'The batch size {(batch_size * config["pipe_size"])} must be divisible by the number of micro_batch {num_micros}'
+        input_requires_grad = []
+        with torch.enable_grad():
+            for arg in args:
+                if torch.is_tensor(arg):
+                    arg_all = all_gather(arg, config['pipe_comm'])
+                    if arg.shape[0] == batch_size:
+                        batch_related_rule.append(True)
+                        arg_all = arg_all.flatten(0, 1).chunk(num_micros, dim=0)
+                        arg_all = [tensor.detach().requires_grad_(arg.requires_grad) for tensor in arg_all]
+                    else:
+                        batch_related_rule.append(False)
+                        # assert num_micros % self.stages == 0, "batch unrelated only support num_micros % stages == 0"
+                        # arg_all = [arg_all[i // (num_micros // self.stages)].detach().requires_grad_(arg.requires_grad) for i in range(num_micros)]
+                        arg_all = [arg_all[0].detach().requires_grad_(arg.requires_grad) for i in range(num_micros)]
+                    input_requires_grad.append(arg.requires_grad)
+                else:
+                    batch_related_rule.append(False)
+                    arg_all = [arg for _ in range(num_micros)]
+                    input_requires_grad.append(False)
+                for i in range(num_micros):
+                    args_list[i].append(arg_all[i])
+            outputs = []
+            if self.return_hidden_states:
+                middles = []
+            hidden_state_list = all_gather(hidden_state, config["pipe_comm"]).flatten(0, 1).detach().requires_grad_()
+            ctx.hidden_state_list = hidden_state_list
+            hidden_state_list = hidden_state_list.chunk(num_micros, dim=0)
+            for micro_idx, (hidden_state, arg) in enumerate(zip(hidden_state_list, args_list)):
+                placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled())
+                output, middle = OpMicroForward.apply(placeholder, self, layers_dict, save_list, hidden_state, *arg)
+                outputs.append(output)
+                if self.return_hidden_states:
+                    middles.append(middle)
+        if len(batch_related) == 0:
+            ctx.batch_related = batch_related_rule
+        else:
+            ctx.batch_related = batch_related_origin
+        ctx.args_list = args_list
+        ctx.input_requires_grad = input_requires_grad
+        ctx.output_list = outputs
+        if self.return_hidden_states:
+            ctx.middle_list = middles
+
+        with torch.enable_grad():
+            last_hidden = torch.cat(outputs, dim=0)
+            last_hidden_shape = last_hidden.shape
+            last_hidden = broadcast(last_hidden, config["pipe_size"] - 1, config["pipe_comm"])
+            last_hidden = last_hidden.chunk(self.stages, dim=0)
+            last_hidden = last_hidden[self.stage_id].clone()
+
+        if self.return_hidden_states:
+            middle_hiddens = []
+            with torch.enable_grad():
+                for stage_id in range(self.stages):
+                    if self.stage_id == stage_id:
+                        middle_hidden = torch.cat(middles, dim=1) # [(layers, micro_batch, ...), ] -> (layers, full_batch, ...)
+                    else:
+                        middle_shape = (self.get_part_len_by_stage_id(stage_id),)+last_hidden_shape
+                        middle_hidden = torch.zeros(middle_shape, device=last_hidden.device, dtype=last_hidden.dtype)
+                    middle_hidden = broadcast(middle_hidden, stage_id, config["pipe_comm"])
+                    middle_hidden = middle_hidden.chunk(self.stages, dim=1)
+                    middle_hidden = middle_hidden[self.stage_id].clone()
+                    middle_hiddens.append(middle_hidden)
+                middle_hiddens = torch.cat(middle_hiddens, dim=0)
+            return last_hidden, middle_hiddens
+        else:
+            return last_hidden, None
+
+
+    @staticmethod
+    def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle : torch.Tensor):
+        ipt = ctx.hidden_state_list
+        args_list = ctx.args_list
+        input_requires_grad = ctx.input_requires_grad
+        grad_hidden_state_list = all_gather(grad_hidden_state, config["pipe_comm"]).flatten(start_dim=0, end_dim=1).chunk(ctx.num_micros, dim=0)
+        if ctx.self.return_hidden_states:
+            for stage_id in range(ctx.self.stages):
+                layer_range = ctx.self.get_range_by_stage_id(stage_id)
+                grad_middle_state = grad_middle[layer_range]
+                grad_middle_state = all_gather(grad_middle_state.transpose(0,1), config["pipe_comm"]).flatten(start_dim=0, end_dim=1).transpose(0, 1).chunk(ctx.num_micros, dim=1) # (layer, micro_batch, ...)
+                if ctx.self.stage_id == stage_id:
+                    grad_middle_state_list = grad_middle_state
+
+            for m in range(ctx.num_micros):
+                output = ctx.output_list[m]
+                middle = ctx.middle_list[m]
+                grad_hidden_state = grad_hidden_state_list[m]
+                grad_middle_state = grad_middle_state_list[m]
+                torch.autograd.backward(
+                    [output, middle],
+                    [grad_hidden_state, grad_middle_state],
+                )
+        else:
+            for m in range(ctx.num_micros):
+                output = ctx.output_list[m]
+                grad_hidden_state = grad_hidden_state_list[m]
+                torch.autograd.backward(
+                    [output],
+                    [grad_hidden_state],
+                )
+        grads = []
+        for idx,requires_grad in enumerate(input_requires_grad):
+            if requires_grad:
+                grad = torch.cat([args_list[m][idx].grad for m in range(ctx.num_micros)], dim=0)
+                grad = all_reduce(grad, "sum", config["pipe_comm"])
+                split_size = ctx.self.stages if ctx.batch_related[idx] else ctx.num_micros
+                grad = grad.chunk(split_size)
+                if ctx.batch_related[idx]:
+                    grads.append(grad[ctx.self.stage_id])
+                else:
+                    grads.append(grad[0])
+            else:
+                grads.append(None)
+        grad = broadcast(ipt.grad, 0, config["pipe_comm"]).chunk(ctx.self.stages)
+        grad = grad[ctx.self.stage_id]
+
+        return (None, None, None, grad) + tuple(grads) + (None,)
+
+class PipelineTransformerBlockList(torch.nn.Module):
+    r"""
+    TransformerBlockList is a list of CheckpointBlocks.
+
+    This is designed to reduce the communication overhead by overlapping the computation and reduce_scatter operation during backward pass.
+
+    It is similar to `torch.nn.ModuleList` but with the difference when calling .forward() and .backward().
+
+    Example:
+        >>> module_list = [ ... ]
+        >>> normal_module_list = torch.nn.ModuleList(module_list)
+        >>> transformer_module_list = TransformerBlockList(module_list)
+        >>> # Calling normal module list
+        >>> for layer in normal_module_list:
+        >>>     hidden_state = layer.forward(hidden_state, ...)
+        >>> # Calling transformer module list
+        >>> hidden_state = transformer_module_list(hidden_state, ...)
+
+    """
+    _modules: Dict[str, CheckpointBlock]
+
+    def __init__(self, modules: Iterable[CheckpointBlock]) -> None:
+        super().__init__()
+        
+        self._modules = {}
+        rank = config['rank']
+        topo = config['topology']
+        self.layer_ids = []
+        pipe_group = topo.pp_group
+        self.stages = topo.stages
+        self.stage_id = topo.stage_id
+        self.pipe_idx = topo.pipe_idx 
+        for idx, module in enumerate(modules):
+            if not isinstance(module, CheckpointBlock):
+                module = CheckpointBlock(module)
+            self._modules[str(idx)] = module
+
+        self.layer_ids = self.get_range_by_stage_id(self.stage_id)
+        self.partition_modules(self.layer_ids)
+        self.next_rank = pipe_group[self.pipe_idx, self.stage_id + 1] if self.stage_id < config['pipe_size'] - 1 else -1
+        self.prev_rank = pipe_group[self.pipe_idx, self.stage_id - 1] if self.stage_id > 0 else -1
+        # self.micro_batches = config['num_micro_batches']
+            
+        self.save_list = [(i, i) for i in range(len(self.layer_ids))]
+            
+    def __len__(self) -> int:
+        return len(self._modules)
+
+    def __iter__(self) -> Iterator[CheckpointBlock]:
+        return iter(self._modules.values())
+
+    def __getitem__(self, index: Union[int, str]) -> CheckpointBlock:
+        return self._modules[str(index)]
+
+    def forward(self, hidden_state, *args, batch_related=[], return_hidden_states=False):
+        self.return_hidden_states = return_hidden_states
+        placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled())
+        args = list(args)
+        args.append(batch_related)
+        hidden_state, middle_states = OpPipeTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args)
+        if return_hidden_states:
+            return hidden_state, middle_states
+        else:
+            return hidden_state
+
+    def get_range_by_stage_id(self, stage_id : int) -> List[int]:
+        part_lens = [0]+[self.get_part_len_by_stage_id(i) for i in range(stage_id+1)]
+        start = sum(part_lens[:stage_id+1])
+        end = start + part_lens[stage_id+1]
+        return range(start, end)
+
+    def get_part_len_by_stage_id(self, stage_id : int) -> int:
+        return len(self) // self.stages + (stage_id < (len(self) % self.stages))
+
+    def get_stage_by_layer_id(self, layer_id : int) -> int:
+        part_len = len(self) // self.stages
+        rest = len(self) % self.stages
+        if layer_id // (part_len + 1) < rest:
+            return layer_id // (part_len + 1)
+        else:
+            return rest + (layer_id - rest * (part_len+1)) // part_len
+
+    def partition_modules(self, idxs) -> None:
+        for i in range(len(self)):
+            contiguous_params = {}
+            for kw, val in self[i]._storage_info.items():
+                storage_type = val["storage_type"]
+                contiguous_params[kw] = storage_type(round_up(val["total"], config["world_size"] // config["pipe_size"]))
+                nccl.allGather(
+                    self[i]._storage_params[kw].storage(),
+                    contiguous_params[kw],
+                    config["comm"]
+                )
+
+            if i not in idxs:
+                for name, param in self[i]._module.named_parameters():
+                    param.data = torch.tensor([], dtype = param.dtype, device = param.device)
+                for kw, val in self[i]._storage_info.items():
+                    val["begin"] = self.stage_id
+                    val["end"] = self.stage_id + 1
+                    val["partition_size"] = 1
+                    val["total"] = val["world_size"]
+                    dtype = self[i]._storage_params[kw].dtype
+                    device = self[i]._storage_params[kw].device
+                    self[i]._storage_params[kw] = \
+                        torch.nn.Parameter(torch.tensor([0], dtype = dtype, device=device))
+            else:
+                for kw, val in self[i]._storage_info.items():
+                    storage_type = val["storage_type"]
+                    val["world_size"] = config["world_size"] // config["pipe_size"]
+                    partition_size = round_up(val["total"], val["world_size"]) // val["world_size"]
+                    val["partition_size"] = partition_size
+                    val["begin"] = config['zero_rank'] * partition_size
+                    val["end"] = (config['zero_rank'] + 1) * partition_size
+                    storage_param_buffer = storage_type(partition_size)
+                    dtype = storage_param_buffer.dtype
+                    device = storage_param_buffer.device
+                    self[i]._storage_params[kw] = torch.nn.Parameter(
+                        torch.tensor([], dtype=dtype, device=device).set_(storage_param_buffer)
+                    )
+                    if val["requires_grad"]:
+                        self[i]._storage_params[kw].requires_grad_(True)
+                    else:
+                        self[i]._storage_params[kw].requires_grad_(False)
+                ordered_parameters = list(self[i]._module.named_parameters())
+                for idx, named_param in enumerate(ordered_parameters):
+                    name, param = named_param
+                    param_info = self[i]._param_info[idx]
+                    kw_name = _get_param_kw(param)
+                    storage_info = self[i]._storage_info[kw_name]
+                    storage_st = storage_info["begin"]
+                    storage_end = storage_info["end"]
+                    param_st = param_info["offset"]
+                    param_end = param_st + param_info["size"]
+                    if not (param_st >= storage_end or param_end <= storage_st):
+                        # copy offset in parameter storage
+                        offset_st = max(storage_st - param_st, 0)
+                        offset_end = min(storage_end - param_st, param_info["size"])
+                        assert offset_st < offset_end
+                        to_offset_st = offset_st + param_st - storage_st
+                        to_offset_end = offset_end + param_st - storage_st
+                        d_dtype = self[i]._storage_params[kw_name].dtype
+                        d_device = self[i]._storage_params[kw_name].device
+                        param.data = torch.tensor([], dtype=param.dtype, device=param.device).set_(self[i]._storage_params[kw_name].storage(), to_offset_st, (to_offset_end - to_offset_st,))
+                        param_info["begin"] = to_offset_st
+                        param_info["end"] = (to_offset_end - to_offset_st,)
+                        param.data[:] = \
+                            torch.tensor([], dtype=d_dtype, device=d_device).set_(contiguous_params[kw], storage_st+to_offset_st, (to_offset_end - to_offset_st,))[:]
+                    else:
+                        param.data = torch.tensor([], dtype=param.dtype, device=param.device)
+            del contiguous_params
+    
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        for name, module in self._modules.items():
+            idx = int(name)
+            name = prefix + name + '.'
+            
+            dst = OrderedDict() # creates an temporary ordered dict
+            dst._metadata = OrderedDict()
+
+            if idx in self.layer_ids:
+                with torch.no_grad():
+                    with CheckpointBlockContext(module, pipe=True):
+                        module._module.state_dict(destination=dst, prefix=name, keep_vars=False)
+                if config["zero_rank"] == 0:
+                    if config["rank"] == 0:
+                        destination.update(dst)
+                    else:
+                        assert list(dst.keys()) == [name+n for n, parameter in module._module.named_parameters()]
+                        for key, tensor in dst.items():
+                            send_activations(tensor.cuda(), 0, config['pipe_comm'])
+            if config['rank'] == 0 and idx not in self.layer_ids:
+                for n, parameter in module._module.named_parameters():
+                    destination[name+n] = recv_activations(self.get_stage_by_layer_id(idx), config['pipe_comm'])
+
+class PipeContext:
+    def __init__(self, module, hidden_state, backward=False):
+        self.module = module
+        self.stage_id = module.stage_id
+        self.stages = module.stages
+        self.next_rank = module.next_rank
+        self.prev_rank = module.prev_rank
+        self.hidden_state = [hidden_state]
+        self.backward = backward
+        self.send_buffer = {}
+    def enter(self):
+        if self.backward:
+            if self.stage_id != self.stages -1:
+                self.hidden_state[0] = recv_activations(self.stage_id + 1, config['pipe_comm'])
+        else:
+            if self.stage_id != 0:
+                self.hidden_state[0] = recv_activations(self.stage_id - 1, config['pipe_comm'])
+        return self.hidden_state
+    def exit(self):
+        if self.backward:
+            if self.stage_id != 0:
+                send_activations(self.hidden_state[0], self.stage_id - 1, config['pipe_comm'])
+        else:
+            if self.stage_id != self.stages - 1:
+                send_activations(self.hidden_state[0], self.stage_id + 1, config['pipe_comm'])
+    def __enter__(self):
+        return self.enter()
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.exit()
\ No newline at end of file
diff --git a/bmtrain/store.py b/bmtrain/store.py
index bf1c00f7..c0914f9a 100644
--- a/bmtrain/store.py
+++ b/bmtrain/store.py
@@ -1,6 +1,8 @@
 from collections import OrderedDict
 from typing import Dict
 import torch
+
+from .pipe_layer import PipelineTransformerBlockList
 from .global_var import config
 from .block_layer import CheckpointBlock
 from . import nccl
@@ -24,15 +26,19 @@ def _save_to_rank0(model : torch.nn.Module, destination=None, prefix=''):
         destination = OrderedDict()
         destination._metadata = OrderedDict()
     destination._metadata[prefix[:-1]] = local_metadata = dict(version=model._version)
-    _save_to_state_dict(model, destination, prefix)
-    for name, module in model._modules.items():
-        if module is not None:
-            _save_to_rank0(module, destination, prefix + name + '.')
-    for hook in model._state_dict_hooks.values():
-        hook_result = hook(model, destination, prefix, local_metadata)
-        if hook_result is not None:
-            destination = hook_result
+    if not isinstance(model, PipelineTransformerBlockList):
+        _save_to_state_dict(model, destination, prefix)
+        for name, module in model._modules.items():
+            if module is not None:
+                _save_to_rank0(module, destination, prefix + name + '.')
+        for hook in model._state_dict_hooks.values():
+            hook_result = hook(model, destination, prefix, local_metadata)
+            if hook_result is not None:
+                destination = hook_result
+    else:
+        model._save_to_state_dict(destination, prefix, False)
     return destination
+        
 
 
 def save(model : torch.nn.Module, file_name : str):
@@ -66,9 +72,28 @@ def save(model : torch.nn.Module, file_name : str):
 
 _pickler = pickle.Pickler
 _unpickler = pickle.Unpickler
-
-def broadcast_object(obj):
-    if config['rank'] == 0:
+def allgather_object(obj, comm):
+    f = io.BytesIO()
+    _pickler(f).dump(obj)
+    byte_storage = torch.ByteStorage.from_buffer(f.getvalue())
+    # Do not replace `torch.ByteTensor` or `torch.LongTensor` with torch.tensor and specifying dtype.
+    # Otherwise, it will casue 100X slowdown.
+    # See:
+    byte_tensor = torch.ByteTensor(byte_storage).cuda()
+    all_bytes_tensors = torch.empty(byte_tensor.numel() * nccl.commCount(comm), dtype=torch.uint8, device="cuda")
+    nccl.allGather(
+        byte_tensor.storage(),
+        all_bytes_tensors.storage(),
+        comm
+    )
+    obj_list = []
+    for i in range(nccl.commCount(comm)):
+        buf = all_bytes_tensors[i*byte_tensor.numel():(i+1)*byte_tensor.numel()].cpu().numpy().tobytes()
+        obj = _unpickler(io.BytesIO(buf)).load()
+        obj_list.append(obj)
+    return obj_list
+def broadcast_object(obj, comm, src = 0):
+    if nccl.commRank(comm) == src:
         f = io.BytesIO()
         _pickler(f).dump(obj)
         byte_storage = torch.ByteStorage.from_buffer(f.getvalue())
@@ -81,30 +106,30 @@ def broadcast_object(obj):
         nccl.broadcast(
             local_size.storage(),
             local_size.storage(),
-            0,
-            config['comm']
+            src,
+            comm
         )
         nccl.broadcast(
             byte_tensor.storage(),
             byte_tensor.storage(),
-            0,
-            config['comm']
+            src,
+            comm
         )
     else:
         local_size = torch.LongTensor([0]).cuda()
         nccl.broadcast(
             local_size.storage(),
             local_size.storage(),
-            0,
-            config['comm']
+            src,
+            comm
         )
         byte_tensor_size = local_size[0].item()
         byte_tensor = torch.empty(int(byte_tensor_size), dtype=torch.uint8, device="cuda")
         nccl.broadcast(
             byte_tensor.storage(),
             byte_tensor.storage(),
-            0,
-            config['comm']
+            src,
+            comm
         )
         buf = byte_tensor.cpu().numpy().tobytes()
         obj = _unpickler(io.BytesIO(buf)).load()
@@ -114,7 +139,7 @@ def broadcast_object(obj):
 class DistributedStateDictWrapper(Mapping):
     def __init__(self, state_dict : Dict) -> None:
         self._state_dict = state_dict
-        self._metadata = broadcast_object(getattr(state_dict, "_metadata", None))
+        self._metadata = broadcast_object(getattr(state_dict, "_metadata", None), config["comm"])
     
     def __getitem__(self, key : str):
         tmp_shape = torch.zeros(32, device="cuda", dtype=torch.int32)
@@ -169,13 +194,13 @@ def copy(self):
         return self
 
     def __len__(self):
-        return broadcast_object(len(self._state_dict))
+        return broadcast_object(len(self._state_dict), config["comm"])
     
     def __contains__(self, key : str):
-        return broadcast_object(key in self._state_dict)
+        return broadcast_object(key in self._state_dict, config["comm"])
     
     def keys(self):
-        return broadcast_object(list(self._state_dict.keys()))
+        return broadcast_object(list(self._state_dict.keys()),config["comm"])
 
     def __iter__(self):
         # pytorch 1.12.0 updated the load_state_dict method, which needs the state_dict to be a `Mapping`.
diff --git a/csrc/include/nccl.h b/csrc/include/nccl.h
index fc6247b1..4bdda76a 100644
--- a/csrc/include/nccl.h
+++ b/csrc/include/nccl.h
@@ -351,6 +351,7 @@ ncclResult_t pncclGroupStart();
 ncclResult_t  ncclGroupEnd();
 ncclResult_t pncclGroupEnd();
 
+
 #ifdef __cplusplus
 } // end extern "C"
 #endif
diff --git a/csrc/nccl.cpp b/csrc/nccl.cpp
index 5a4a275a..2d0ac8de 100644
--- a/csrc/nccl.cpp
+++ b/csrc/nccl.cpp
@@ -129,7 +129,40 @@ void pyNCCLReduceScatter(
         reinterpret_cast<cudaStream_t>(stream)
     ));
 }
-
+void pyNCCLSend(
+    std::uintptr_t sendbuff,
+    size_t sendcount,
+    int data_type,
+    int peer,
+    std::uintptr_t comm,
+    std::uintptr_t stream
+) {
+    checkNCCLStatus(ncclSend(
+        reinterpret_cast<void*>(sendbuff),
+        sendcount,
+        static_cast<ncclDataType_t>(data_type),
+        peer,
+        reinterpret_cast<ncclComm_t>(comm),
+        reinterpret_cast<cudaStream_t>(stream)
+    ));
+}
+void pyNCCLRecv(
+    std::uintptr_t recvbuff,
+    size_t recvcount,
+    int data_type,
+    int peer,
+    std::uintptr_t comm,
+    std::uintptr_t stream
+) {
+    checkNCCLStatus(ncclRecv(
+        reinterpret_cast<void*>(recvbuff),
+        recvcount,
+        static_cast<ncclDataType_t>(data_type),
+        peer,
+        reinterpret_cast<ncclComm_t>(comm),
+        reinterpret_cast<cudaStream_t>(stream)
+    ));
+}
 void pyNCCLGroupStart() {
     checkNCCLStatus(ncclGroupStart());
 }
@@ -137,7 +170,20 @@ void pyNCCLGroupStart() {
 void pyNCCLGroupEnd() {
     checkNCCLStatus(ncclGroupEnd());
 }
-
+int pyNCCLCommCount(
+    std::uintptr_t comm
+){
+    int res;
+    checkNCCLStatus(ncclCommCount(reinterpret_cast<ncclComm_t>(comm),&res));
+    return res;
+}
+int pyNCCLCommUserRank(
+    std::uintptr_t comm
+){
+    int rank;
+    checkNCCLStatus(ncclCommUserRank(reinterpret_cast<ncclComm_t>(comm),&rank));
+    return rank;
+}
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("ncclGetUniqueId", &pyNCCLGetUniqueID, "nccl get unique ID");
     m.def("ncclCommInitRank", &pyNCCLCommInitRank, "nccl init rank");
@@ -149,4 +195,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("ncclReduceScatter", &pyNCCLReduceScatter, "nccl reduce scatter");
     m.def("ncclGroupStart", &pyNCCLGroupStart, "nccl group start");
     m.def("ncclGroupEnd", &pyNCCLGroupEnd, "nccl group end");
+    m.def("ncclSend",&pyNCCLSend,"nccl send");
+    m.def("ncclRecv",&pyNCCLRecv,"nccl recv");
+    m.def("ncclCommCount",&pyNCCLCommCount,"nccl comm count");
+    m.def("ncclCommUserRank",&pyNCCLCommUserRank,"nccl comm user rank");
 }
diff --git a/example/models/gpt.py b/example/models/gpt.py
index 65779e5c..78d77a7d 100644
--- a/example/models/gpt.py
+++ b/example/models/gpt.py
@@ -36,9 +36,10 @@ def forward(self,
         mask_2d = mask[:, None, :] & mask[:, :, None]   # (batch, seq_len, seq_len)
         mask_2d = mask_2d & (pos[:, None, :] >= pos[:, :, None])
 
-        input_emb = self.pos_emb(pos) + self.word_emb(input)
+        out = self.pos_emb(pos) + self.word_emb(input)
 
-        out = self.transformers(input_emb, mask_2d, None)
+        # for layer in self.transformers:
+        out = self.transformers(out, mask_2d, None)
         out = self.layernorm(out)
 
         logits = self.word_emb(out, projection=True)
diff --git a/example/run.sh b/example/run.sh
index c0f5a47e..542e5252 100644
--- a/example/run.sh
+++ b/example/run.sh
@@ -1 +1,3 @@
-torchrun --nnodes=1 --nproc_per_node=8 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost train.py
\ No newline at end of file
+export NCCL_P2P_DISABLE=1
+export CUDA_LAUNCH_BLOCKING=1
+torchrun --nnodes=1 --nproc_per_node=4 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost train.py
diff --git a/tests/test_middle_hidden.py b/tests/test_middle_hidden.py
new file mode 100644
index 00000000..ca0dc9b3
--- /dev/null
+++ b/tests/test_middle_hidden.py
@@ -0,0 +1,188 @@
+import bmtrain as bmt
+import random
+import torch
+from bmtrain import config
+from bmtrain.block_layer import CheckpointBlock, TransformerBlockList
+from bmtrain.pipe_layer import PipelineTransformerBlockList
+import torch.nn.functional as F
+
+class Linear(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, init_weight = None, init_bias = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out = {}
+        if init_weight:
+            self.weight = bmt.DistributedParameter(torch.tensor(init_weight, dtype=torch.float, device="cuda").reshape(out_features, in_features))
+        else:
+            self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=torch.float, device="cuda"), init_method=torch.nn.init.xavier_normal_)
+
+        if init_bias:
+            self.bias = bmt.DistributedParameter(torch.tensor(init_bias, dtype=torch.float, device="cuda").reshape(out_features,))
+        else:
+            self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=torch.float, device="cuda"), init_method=torch.nn.init.zeros_)
+    
+    def forward(self, input):
+        ret = F.linear(input, self.weight, self.bias)
+        return ret
+
+class Model_ZERO(torch.nn.Module):
+    def __init__(self, pre, ms, post) -> None:
+        super().__init__()
+        self.pre = pre
+        self.ms = TransformerBlockList([
+            CheckpointBlock(m)
+            for m in ms
+        ])
+        self.post = post
+    
+    def forward(self, x, return_hidden_states=False):
+        x = self.pre(x)
+        if return_hidden_states:
+            x, o = self.ms(x, return_hidden_states=return_hidden_states)
+            return self.post(x), o
+        else:
+            x = self.ms(x, return_hidden_states=return_hidden_states)
+            return self.post(x)
+
+class Model_PIPE(torch.nn.Module):
+    def __init__(self, pre, ms, post) -> None:
+        super().__init__()
+        self.pre = pre
+        self.ms = PipelineTransformerBlockList([
+            CheckpointBlock(m)
+            for m in ms
+        ])
+        self.post = post
+    
+    def forward(self, x, return_hidden_states=False):
+        x = self.pre(x)
+        if return_hidden_states:
+            x, o = self.ms(x, return_hidden_states=return_hidden_states)
+            return self.post(x), o
+        else:
+            x = self.ms(x, return_hidden_states=return_hidden_states)
+            return self.post(x)
+
+class Model_BLOCK(torch.nn.Module):
+    def __init__(self, pre, ms, post) -> None:
+        super().__init__()
+        self.pre = pre
+        self.ms = torch.nn.ModuleList([
+            CheckpointBlock(m)
+            for m in ms
+        ])
+        self.post = post
+    
+    def forward(self, x, return_hidden_states=False):
+        x = self.pre(x)
+        o = []
+        y = x
+        for m in self.ms:
+            o.append(y)
+            y = m(y)
+        if return_hidden_states:
+            return self.post(y), o
+        else:
+            return self.post(y)
+
+class Model_NORMAL(torch.nn.Module):
+    def __init__(self, pre, ms, post) -> None:
+        super().__init__()
+        self.pre = pre
+        self.ms = torch.nn.ModuleList(ms)
+        self.post = post
+    
+    def forward(self, x, return_hidden_states=False):
+        x = self.pre(x)
+        o = []
+        y = x
+        for m in self.ms:
+            o.append(y)
+            y = m(y)
+        if return_hidden_states:
+            return self.post(y), o
+        else:
+            return self.post(y)
+
+def manual_seed(seed=33):
+    torch.manual_seed(seed)
+    random.seed(seed)
+    try:
+        import numpy as np
+        np.random.seed(seed)
+    except ModuleNotFoundError:
+        pass
+
+def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_middle=False, mix_test=False):
+    manual_seed()
+
+    pre = Linear(dim, dim)
+    post = Linear(dim, dim)
+    ms = [Linear(dim, dim) for i in range(num_layer)]
+
+    inp = torch.randn((batch, seq_len, dim)).cuda()
+    last_weight = torch.randn((batch, seq_len, dim)).cuda()
+    middle_weight = [
+        torch.randn((batch, seq_len, dim)).cuda()
+        for i in range(len(ms))
+    ]
+
+    bmt.init_parameters(pre)
+    bmt.init_parameters(post)
+    for m in ms:
+        bmt.init_parameters(m)
+    m = cls(pre, [m for m in ms], post)
+
+    if only_last:
+        logits = m(inp)
+        loss = (logits * last_weight).sum()
+        loss.backward()
+        bmt.print_rank(f"========================{name}:only last========================")
+        bmt.print_rank(
+            bmt.inspect.format_summary(
+                bmt.inspect.inspect_model(m, '*')
+            )
+        )
+    if only_middle:
+        logits, hidden_states = m(inp, return_hidden_states=True)
+        loss = sum([
+            (hidden_state * middle_weight[i]).sum()
+            for i, hidden_state in enumerate(hidden_states)
+        ])
+        loss.backward()
+        bmt.print_rank(f"========================{name}:only middle========================")
+        bmt.print_rank(
+            bmt.inspect.format_summary(
+                bmt.inspect.inspect_model(m, '*')
+            )
+        )
+    if mix_test:
+        logits, hidden_states = m(inp, return_hidden_states=True)
+        loss = sum([
+            (hidden_state * middle_weight[i]).sum()
+            for i, hidden_state in enumerate(hidden_states)
+        ]) + (logits * last_weight).sum()
+        loss.backward()
+        bmt.print_rank(f"========================{name}:mix========================")
+        bmt.print_rank(
+            bmt.inspect.format_summary(
+                bmt.inspect.inspect_model(m, '*')
+            )
+        )
+
+def test(name, cls, num_layer=4, dim=4096, batch=32, seq_len=256):
+    sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_last=True)
+    bmt.synchronize()
+    sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_middle=True)
+    bmt.synchronize()
+    sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, mix_test=True)
+    bmt.synchronize()
+
+bmt.init_distributed(pipe_size=4)
+
+test("normal", Model_NORMAL)
+test("block", Model_BLOCK)
+test("zero", Model_ZERO)
+test("pipe", Model_PIPE)
\ No newline at end of file
diff --git a/tests/test_send_recv.py b/tests/test_send_recv.py
new file mode 100644
index 00000000..f969d661
--- /dev/null
+++ b/tests/test_send_recv.py
@@ -0,0 +1,27 @@
+import bmtrain as bmt
+import torch
+from bmtrain.global_var import config
+from bmtrain.pipe_comm import send_activations, recv_activations, gather_input
+from bmtrain import nccl
+from torch.distributed.distributed_c10d import recv
+from time import sleep
+def test_send_tensor():
+    bmt.init_distributed()
+    current_stream = torch.cuda.current_stream()
+    groups = [0,2]
+    rank = config['rank']
+def test_gather_input():
+    bmt.init_distributed(pipe_size=2)
+    if config['topology'].get_group_id("pipe") == 0:
+        a = torch.ones((2,1))
+    else:
+        a = torch.zeros((2,1))
+    res = gather_input(a.cuda(),config['pipe_comm'])
+    if config['topology'].get_group_rank("pipe") == 0:
+        print(res)
+        
+def main():
+    test_send_tensor()
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
