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: In this action, 1.prepare inputs for encoding for first stage. 2.do the forward to get hidden states 3.step to update
docstr_tok: ["In","this","action",",","1.prepare","inputs","for","encoding","for","first","stage",".","2.do","the","forward","to","get","hidden","states","3.step","to","update"]
code: def _head_encoding_action(self, model: Module):
      """
      In this action, 1.prepare inputs for encoding for first stage. 2.do the forward to get hidden states 3.step to update
      """
      new_token = self.action_interval_buffer.new_token
      assert (
          new_token is not None
      ), "When first stage in GENERATE phase, the new token should not be None"
      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)
  
      self.action_interval_buffer.hidden_states = output_dict[
          "hidden_states"
      ]
  
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 11 for calling the GenerateSchedule obj's underlying member method code with example usage: obj._head_encoding_action(model) without return types