code: def check_backward(output_grad, rank, logger):
      dist.barrier()
      if gpc.is_last_rank(ParallelMode.PIPELINE):
          grad = output_grad.clone()
      else:
          grad = recv_backward(output_grad.shape)
          logger.info(
              "Rank {} received backward. Correct grad: {}".format(
                  rank, check_equal(grad, output_grad)
              )
          )
      if not gpc.is_first_rank(ParallelMode.PIPELINE):
          send_backward(grad)
          logger.info("Rank {} sent backward.".format(rank))
  
file_path: colossalai/tests/test_legacy/test_trainer/test_pipeline/test_p2p.py
imports: import pytest
  import torch
  import torch.distributed
  from colossalai.legacy.communication import recv_backward, recv_forward, send_backward, send_backward_recv_forward, send_forward, send_forward_recv_backward
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.initialize import launch
  from colossalai.logging import get_dist_logger
  from colossalai.testing import rerun_if_address_is_in_use, spawn
  from colossalai.utils import get_current_device
label_desc: Use image node_id 3 for calling a global function with example usage: check_backward(output_grad, rank, logger) without return types