class_docstr: GenerateSchedule is a class that handles the pipeline parallel inference.
  In our schedule, we place tie weight layer, embedding and lm_head in the same device to save space, so in
  this schedule, the out for each encoding progress is on rank0.
  
  Args:
      stage_manager (`PipelineStageManager`): Pipeline stage manager.
      mb_manager (`MicroBatchManager`): Micro batch manager.
      verbose (bool): Whether to verbose the information of the pipeline.
docstr: Forward one step of the pipeline
  
  Args:
      model (Module): Model to be run.
      data_iter (Iterable): Data iterator.
  
  Returns:
      Union[torch.Tensor, dict]: The intermediate output (dict) of the current stage. If it is the last stage, the output is the loss (Tensor).
docstr_tok: ["Forward","one","step","of","the","pipeline","Args",":","model","(","Module",")",":","Model","to","be","run",".","data_iter","(","Iterable",")",":","Data","iterator",".","Returns",":","Union","[","torch.Tensor",",","dict","]",":","The","intermediate","output","(","dict",")","of","the","current","stage",".","If","it","is","the","last","stage",",","the","output","is","the","loss","(","Tensor",")","."]
code: def generate_step_broadcast(
      self, model: Module, data_iter: Iterable
  ) -> Union[torch.Tensor, dict]:
      """
      Forward one step of the pipeline
  
      Args:
          model (Module): Model to be run.
          data_iter (Iterable): Data iterator.
  
      Returns:
          Union[torch.Tensor, dict]: The intermediate output (dict) of the current stage. If it is the last stage, the output is the loss (Tensor).
      """
      output_sequence = []
      self.load_batch(data_iter)
      model.eval()
  
      whole_timestamp = []
      # run by round
      for _ in range(self.round):
          self.timestamps = (
              [[] for _ in range(self.stage_manager.num_stages)]
              if self.verbose and self.stage_manager.is_first_stage()
              else None
          )
          while self.mb_manager.is_micro_batch_done() is False:
              inputs_dict = None
              new_token = None
              output_dict = None
  
              # First stage and in PREFILL phase, just load the inputs
              if (
                  self.stage_manager.is_first_stage()
                  and self.mb_manager.cur_state is Status.PREFILL
              ):
                  inputs_dict = self.load_micro_batch()
                  if (
                      self.verbose
                      and self.stage_manager.is_first_stage()
                  ):
                      torch.cuda.synchronize()
                      self.timestamps[self.mb_manager.idx].append(
                          time.time()
                      )
                  self.mb_manager.add_descrption(inputs_dict)
                  interval_inputs = {
                      "infer_state": self.mb_manager.cur_infer_state
                  }
                  output_dict = model_forward(
                      model, inputs_dict, interval_inputs
                  )
              # In GENERATE phase
              else:
                  # Get hidden_states from previous stage
                  hidden_states = self.comm.recv_forward()
                  if self.stage_manager.is_first_stage():
                      # First just generate a new token
                      assert (
                          hidden_states is not None
                      ), "When first stage in GENERATE phase, the hidden states should not be None"
                      interval_inputs = {
                          "hidden_states": hidden_states[
                              "hidden_states"
                          ],
                          "infer_state": self.mb_manager.cur_infer_state,
                      }
                      logits = model_forward(
                          model, None, interval_inputs
                      )
                      if (
                          self.verbose
                          and self.stage_manager.is_first_stage()
                      ):
                          torch.cuda.synchronize()
                          self.timestamps[self.mb_manager.idx].append(
                              time.time()
                          )
                      assert (
                          "logits" in logits
                      ), f"When first stage in GENERATE phase, the ouput should have attribute `logits`, but has {logits.keys()}"
                      new_token = self._get_token_id(logits["logits"])
                      self.mb_manager.step(new_token)
                      # If the current micro batch is not DONE, go through blocks
                      if self.mb_manager.cur_state in (
                          Status.GENERATE,
                          Status.COOLDOWN,
                      ):
                          inputs_dict = (
                              self._prepare_inputs_for_new_token(
                                  new_token
                              )
                          )
                          interval_inputs = {
                              "infer_state": self.mb_manager.cur_infer_state
                          }
                          output_dict = model_forward(
                              model, inputs_dict, interval_inputs
                          )
                  else:
                      assert (
                          hidden_states is not None
                      ), "When not first stage, the hidden states should not be None"
                      # inputs_dict = self._prepare_inputs_for_interval_stage()
                      inputs_dict = None
                      if self.mb_manager.cur_state is Status.PREFILL:
                          inputs_dict = self.load_micro_batch()
                          self.mb_manager.add_descrption(inputs_dict)
                      interval_inputs = {
                          "hidden_states": hidden_states[
                              "hidden_states"
                          ],
                          "infer_state": self.mb_manager.cur_infer_state,
                      }
                      output_dict = model_forward(
                          model, inputs_dict, interval_inputs
                      )
  
              # Current microbatch is not DONE, send hidden_state to next stage
              if (
                  not self.stage_manager.is_first_stage()
                  or self.mb_manager.cur_state
                  in (
                      Status.GENERATE,
                      Status.COOLDOWN,
                  )
              ):
                  self.comm.send_forward(
                      {"hidden_states": output_dict["hidden_states"]}
                  )
  
              self.mb_manager.next()
  
          # All microbatch in current round is DONE
          if self.stage_manager.is_first_stage():
              output_sequence.extend(
                  self.mb_manager.export_new_tokens()
              )
          self.mb_manager.clear()
          if self.verbose and self.stage_manager.is_first_stage():
              whole_timestamp.extend(self.timestamps)
  
      return output_sequence, whole_timestamp
  
imports: import time
  from functools import partial
  from typing import Any, Iterable, Optional, Union
  import torch
  import torch.cuda
  from torch.nn import Module
  from torch.utils._pytree import tree_map
  from colossalai.inference.engine.microbatch_manager import MicroBatchManager, Status
  from colossalai.pipeline.p2p import PipelineP2PCommunication
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.utils.device import get_current_device
  from ._utils import get_batch_size, get_micro_batch, model_forward, to_device
  from .base import PipelineSchedule
label_desc: Use image node_id 19 for calling the GenerateSchedule obj's underlying member method code with example usage: obj.generate_step_broadcast(model, data_iter) and returns: output_sequence,  whole_timestamp