class_docstr: Model class for prompt-tuning or p-tuning a pretrained Megatron GPT model. 
  
  Prompt Tuning initalizes virtual prompt embeddings directly from a copy of
  certain token embeddings from the the pretrained GPT model's vocabulary
  and directly tunes these embedding weights. The token embeddings used in 
  initalization are specified by the user in the config file. The model can 
  be prompt-tuned for multiple tasks at once. virtual prompts are stored in a 
  prompt table and can be added or deleted without disrupting virtual prompts 
  for other tasks. 
  
  P-tuning initializes an LSTM encoder model that generates virtual prompt
  embeddings for every task. Each task shares the same encoder. After ptuning
  is compelete, the learned virtual prompts can be saved to the prompt table
  using add_ptuned_prompts_to_prompt_table(). Thus, if a user wants to add a 
  new virtual prompt via p-tuning, they do not need to retrain on all previous 
  tasks. This gives p-tuning the same task flexiblity as prompt-tuning.
class_docstr_tok: ["Model","class","for","prompt-tuning","or","p-tuning","a","pretrained","Megatron","GPT","model",".","Prompt","Tuning","initalizes","virtual","prompt","embeddings","directly","from","a","copy","of","certain","token","embeddings","from","the","the","pretrained","GPT","model","'s","vocabulary","and","directly","tunes","these","embedding","weights",".","The","token","embeddings","used","in","initalization","are","specified","by","the","user","in","the","config","file",".","The","model","can","be","prompt-tuned","for","multiple","tasks","at","once",".","virtual","prompts","are","stored","in","a","prompt","table","and","can","be","added","or","deleted","without","disrupting","virtual","prompts","for","other","tasks",".","P-tuning","initializes","an","LSTM","encoder","model","that","generates","virtual","prompt","embeddings","for","every","task",".","Each","task","shares","the","same","encoder",".","After","ptuning","is","compelete",",","the","learned","virtual","prompts","can","be","saved","to","the","prompt","table","using","add_ptuned_prompts_to_prompt_table","(",")",".","Thus",",","if","a","user","wants","to","add","a","new","virtual","prompt","via","p-tuning",",","they","do","not","need","to","retrain","on","all","previous","tasks",".","This","gives","p-tuning","the","same","task","flexiblity","as","prompt-tuning","."]
docstr: Dataloader produces a global batch which is turned into an iterator of microbatches.
  The iterator of microbatches is then piped through the pipeline using Core's fwd/bwd functions.
docstr_tok: ["Dataloader","produces","a","global","batch","which","is","turned","into","an","iterator","of","microbatches",".","The","iterator","of","microbatches","is","then","piped","through","the","pipeline","using","Core","'s","fwd\/bwd","functions","."]
code: def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only):
      """
      Dataloader produces a global batch which is turned into an iterator of microbatches.
      The iterator of microbatches is then piped through the pipeline using Core's fwd/bwd functions.
      """
      # Get seq length of batch
      batch = next(dataloader_iter)
      _, seq_length = batch[0].shape
      data_iter = get_iterator_k_split(batch, get_num_microbatches())
  
      fwd_bwd_function = get_forward_backward_func()
  
      losses_reduced_per_micro_batch = fwd_bwd_function(
          forward_step_func=self.get_forward_output_and_loss_func(),
          data_iterator=data_iter,
          model=[self],
          num_microbatches=get_num_microbatches(),
          forward_only=forward_only,
          seq_length=seq_length,
          micro_batch_size=get_micro_batch_size(),
      )
  
      # only the last stages of the pipeline return losses
      if losses_reduced_per_micro_batch:
          # average loss across micro batches
          loss_tensors_list = [
              loss_reduced["avg"]
              for loss_reduced in losses_reduced_per_micro_batch
          ]
          loss_tensor = torch.concat(loss_tensors_list)
          loss_mean = loss_tensor.mean()
      else:
          # we're not on the last pipeline stage so no losses
          loss_mean = torch.tensor(0.0).cuda()
  
      return loss_mean
  
imports: import itertools
  import os
  from functools import partial
  from typing import Any, List, Optional, Union
  import torch
  from omegaconf.dictconfig import DictConfig
  from omegaconf.omegaconf import open_dict
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
  from nemo.collections.nlp.data.language_modeling.megatron.gpt_prompt_learning_dataset import GPTPromptLearningDataset
  from nemo.collections.nlp.metrics.prompt_learning_metrics import AccuracyScore, BLEUScore, ROUGEScores
  from nemo.collections.nlp.models.language_modeling.megatron_base_prompt_learning_model import MegatronBasePromptLearningModel
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
  from nemo.collections.nlp.modules.common import VirtualPromptPlaceholderToken, VirtualPromptSource, VirtualPromptStyle
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_iterator_k_split
  from nemo.collections.nlp.modules.common.text_generation_utils import get_default_length_params, get_default_sampling_params, megatron_gpt_generate
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, SamplingParam
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler, NLPSaveRestoreConnector
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 5 for calling the MegatronGPTPromptLearningModel obj's underlying member method code with example usage: obj.fwd_bwd_step(dataloader_iter, batch_idx, forward_only) and returns: loss_mean