code: def _communicate(
      tensor_send_next: Optional[torch.Tensor],
      tensor_send_prev: Optional[torch.Tensor],
      recv_prev: bool,
      recv_next: bool,
      tensor_shape: Optional[Shape] = None,
      override_scatter_gather_tensors_in_pipeline: bool = False,
      dtype_: Optional[torch.dtype] = None,
      *,
      scatter_gather_tensors_in_pipeline: bool = True,
      params_dtype: Optional[torch.dtype] = None,
      fp32_residual_connection: bool = False,
      async_comm: bool = False,
      sequence_parallel_enabled: bool = False,
      sync_batch_comm: bool = True,
      overlap_p2p_comm: bool = False,
      batch_p2p_comm: bool = True,
  ) -> Tuple[
      Union[torch.Tensor, FutureTensor, None],
      Union[torch.Tensor, FutureTensor, None],
  ]:
      """Base function for communication of tensors between stages.
  
  
      .. note::
          Reference https://github.com/NVIDIA/Megatron-LM/blob/cfd2e2160700b7f2c1bf35298ac14bc341f4c759/megatron/p2p_communication.py#L24-L159
  
      dtype logic: If none of ``dtype_``, ``params_dtype``, ``fp32_residual_connection`` is specified,
      torch.float32 is used.
  
      See https://github.com/NVIDIA/Megatron-LM/blob/d41696840ed0a7edb7e0499eb82a48ae112d9bb3/megatron/arguments.py#L145-L159
      for the details of arguments of ``dtype_``, ``params_dtype``, ``fp32_residual_connection``.
  
      Args:
          tensor_send_next: tensor to send to next rank (no tensor sent if set to None).
          tensor_send_prev: tensor to send to prev rank (no tensor sent if set to None).
          recv_prev: boolean for whether tensor should be received from previous rank.
          recv_next: boolean for whether tensor should be received from next rank.
          tensor_shape: optional, use when the input sequence contains less tokens than the default sequence length
          override_scatter_gather_tensors_in_pipeline:
              optional, this is used when tensor_shape is provided to override scatter gather tensors
          dtype_: This is used when tensor_shape is provided and what is the type of tensor_shape
  
      Keyword args:
          scatter_gather_tensors_in_pipeline: Optional. If :obj:`True`, use scatter/gather to optimize communication of tensors.
          params_dtype: Optional and legacy. Defaults to torch.float. If you manually call `.half()` or `.bfloat16()` on
              your model deliberately, pass this argument.
          fp32_residual_connection: Optional. If :obj:`True`, move residual connections to fp32.
          sequence_parallel_enabled: Set to :obj:`True` if sequence parallel is enabled.
              This argument is here for consistency with Megatron-LM.
              This argument has an effect on the communication optimization, not on tensor_shape update.
          sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication.
              To disable, https://github.com/pytorch/pytorch/pull/82450 would be required.
          overlap_p2p_comm: If :obj:`True`, returns cuda wait handles to scheduler instead of completing
              the communication within the p2p transfer API instance. The scheduler manages the communication completion
              to overlap with computation.
          batch_p2p_comm: If :obj:`True`, use the batched send and receive api to conduct the communication of
              a collection of send and receive operations between peer. If :obj:`False`, conduct each send and recv operation
              individually.
  
      Returns:
          tuple containing
  
          - tensor_recv_prev: `torch.Tensor` if `recv_prev` is :obj:`True`, `None` otherwise.
          - tensor_recv_next: `torch.Tensor` if `recv_next` is :obj:`True`, `None` otherwise.
      """
      if async_comm and sequence_parallel_enabled:
          import warnings  # NOQA
  
          class ExperimentalWarning(UserWarning):
              pass  # NOQA
  
          warnings.warn(
              "The combination of `async_comm` and `sequence_parallel_enabled` is not well tested.",
              ExperimentalWarning,
          )
      # Create placeholder tensors for receive in forward and backward directions if needed.
      tensor_recv_prev = None
      tensor_recv_next = None
      if tensor_shape is None:
          # In megatron, `tensor_shape` is set to `(args.seq_length, args.micro_batch_size, args.hidden_size)`
          raise RuntimeError(
              "`tensor_shape` must be specified. Common `tensor_shape` is `(seq_length, micro_batch_size, hidden_size)`"
          )
  
      tensor_parallel_size = (
          parallel_state.get_tensor_model_parallel_world_size()
      )
      override_scatter_gather_tensors_in_pipeline_ = False
      # TODO(mkozuki): Demystify hardcode False of `scatter_gather_tensors_in_pipeline` and add a testcase if possible.
      # NOTE(mkozuki): This is super strange and doesn't make sense to me. I have no idea what is happening here.
      # However, I can say that this hardcoding override is necessary for sequence parallel in nemo megatron to work.
      # I've not managed to reproduce the hang using standalone GPT with sequence parallel.
      # The hang in NeMo Megatron happens in the 3rd iteration, the last iteration of stead phase inside
      # forward_backward_pipelining_without_interleaving, pipeline parallel rank of 0 (tensor model parallel world
      # size of 2 and pipeline model parallel world size of 2). The commit then of APEX and NeMo were
      # https://github.com/NVIDIA/apex/pull/1396/commits/3060c98dd8ba42abf7702ea9d2cff0f39ea74f45 and
      # https://github.com/NVIDIA/NeMo/pull/4232/commits/1cb32dfca2ab9b20f53ebdb84476c34cb42f0205.
      # The PyTorch version was 1.13.0a0+git2d354cd, for what is worth.
      # Currently, indiscriminately this is set to `False`, which can lead to an unexpected performance regression
      # for non sequence parallel case.
      scatter_gather_tensors_in_pipeline = False
      if (
          scatter_gather_tensors_in_pipeline
          and not sequence_parallel_enabled
      ):
          tensor_chunk_size = int(reduce(operator.mul, tensor_shape, 1))
          if tensor_chunk_size % tensor_parallel_size == 0:
              tensor_chunk_shape = [
                  tensor_chunk_size // tensor_parallel_size
              ]
          else:
              tensor_chunk_shape = tensor_shape
              override_scatter_gather_tensors_in_pipeline_ = True
      else:
          tensor_chunk_shape = tensor_shape
  
      # The dtype logic below is copied from NVIDIA/Megatron-LM repo:
      # https://github.com/NVIDIA/Megatron-LM/blob/d41696840ed0a7edb7e0499eb82a48ae112d9bb3/megatron/p2p_communication.py#L74-L81
      dtype = params_dtype or torch.float
      if fp32_residual_connection:
          dtype = torch.float
      requires_grad = True
      if dtype_ is not None:
          dtype = dtype_
          # TODO(mkozuki): Figure out why this logic of requires_grad isn't working
          # when sequence_parallel_enabled=True. Otherwise, `x.retain_grad()` of
          # https://github.com/crcrpar/apex/blob/069832078a652b4bd8a99db84faf953a81415ab3/apex/transformer/pipeline_parallel/schedules/common.py#L360
          # fails.
          # requires_grad = False
  
      if recv_prev:
          tensor_recv_prev = torch.empty(
              tensor_chunk_shape,
              requires_grad=requires_grad,
              device=torch.cuda.current_device(),
              dtype=dtype,
          )
      if recv_next:
          tensor_recv_next = torch.empty(
              tensor_chunk_shape,
              requires_grad=requires_grad,
              device=torch.cuda.current_device(),
              dtype=dtype,
          )
  
      # Split tensor into smaller chunks if using scatter-gather optimization.
      scatter_gather_optimization_doable = (
          not override_scatter_gather_tensors_in_pipeline_
          and scatter_gather_tensors_in_pipeline
          and not sequence_parallel_enabled
      )
      if scatter_gather_optimization_doable:
          if tensor_send_next is not None:
              tensor_send_next = split_tensor_into_1d_equal_chunks(
                  tensor_send_next
              )
  
          if tensor_send_prev is not None:
              tensor_send_prev = split_tensor_into_1d_equal_chunks(
                  tensor_send_prev
              )
  
      # Send tensors in both the forward and backward directions as appropriate.
      (
          tensor_send_prev_req,
          tensor_recv_prev_req,
          tensor_send_next_req,
          tensor_recv_next_req,
          wait_handles,
      ) = _run_p2pops(
          tensor_send_prev,
          tensor_send_next,
          tensor_recv_prev,
          tensor_recv_next,
          async_comm,
          overlap_p2p_comm,
          batch_p2p_comm,
      )
  
      if async_comm:
          tensor_recv_prev_waitfunc = None
          tensor_recv_next_waitfunc = None
          # TODO: investigate whether this is necessary for correctness (ref: https://github.com/pytorch/pytorch/issues/38642)
          # see also: sync added for async_comm callbacks below in gather_recv_prev_wait and gather_recv_next_wait
          if tensor_recv_prev_req is not None:
  
              def tensor_recv_prev_wait():
                  tensor_recv_prev_req.wait()
                  torch.cuda.synchronize()
  
              tensor_recv_prev_waitfunc = tensor_recv_prev_wait
          if tensor_recv_next_req is not None:
  
              def tensor_recv_next_wait():
                  tensor_recv_next_req.wait()
                  torch.cuda.synchronize()
  
              tensor_recv_next_waitfunc = tensor_recv_next_wait
      else:
          if sync_batch_comm:
              # To protect against race condition when using batch_isend_irecv().
              torch.cuda.synchronize()
  
      # If using scatter-gather optimization, gather smaller chunks.
      if scatter_gather_optimization_doable:
          if not async_comm:
              if recv_prev:
                  tensor_recv_prev = (
                      gather_split_1d_tensor(tensor_recv_prev)
                      .view(tensor_shape)
                      .requires_grad_()
                  )
  
              if recv_next:
                  tensor_recv_next = (
                      gather_split_1d_tensor(tensor_recv_next)
                      .view(tensor_shape)
                      .requires_grad_()
                  )
          else:
  
              def gather_recv_prev_wait():
                  tensor_recv_prev_req.wait()
                  # From @Deepak's PR https://github.com/NVIDIA/Megatron-LM/commit/27fc468964064eeb33b703c9a0b2af938d80dd14
                  # A sync seems to be needed before gather otherwise losses jump around e.g., in run_gpt_minimal_test
                  torch.cuda.synchronize()
                  return (
                      gather_split_1d_tensor(tensor_recv_prev)
                      .view(tensor_shape)
                      .requires_grad_()
                  )
  
              def gather_recv_next_wait():
                  tensor_recv_next_req.wait()
                  torch.cuda.synchronize()
                  return (
                      gather_split_1d_tensor(tensor_recv_next)
                      .view(tensor_shape)
                      .requires_grad_()
                  )
  
              tensor_recv_prev_waitfunc = gather_recv_prev_wait
              tensor_recv_next_waitfunc = gather_recv_next_wait
      if async_comm:
          future_tensor_recv_prev = None
          future_tensor_recv_next = None
          if tensor_recv_prev is not None:
              future_tensor_recv_prev = FutureTensor(
                  tensor_recv_prev, tensor_recv_prev_waitfunc
              )
          if tensor_recv_next is not None:
              future_tensor_recv_next = FutureTensor(
                  tensor_recv_next, tensor_recv_next_waitfunc
              )
          return future_tensor_recv_prev, future_tensor_recv_next, None
      return tensor_recv_prev, tensor_recv_next, wait_handles
  
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 2 for calling a global function with example usage: _communicate(tensor_send_next, tensor_send_prev, recv_prev, recv_next, tensor_shape, override_scatter_gather_tensors_in_pipeline, dtype_) and returns: tensor_recv_prev,  tensor_recv_next,  wait_handles, future_tensor_recv_prev,  future_tensor_recv_next,  None, gather_split_1d_tensor, gather_split_1d_tensor