code: def main_training_process(cfg, setup):
      """This function controls the central training loop."""
      local_time = time.time()
      model = cramming.construct_model(cfg.arch, cfg.data.vocab_size)
      dataset, tokenizer = cramming.load_pretraining_corpus(
          cfg.data, cfg.impl
      )
      checkpoint_rendevous = os.path.join(
          cfg.base_dir, cfg.name, "intermediate_state.pth"
      )
      if cfg.impl.resume_run_after_preempt and os.path.isfile(
          checkpoint_rendevous
      ):
          try:
              metadata = torch.load(
                  checkpoint_rendevous, map_location=torch.device("cpu")
              )["metadata"]
              initial_step, elapsed_time = (
                  metadata["step"],
                  metadata["elapsed"],
              )
          except RuntimeError:
              log.info("Checkpoint file unreadable or corrupted.")
              os.remove(checkpoint_rendevous)
              initial_step, elapsed_time = 0, 0.0
      else:
          initial_step, elapsed_time = 0, 0.0
  
      model_engine, _, _, dataloader = cramming.load_backend(
          model,
          dataset,
          tokenizer,
          cfg.train,
          cfg.impl,
          elapsed_time,
          setup=setup,
      )
      if cfg.impl.resume_run_after_preempt and os.path.isfile(
          checkpoint_rendevous
      ):
          log.info(
              f"Loading intermediate checkpoint from previous run onto device {cfg.impl.local_rank}..."
          )
          model_engine.load_training_checkpoint(checkpoint_rendevous)
      model_engine.train(cfg.train.pretrain_in_train_mode)
      stats = defaultdict(list)
  
      # Start the clocks now:
      wallclock_timer = time.time() - elapsed_time
      train_time = time.time()
      training_allowed, no_recovery_necessary = True, True
      loss_vals = []
  
      # Launch training
      for step, batch in enumerate(dataloader, initial_step + 1):
          # Heavy lifting is moved to engines
          device_batch = model_engine.to_device(batch)
          loss = model_engine.step(device_batch)
          loss_vals.append(loss.detach())
  
          # Check stopping criteria
          if (
              check_deadline(wallclock_timer, cfg.budget)
              or step == cfg.train.steps
          ):
              training_allowed = False
              log.info("Reached deadline. Stopping training ...")
  
          time_progress = (
              (time.time() - wallclock_timer) / 3600 / cfg.budget
          )
          step_progress = step / cfg.train.steps
          progress = max(time_progress, step_progress)
  
          # Collect stats and print to console and upload to wandb
          if step % cfg.impl.print_loss_every_nth_step == 0:
              loss_vals, train_time = collect_stats(
                  step,
                  loss_vals,
                  train_time,
                  stats,
                  model_engine,
                  dataloader,
                  cfg,
                  progress,
              )
              if check_early_termination(
                  wallclock_timer,
                  stats["loss"][-1],
                  cfg.impl.early_termination,
              ):
                  training_allowed = False
                  log.info(
                      "Loss higher than allowed threshold. Stopping training early..."
                  )
  
          # Checkpointing is triggered from stopping criteria and normal intervals
          if (
              cfg.impl.save_intermediate_checkpoints
              and step % cfg.impl.save_every_nth_step == 0
          ):
              if (
                  loss.detach().isfinite()
                  and cramming.utils.is_main_process()
                  and not cfg.dryrun
              ):
                  model_engine.save_training_checkpoint(
                      checkpoint_rendevous,
                      metadata=dict(
                          step=step,
                          elapsed=time.time() - wallclock_timer,
                      ),
                  )
  
          if not loss.detach().isfinite():
              (
                  training_allowed,
                  no_recovery_necessary,
              ) = engage_troubleshooting(
                  model_engine,
                  step,
                  training_allowed,
                  no_recovery_necessary,
                  cfg,
              )
  
          communicate_flags(training_allowed, no_recovery_necessary)
  
          if (cfg.dryrun and step > 2) or not training_allowed:
              break
  
          if not no_recovery_necessary:  # synced across devices
              log.info(
                  f"Attempting reload of checkpoint on device {cfg.impl.local_rank}."
              )
              model_engine.load_training_checkpoint(
                  checkpoint_rendevous
              )
              no_recovery_necessary = True
  
      # Save to summary:
      cramming.utils.save_summary(
          "pretrain", cfg, stats, time.time() - local_time, setup
      )
      if cramming.utils.is_main_process():
          # Save final checkpoint? Might have to recover the latest checkpoint first
          if (
              not loss.detach().isfinite()
              and cfg.impl.save_intermediate_checkpoints
          ):
              model_engine.load_training_checkpoint(
                  checkpoint_rendevous
              )
              loss = torch.as_tensor(
                  16.0
              )  # fake value for model file name
          if loss.detach().isfinite():
              now = datetime.datetime.now()
              long_checkpoint_id = f"{''.join(cfg.arch.architectures)}_{now.strftime('%Y-%m-%d')}_{loss:2.4f}"
              model_engine.save_final_model(
                  os.path.join(cfg.base_dir, cfg.name),
                  long_checkpoint_id,
                  tokenizer,
                  cfg.arch,
                  cfg.dryrun,
              )
  
              if cfg.impl.push_to_huggingface_hub:
                  model_engine.push_to_hub(
                      tokenizer, cfg, dryrun=cfg.dryrun
                  )
      metrics = dict(
          num_params=sum([p.numel() for p in model.parameters()])
      )
      return metrics
  
file_path: bert-ultrafast/training/pretrain.py
imports: import torch
  import hydra
  import os
  import time
  import datetime
  import logging
  from collections import defaultdict
  import cramming
label_desc: Use image node_id 1 for calling a global function with example usage: main_training_process(cfg, setup) and returns: metrics