code: def _no_pipelining_forward_step(
      model,
      tokens,
      position_ids,
      attention_mask,
      inference_params,
      recv_buffer=None,
  ):
      """If recv_buffer is none, we will allocate one on the fly."""
      # Run a simple forward pass.
      output_tensor = _forward_step_helper(
          model,
          tokens,
          position_ids,
          attention_mask,
          inference_params,
          recv_buffer=recv_buffer,
      )
      # Update the sequence length offset.
      inference_params.sequence_len_offset += tokens.size(1)
  
      logits = None
      if mpu.is_pipeline_last_stage():
          logits = output_tensor
  
      return logits
  
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 4 for calling a global function with example usage: _no_pipelining_forward_step(model, tokens, position_ids, attention_mask, inference_params, recv_buffer) and returns: logits