code: def _forward_step_helper(
      model,
      tokens,
      position_ids,
      attention_mask,
      inference_params,
      recv_buffer=None,
  ):
      """Single forward step. Update the allocate memory flag so
      only the first time the memory is allocated."""
      batch_size = tokens.size(0)
      sequence_length = tokens.size(1)
      if recv_buffer is None:
          recv_buffer = _allocate_recv_buffer(
              batch_size, sequence_length
          )
  
      # Receive from previous stage.
      recv_from_prev_pipeline_rank_(recv_buffer)
  
      # Forward pass through the model.
      model.set_input_tensor(recv_buffer)
      output_tensor = model(
          tokens,
          position_ids,
          attention_mask,
          inference_params=inference_params,
      )
  
      # Send output to the next stage.
      send_to_next_pipeline_rank(output_tensor)
  
      return output_tensor
  
file_path: megatron-lm/megatron/text_generation/forward_step.py
imports: from collections.abc import Iterable
  import torch
  from megatron import get_args
  from megatron.core import mpu, InferenceParams
  from .communication import send_to_next_pipeline_rank, recv_from_prev_pipeline_rank_
label_desc: Use image node_id 3 for calling a global function with example usage: _forward_step_helper(model, tokens, position_ids, attention_mask, inference_params, recv_buffer) and returns: output_tensor