code: def pretrain(
      train_valid_test_dataset_provider,
      model_provider,
      model_type,
      forward_step_func,
      process_non_loss_data_func=None,
      extra_args_provider=None,
      args_defaults={},
  ):
      """Main training program.
  
      This function will run the followings in the order provided:
          1) initialize Megatron.
          2) setup model, optimizer and lr schedule using the model_provider.
          3) call train_val_test_data_provider to get train/val/test datasets.
          4) train the modle using the forward_step_func.
  
      Arguments:
          train_valid_test_dataset_provider: a function that takes the size of
              train/valid/test dataset and returns `train, valid, test` datasets.
          model_provider: a function that returns a vanilla version of the
              model. By vanilla we mean a simple model on cpu with no fp16 or ddp.
          model_type: an enum that specifies the type of model being trained.
          forward_step_func: a function that takes a `data iterator` and `model`,
              and returns a `loss` scalar with a dictionary with key:values being
              the info we would like to monitor during training, for example
              `lm-loss: value`. We also require that this function add
              `batch generator` to the timers class.
          process_non_loss_data_func: a function to post process outputs of the
              network. It can be used for dumping output tensors (e.g images) to
              tensorboard. It takes `collected data`(list of tensors),
              `current iteration index` and `tensorboard writer` as arguments.
          extra_args_provider: a function that takes a parser and adds arguments
              to it. It is used for programs to add their own arguments.
          args_defaults: a dictionary from argument-name to argument-value. It
              to set already parse arguments.
      """
  
      # Initalize and get arguments, timers, and Tensorboard writer.
      initialize_megatron(
          extra_args_provider=extra_args_provider,
          args_defaults=args_defaults,
      )
      # Set pytorch JIT layer fusion options and warmup JIT functions.
      set_jit_fusion_options()
  
      # Adjust the startup time so it reflects the largest value.
      # This will be closer to what scheduler will see (outside of
      # image ... launches.
      global _TRAIN_START_TIME
      start_time_tensor = torch.tensor(
          [_TRAIN_START_TIME], dtype=torch.double, device="cuda"
      )
      torch.distributed.all_reduce(
          start_time_tensor, op=torch.distributed.ReduceOp.MIN
      )
      _TRAIN_START_TIME = start_time_tensor.item()
      print_rank_0(
          "time to initialize megatron (seconds): {:.3f}".format(
              time.time() - _TRAIN_START_TIME
          )
      )
      print_datetime("after megatron is initialized")
  
      args = get_args()
      timers = get_timers()
  
      # Model, optimizer, and learning rate.
      timers("model-and-optimizer-setup", log_level=0).start(
          barrier=True
      )
      model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
          model_provider, model_type
      )
      timers("model-and-optimizer-setup").stop()
      print_datetime(
          "after model, optimizer, and learning rate "
          "scheduler are built"
      )
      config = get_model_config(model[0])
  
      # Data stuff.
      timers(
          "train/valid/test-data-iterators-setup", log_level=0
      ).start(barrier=True)
      if args.virtual_pipeline_model_parallel_size is not None:
          train_data_iterator = []
          valid_data_iterator = []
          test_data_iterator = []
          for i in range(len(model)):
              mpu.set_virtual_pipeline_model_parallel_rank(i)
              iterators = build_train_valid_test_data_iterators(
                  train_valid_test_dataset_provider
              )
              train_data_iterator.append(iterators[0])
              valid_data_iterator.append(iterators[1])
              test_data_iterator.append(iterators[2])
      else:
          (
              train_data_iterator,
              valid_data_iterator,
              test_data_iterator,
          ) = build_train_valid_test_data_iterators(
              train_valid_test_dataset_provider
          )
      timers("train/valid/test-data-iterators-setup").stop()
      print_datetime("after dataloaders are built")
  
      # Print setup timing.
      print_rank_0("done with setup ...")
      timers.log(
          [
              "model-and-optimizer-setup",
              "train/valid/test-data-iterators-setup",
          ],
          barrier=True,
      )
  
      if not args.skip_train:
          print_rank_0("training ...")
  
          if (
              args.dataloader_type == "cyclic"
              and args.retro_add_retriever
          ):
              args.train_iters = args.retro_cyclic_train_iters
              print_rank_0(
                  "retro cyclic train iters : %d" % args.train_iters
              )
  
          iteration = 0
          if args.do_train and args.train_iters > 0:
              iteration = train(
                  forward_step_func,
                  model,
                  optimizer,
                  opt_param_scheduler,
                  train_data_iterator,
                  valid_data_iterator,
                  process_non_loss_data_func,
                  config,
              )
  
          print_datetime("after training is done")
  
          if args.save and iteration != 0:
              save_checkpoint(
                  iteration, model, optimizer, opt_param_scheduler
              )
      else:
          print_rank_0("skipping training (--skip-train is on) ...")
  
          iteration = args.iteration
  
      if args.do_valid:
          prefix = f"iteration {iteration} on validation set"
          evaluate_and_print_results(
              prefix,
              forward_step_func,
              valid_data_iterator,
              model,
              iteration,
              process_non_loss_data_func,
              config,
              verbose=True,
              write_to_tensorboard=not args.skip_train,
          )
  
      if args.do_test:
          prefix = f"iteration {iteration} on test set"
          evaluate_and_print_results(
              prefix,
              forward_step_func,
              test_data_iterator,
              model,
              iteration,
              process_non_loss_data_func,
              config,
              verbose=True,
              write_to_tensorboard=not args.skip_train,
          )
  
imports: import gc
  from datetime import datetime
  import math
  import logging
  import sys
  from .log_handler import CustomHandler
  from .theoretical_memory_usage import report_theoretical_memory
  import time
  import torch
  from megatron import get_args
  from megatron import get_signal_handler
  from megatron import get_timers
  from megatron import get_tensorboard_writer
  from megatron import get_wandb_writer
  from megatron import get_current_global_batch_size
  from megatron import get_num_microbatches
  from megatron import is_last_rank
  from megatron import update_num_microbatches
  from megatron.core import mpu, tensor_parallel
  from megatron.core.utils import get_model_config
  from megatron import print_rank_0
  from megatron import print_rank_last
  from megatron.checkpointing import load_checkpoint
  from megatron.checkpointing import save_checkpoint
  from megatron.model import Float16Module
  from megatron.model import GPTModel
  from megatron.core.distributed import DistributedDataParallel
  from megatron.core.distributed import finalize_model_grads
  from megatron.core.enums import ModelType
  from megatron.optimizer import get_megatron_optimizer
  from megatron.initialize import initialize_megatron
  from megatron.initialize import write_args_to_tensorboard
  from megatron.initialize import set_jit_fusion_options
  from megatron.optimizer_param_scheduler import OptimizerParamScheduler
  from megatron.utils import check_adlr_autoresume_termination
  from megatron.utils import unwrap_model
  from megatron.data.data_samplers import build_pretraining_data_loader
  from megatron.utils import calc_params_l2_norm
  from megatron.core.pipeline_parallel import get_forward_backward_func
  from megatron.utils import report_memory
  from megatron.model.vision.knn_monitor import compute_feature_bank
label_desc: Use image node_id 3 for calling a global function with example usage: pretrain(train_valid_test_dataset_provider, model_provider, model_type, forward_step_func, process_non_loss_data_func, extra_args_provider, args_defaults) without return types