code: def _run_p2pops(
      tensor_send_prev: Union[torch.Tensor, None],
      tensor_send_next: Union[torch.Tensor, None],
      tensor_recv_prev: Union[torch.Tensor, None],
      tensor_recv_next: Union[torch.Tensor, None],
      async_comm: bool = False,
      overlap_p2p_comm: bool = False,
      batch_p2p_comm: bool = True,
  ):
      p2p_group = parallel_state.get_pipeline_model_parallel_group()
      default_group = parallel_state.get_model_parallel_group()
  
      need_to_sync = p2p_group.name() != default_group.name()
      reqs = []
      ops = []
  
      if batch_p2p_comm and p2p_group.name() == "nccl":
          if tensor_send_prev is not None:
              send_prev_op = torch.distributed.P2POp(
                  op=torch.distributed.isend,
                  tensor=tensor_send_prev,
                  peer=parallel_state.get_pipeline_model_parallel_prev_rank(),
                  group=p2p_group,
              )
              ops.append(send_prev_op)
          if tensor_recv_prev is not None:
              recv_prev_op = torch.distributed.P2POp(
                  op=torch.distributed.irecv,
                  tensor=tensor_recv_prev,
                  peer=parallel_state.get_pipeline_model_parallel_prev_rank(),
                  group=p2p_group,
              )
              ops.append(recv_prev_op)
          if tensor_send_next is not None:
              send_next_op = torch.distributed.P2POp(
                  op=torch.distributed.isend,
                  tensor=tensor_send_next,
                  peer=parallel_state.get_pipeline_model_parallel_next_rank(),
                  group=p2p_group,
              )
              ops.append(send_next_op)
          if tensor_recv_next is not None:
              recv_next_op = torch.distributed.P2POp(
                  op=torch.distributed.irecv,
                  tensor=tensor_recv_next,
                  peer=parallel_state.get_pipeline_model_parallel_next_rank(),
                  group=p2p_group,
              )
              ops.append(recv_next_op)
          if len(ops) > 0:
              # sync before communication if needed
              if need_to_sync:
                  torch.cuda.synchronize()
              reqs = torch.distributed.batch_isend_irecv(ops)
      else:
          # sync before communication if needed
          if need_to_sync and any(
              [
                  tensor_send_prev is not None,
                  tensor_recv_prev is not None,
                  tensor_send_next is not None,
                  tensor_recv_next is not None,
              ]
          ):
              torch.cuda.synchronize()
  
          if tensor_send_prev is not None:
              send_prev_req = torch.distributed.isend(
                  tensor=tensor_send_prev,
                  dst=parallel_state.get_pipeline_model_parallel_prev_rank(),
                  group=p2p_group,
              )
              reqs.append(send_prev_req)
          if tensor_recv_prev is not None:
              recv_prev_req = torch.distributed.irecv(
                  tensor=tensor_recv_prev,
                  src=parallel_state.get_pipeline_model_parallel_prev_rank(),
                  group=p2p_group,
              )
              reqs.append(recv_prev_req)
          if tensor_send_next is not None:
              send_next_req = torch.distributed.isend(
                  tensor=tensor_send_next,
                  dst=parallel_state.get_pipeline_model_parallel_next_rank(),
                  group=p2p_group,
              )
              reqs.append(send_next_req)
          if tensor_recv_next is not None:
              recv_next_op = torch.distributed.irecv(
                  tensor=tensor_recv_next,
                  src=parallel_state.get_pipeline_model_parallel_next_rank(),
                  group=p2p_group,
              )
              reqs.append(recv_next_op)
  
      if len(reqs) > 0:
          if overlap_p2p_comm:
              return (None, None, None, None, reqs)
  
          if async_comm:
              if len(ops) == 0 or len(reqs) == len(ops):
                  tensor_send_prev_req = (
                      None if tensor_send_prev is None else reqs.pop(0)
                  )
                  tensor_recv_prev_req = (
                      None if tensor_recv_prev is None else reqs.pop(0)
                  )
                  tensor_send_next_req = (
                      None if tensor_send_next is None else reqs.pop(0)
                  )
                  tensor_recv_next_req = (
                      None if tensor_recv_next is None else reqs.pop(0)
                  )
              elif len(reqs) == 1:
                  tensor_send_prev_req = (
                      None if tensor_send_prev is None else reqs[0]
                  )
                  tensor_recv_prev_req = (
                      None if tensor_recv_prev is None else reqs[0]
                  )
                  tensor_send_next_req = (
                      None if tensor_send_next is None else reqs[0]
                  )
                  tensor_recv_next_req = (
                      None if tensor_recv_next is None else reqs[0]
                  )
              else:
                  assert (
                      False
                  ), "failed to manage p2p requests and handles"
              return (
                  tensor_send_prev_req,
                  tensor_recv_prev_req,
                  tensor_send_next_req,
                  tensor_recv_next_req,
                  None,
              )
          else:
              for req in reqs:
                  req.wait()
              return (None, None, None, None, None)
      return (None, None, None, None, None)
  
imports: from functools import reduce
  import operator
  from typing import Union, Optional, Tuple
  import torch
  from apex.transformer import parallel_state
  from apex.transformer.log_util import get_transformer_logger
  from apex.transformer.utils import split_tensor_into_1d_equal_chunks
  from apex.transformer.utils import gather_split_1d_tensor
  from apex.transformer.pipeline_parallel.utils import Shape
  from apex.transformer.pipeline_parallel._timers import _Timers
label_desc: Use image node_id 1 for calling a global function with example usage: _run_p2pops(tensor_send_prev, tensor_send_next, tensor_recv_prev, tensor_recv_next, async_comm, overlap_p2p_comm, batch_p2p_comm) and returns: None,  None,  None,  None,  None, None,  None,  None,  None,  reqs, tensor_send_prev_req,  tensor_recv_prev_req,  tensor_send_next_req,  tensor_recv_next_req,  None, None,  None,  None,  None,  None