code: def _with_pipelining_forward_step(
      model,
      tokens,
      position_ids,
      attention_mask,
      inference_params,
      micro_batch_size,
  ):
      """No interleaving is supported."""
      sequence_length = tokens.size(1)
      batch_size = tokens.size(0)
  
      # Divide the batch dimension into micro batches.
      num_micro_batches, last_chunk = divmod(
          batch_size, micro_batch_size
      )
      if last_chunk > 0:
          num_micro_batches += 1
  
      # Preallocate memory for output logits.
      logits = None
      if mpu.is_pipeline_last_stage():
          args = get_args()
          logits = torch.empty(
              (batch_size, sequence_length, args.padded_vocab_size),
              dtype=torch.float32,
              device=torch.cuda.current_device(),
          )
  
      # Preallocate recv buffer.
      recv_buffer = _allocate_recv_buffer(
          micro_batch_size, sequence_length
      )
  
      for micro_batch_index in range(num_micro_batches):
          # Slice among the batch dimenion.
          start = micro_batch_index * micro_batch_size
          end = min(start + micro_batch_size, batch_size)
          this_micro_batch_size = end - start
          tokens2use = tokens[start:end, ...]
          position_ids2use = position_ids[start:end, ...]
  
          # Run a simple forward pass.
          if this_micro_batch_size != micro_batch_size:
              recv_buffer = None
          output = _forward_step_helper(
              model,
              tokens2use,
              position_ids2use,
              attention_mask,
              inference_params,
              recv_buffer=recv_buffer,
          )
  
          # Adjust the batch size offset to account for the micro-batch.
          inference_params.batch_size_offset += this_micro_batch_size
  
          # Copy logits.
          if mpu.is_pipeline_last_stage():
              logits[start:end, ...] = output
  
      # Once we are done with all the micro-batches, we can
      # adjust the sequence length offset.
      inference_params.sequence_len_offset += sequence_length
      # and reset the batch size offset
      inference_params.batch_size_offset = 0
  
      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 5 for calling a global function with example usage: _with_pipelining_forward_step(model, tokens, position_ids, attention_mask, inference_params, micro_batch_size) and returns: logits