class_docstr: A training engine hybrid pipeline, data, and model parallel training.
  
  This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule`
  is provided.
class_docstr_tok: ["A","training","engine","hybrid","pipeline",",","data",",","and","model","parallel","training",".","This","engine","is","created","by","``","deepspeed.initialize","(",")","``","when","a",":","class",":","`","PipelineModule","`","is","provided","."]
code: def _exec_recv_grads(self, buffer_id):
      if self.wall_clock_breakdown():
          self.timers(PIPE_RECV_GRAD_TIMER).start()
  
      outputs = self.pipe_buffers["outputs"][buffer_id]
      # XXX these shapes are hardcoded for Megatron
      # Restore partitioned output if it was partitioned and we are sending full gradients
      if self.is_pipe_partitioned and not self.is_grad_partitioned:
          if self.pipe_partition_grad_meta_cache is None:
              self.pipe_partition_grad_meta_cache = outputs[0].to("cpu")
          part_output = PartitionedTensor.from_meta(
              meta=self.pipe_partition_grad_meta_cache,
              local_part=outputs[1],
              group=self.grid.get_slice_parallel_group(),
          )
          outputs[0].data = part_output.full()
          outputs = (outputs[0], *outputs[2:])
          # save for backward
          self.pipe_buffers["outputs"][buffer_id] = outputs
  
      # Allocate gradient if necessary
      if self.grad_layer is None:
          if isinstance(outputs, torch.Tensor):
              s = list(outputs.size())
              self.grad_layer = self._allocate_buffer(
                  s, dtype=outputs.dtype, num_buffers=1
              )[0]
          else:
              # XXX This is a HACK
              # When we exchange activations/gradients, the two pipe stages
              # need to issue the send/recv with the same buffer sizes or
              # else there is a deadlock. The is_floating_point() filter is
              # used to avoid sending gradients for tensors that do not
              # produce gradients. When TP>1, we partition the first
              # activations/gradients across TP ranks to save communication
              # volume and memory. That partitioned tensor is represented as
              # two tensors: a 1/TPth chunk of the original data and also a
              # small LongTensor storing the metadata used to reconstruct on
              # the other side. When combined, the floating point filter also
              # filtered out the metadata tensor. This quick (hacky) fix just
              # branches on is_grad_partitioned so we don't filter out the
              # metadata tensor.
              if self.is_grad_partitioned:
                  sizes_and_dtypes = [
                      (list(t.size()), t.dtype) for t in outputs[:2]
                  ] + [
                      (list(t.size()), t.dtype)
                      for t in outputs[2:]
                      if t.is_floating_point()
                  ]
              else:
                  sizes_and_dtypes = [
                      (list(t.size()), t.dtype)
                      for t in outputs
                      if t.is_floating_point()
                  ]
              self.grad_layer = self._allocate_buffers(
                  sizes_and_dtypes, num_buffers=1
              )[0]
  
      if isinstance(self.grad_layer, torch.Tensor):
          p2p.recv(self.grad_layer, self.next_stage)
      else:
          assert isinstance(outputs, tuple)
          for idx, buffer in enumerate(self.grad_layer):
              # XXX GPT-2 hack
              if (
                  self.is_grad_partitioned
                  and idx == 0
                  and buffer.dtype != torch.long
              ):
                  buffer.data = torch.zeros(
                      buffer.size(),
                      dtype=torch.long,
                      device=self.device,
                  )
              p2p.recv(buffer, self.next_stage)
  
      if self.wall_clock_breakdown():
          self.timers(PIPE_RECV_GRAD_TIMER).stop()
  
imports: from types import MethodType
  import torch
  from deepspeed import comm
  from deepspeed.utils import logger
  from deepspeed.utils.timer import ThroughputTimer
  from deepspeed.accelerator import get_accelerator
  from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
  from ..engine import DeepSpeedEngine, MEMORY_OPT_ALLREDUCE_SIZE
  from deepspeed.utils.timer import FORWARD_MICRO_TIMER, FORWARD_GLOBAL_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_MICRO_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, STEP_MICRO_TIMER, STEP_GLOBAL_TIMER
  from ..utils import PartitionedTensor
  from ..dataloader import RepeatingLoader
  from ..zero.config import ZeroStageEnum
  from ..activation_checkpointing import checkpointing
  from .module import PipelineModule, PipelineError
  from .None import p2p
  from .None import schedule
label_desc: Use image node_id 32 for calling the PipelineEngine obj's underlying member method code with example usage: obj._exec_recv_grads(buffer_id) without return types