class_docstr: Megatron Vision Transformer Model.
class_docstr_tok: ["Megatron","Vision","Transformer","Model","."]
docstr: Our dataloaders produce a micro-batch and then we fetch
  a number of microbatches depending on the global batch size and model parallel size
  from the dataloader to produce a list of microbatches.
  Batch should be a list of microbatches and those microbatches should on CPU.
  Microbatches are then moved to GPU during the pipeline.
  The list of microbatches is then piped through the pipeline using Apex fwd/bwd functions.
docstr_tok: ["Our","dataloaders","produce","a","micro-batch","and","then","we","fetch","a","number","of","microbatches","depending","on","the","global","batch","size","and","model","parallel","size","from","the","dataloader","to","produce","a","list","of","microbatches",".","Batch","should","be","a","list","of","microbatches","and","those","microbatches","should","on","CPU",".","Microbatches","are","then","moved","to","GPU","during","the","pipeline",".","The","list","of","microbatches","is","then","piped","through","the","pipeline","using","Apex","fwd\/bwd","functions","."]
code: def training_step(self, dataloader_iter, batch_idx):
      """
      Our dataloaders produce a micro-batch and then we fetch
      a number of microbatches depending on the global batch size and model parallel size
      from the dataloader to produce a list of microbatches.
      Batch should be a list of microbatches and those microbatches should on CPU.
      Microbatches are then moved to GPU during the pipeline.
      The list of microbatches is then piped through the pipeline using Apex fwd/bwd functions.
      """
      # Initialize userbuffer communicators.
      if self.initialize_ub:
          self.initialize_ub_func()
  
      # we zero grads here because we also call backward in the megatron-core fwd/bwd functions
      self._optimizer.zero_grad()
  
      loss_mean, _ = self.fwd_bwd_step(
          dataloader_iter, batch_idx, False
      )
  
      # when using sequence parallelism, the sequence parallel layernorm grads must be all-reduced
      if self.cfg.get(
          "tensor_model_parallel_size", 1
      ) > 1 and self.cfg.get("sequence_parallel", False):
          self.allreduce_sequence_parallel_gradients()
  
      if self.with_distributed_adam:
          # KJJ - Added this block from megatron_gpt_model.  It says it's not necessary
          #  and it's not clear if the remaining "if not" logic is still needed.
          #  keeping it for now, but might need to delete one or both of these.
  
          # synchronize asynchronous grad reductions
          # note: not necessary, but reduces performance degradation
          # from multiple simultaneous NCCL calls
          self._optimizer._finish_bucket_grad_sync()
  
          # launch grad reductions
          # Note: grads in first pipeline stage have already been
          # reduced
          if not parallel_state.is_pipeline_first_stage():
              self.reduce_overlap_gradients()
      elif self.megatron_amp_O2:
          # # when using pipeline parallelism grads must be all-reduced after the pipeline (not asynchronously)
          self._optimizer.allreduce_main_grads()
      else:
          # async grad allreduce is not currently implemented for O1/autocasting mixed precision training
          # so we all-reduce gradients after the pipeline
          self.allreduce_gradients()  # @sangkug we think this is causing memory to blow up (hurts perf)
  
      ## logging
      # we can only log on one rank if it is rank zero so we broadcast from last rank
      # we can avoid this broadcast by updating the PTL log function to accept specific ranks
      torch.distributed.broadcast(loss_mean, get_last_rank())
  
      if self.cfg.precision in [16, "16", "16-mixed"]:
          loss_scale = self.trainer.precision_plugin.scaler._scale
          if loss_scale is not None:
              self.log("loss_scale", loss_scale, batch_size=1)
  
      self.log(
          "reduced_train_loss",
          loss_mean,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      lr = self._optimizer.param_groups[0]["lr"]
      self.log("lr", lr, rank_zero_only=True, batch_size=1)
      self.log(
          "global_step",
          self.trainer.global_step + 1,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      self.log(
          "consumed_samples",
          self.compute_consumed_samples(
              self.trainer.global_step + 1 - self.init_global_step
          ),
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
  
      return loss_mean
  
imports: import itertools
  from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module, MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group, get_all_params_for_weight_decay_optimization, get_linear_layer, get_params_for_weight_decay_optimization, init_method_normal, scaled_init_method_normal
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank, torch_dtype_from_precision
  from nemo.collections.vision.data.megatron.data_samplers import MegatronVisionPretrainingRandomSampler
  from nemo.collections.vision.data.megatron.vit_dataset import build_train_valid_datasets
  from nemo.collections.vision.modules.vit.vit_backbone import VitBackbone, VitMlpHead
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.utils import logging
label_desc: Use image node_id 9 for calling the MegatronVitClassificationModel obj's underlying member method code with example usage: obj.training_step(dataloader_iter, batch_idx) and returns: loss_mean