docstr: Runs interleaved schedule, with communication between pipeline stages.
docstr_tok: ["Runs","interleaved","schedule",",","with","communication","between","pipeline","stages","."]
code: def run_forward_backward(
      self,
      model_chunk: Union[ModuleList, Module],
      data_iter: Iterable,
      criterion: Callable[..., Any],
      optimizer: Optional[OptimizerWrapper] = None,
      return_loss: bool = False,
      return_outputs: bool = False,
  ) -> Dict:
      """
      Runs interleaved schedule, with communication between pipeline stages.
      """
      assert not self.forward_only
  
      self.load_batch(data_iter)
  
      num_microbatch = self.num_microbatch * self.num_model_chunks
      num_warmup_microbatch = (
          self.stage_manager.num_stages - self.stage_manager.stage - 1
      ) * 2
      num_warmup_microbatch += (
          self.num_model_chunks - 1
      ) * self.stage_manager.num_stages
      num_warmup_microbatch = min(num_warmup_microbatch, num_microbatch)
      num_microbatch_remaining = num_microbatch - num_warmup_microbatch
  
      # Input, output tensors only need to be saved when doing backward passes
      input_objs = [[] for _ in range(self.num_model_chunks)]
      output_objs = [[] for _ in range(self.num_model_chunks)]
  
      outputs = (
          []
          if return_outputs
          and self.stage_manager.is_last_stage(ignore_chunk=True)
          else None
      )
  
      accum_loss = None
      if return_loss and self.stage_manager.is_last_stage(
          ignore_chunk=True
      ):
          accum_loss = torch.scalar_tensor(
              0, device=get_current_device()
          )
  
      model_chunk_id = self.get_model_chunk_id(0, is_forward=True)
      input_obj = self.recv_forward(model_chunk_id)
      # Run warmup forward passes.
      for i in range(num_warmup_microbatch):
          last_iteration = i == num_warmup_microbatch - 1
          model_chunk_id = self.get_model_chunk_id(i, is_forward=True)
          output_obj = self.forward_step(
              model_chunk,
              model_chunk_id,
              input_obj,
              criterion,
              accum_loss,
              outputs,
          )
          input_objs[model_chunk_id].append(input_obj)
          output_objs[model_chunk_id].append(output_obj)
  
          if last_iteration and num_microbatch_remaining == 0:
              self.send_forward(model_chunk_id, output_obj)
          else:
              input_obj = self.send_forward_recv_forward(
                  model_chunk_id_send=model_chunk_id,
                  model_chunk_id_recv=self.get_model_chunk_id(
                      i + 1, is_forward=True
                  ),
                  output_tensor=output_obj,
                  send_prior=self.stage_manager.stage % 2 == 0,
              )
  
      if num_microbatch_remaining > 0:
          model_chunk_id = self.get_model_chunk_id(0, is_forward=False)
          output_obj_grad = self.recv_backward(model_chunk_id)
  
      # Run 1F1B in steady state.
      for i in range(num_microbatch_remaining):
          last_iteration = i == num_microbatch_remaining - 1
  
          model_chunk_id = self.get_model_chunk_id(
              i + num_warmup_microbatch, is_forward=True
          )
          output_obj = self.forward_step(
              model_chunk,
              model_chunk_id,
              input_obj,
              criterion,
              accum_loss,
              outputs,
          )
          # Add input_obj and output_obj to end of list.
          input_objs[model_chunk_id].append(input_obj)
          output_objs[model_chunk_id].append(output_obj)
  
          model_chunk_id = self.get_model_chunk_id(i, is_forward=False)
          # Pop output_obj and output_obj from the start of the list for the backward pass.
          _input_obj = input_objs[model_chunk_id].pop(0)
          _output_obj = output_objs[model_chunk_id].pop(0)
          input_obj_grad = self.backward_step(
              optimizer, _input_obj, _output_obj, output_obj_grad
          )
  
          # NOTE: perform 2x communication for forward and backward
          def send_forward_recv_backward():
              if (
                  last_iteration
                  and num_microbatch == num_microbatch_remaining
              ):
                  model_chunk_id = self.get_model_chunk_id(
                      i + num_warmup_microbatch, is_forward=True
                  )
                  self.send_forward(model_chunk_id, output_obj)
              else:
                  output_obj_grad = self.send_forward_recv_backward(
                      model_chunk_id_send=self.get_model_chunk_id(
                          i + num_warmup_microbatch, is_forward=True
                      ),
                      model_chunk_id_recv=self.get_model_chunk_id(
                          i + 1, is_forward=False
                      ),
                      output_tensor=output_obj,
                      send_prior_fallback=self.stage_manager.stage % 2
                      == 0,
                  )
                  return output_obj_grad
  
          def send_backward_recv_forward():
              if last_iteration:
                  model_chunk_id = self.get_model_chunk_id(
                      i, is_forward=False
                  )
                  self.send_backward(model_chunk_id, input_obj_grad)
              else:
                  input_obj = self.send_backward_recv_forward(
                      model_chunk_id_send=self.get_model_chunk_id(
                          i, is_forward=False
                      ),
                      model_chunk_id_recv=self.get_model_chunk_id(
                          i + num_warmup_microbatch + 1, is_forward=True
                      ),
                      input_tensor_grad=input_obj_grad,
                      send_prior_fallback=self.stage_manager.stage % 2
                      == 0
                      and i > 0,
                  )
                  return input_obj
  
          if self.stage_manager.stage % 2 == 0:
              output_obj_grad = send_forward_recv_backward()
              input_obj = send_backward_recv_forward()
          else:
              input_obj = send_backward_recv_forward()
              output_obj_grad = send_forward_recv_backward()
  
      if num_microbatch_remaining == 0:
          model_chunk_id = self.get_model_chunk_id(0, is_forward=False)
          output_obj_grad = self.recv_backward(model_chunk_id)
      # Run cooldown backward passes.
      for i in range(num_microbatch_remaining, num_microbatch):
          last_iteration = i == num_microbatch - 1
          model_chunk_id = self.get_model_chunk_id(i, is_forward=False)
          _input_obj = input_objs[model_chunk_id].pop(0)
          _output_obj = output_objs[model_chunk_id].pop(0)
          # output_obj_grad = self.recv_backward(model_chunk_id)
          input_obj_grad = self.backward_step(
              optimizer, _input_obj, _output_obj, output_obj_grad
          )
  
          if not last_iteration:
              output_obj_grad = self.send_backward_recv_backward(
                  model_chunk_id_send=self.get_model_chunk_id(
                      i, is_forward=False
                  ),
                  model_chunk_id_recv=self.get_model_chunk_id(
                      i + 1, is_forward=False
                  ),
                  input_tensor_grad=input_obj_grad,
                  send_prior=self.stage_manager.stage % 2 == 0
                  and i > num_microbatch_remaining,
              )
          else:
              model_chunk_id = self.get_model_chunk_id(
                  i, is_forward=False
              )
              self.send_backward(model_chunk_id, input_obj_grad)
  
      assert all(len(v) == 0 for v in input_objs) and all(
          len(v) == 0 for v in output_objs
      )
  
      if outputs is not None:
          outputs = merge_batch(outputs)
      return {"loss": accum_loss, "outputs": outputs}
  
imports: from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Union
  import torch
  import torch.cuda
  from torch.nn import Module, ModuleList
  from torch.utils._pytree import tree_map
  from colossalai.interface import OptimizerWrapper
  from colossalai.pipeline.p2p import PipelineP2PCommunication, create_send_metadata
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.utils.device import get_current_device
  from ._utils import detach, get_batch_size, get_micro_batch, merge_batch, model_forward, retain_grad, to_device
  from .base import PipelineSchedule
label_desc: Use image node_id 16 for calling the InterleavedSchedule obj's underlying member method code with example usage: obj.run_forward_backward(model_chunk, data_iter, criterion, optimizer, return_loss, return_outputs) and returns: dict, output_obj_grad, input_obj