code: def train(args):
      (
          _,
          num_workers,
          rank,
          local_rank,
          is_master_node,
          ctx_l,
      ) = init_comm(args.comm_backend, args.gpus)
      level = logging.DEBUG if args.verbose else logging.INFO
      logging_config(
          args.ckpt_dir,
          name="pretrain_bert_" + str(rank),  # avoid race
          level=level,
          console=(local_rank == 0),
      )
      logging.info(args)
      logging.debug("Random seed set to {}".format(args.seed))
      set_seed(args.seed)
      logging.info(
          "Training info: num_buckets: {}, "
          "num_workers: {}, rank: {}".format(
              args.num_buckets, num_workers, rank
          )
      )
      cfg, tokenizer, model = get_pretraining_model(
          args.model_name, ctx_l
      )
      if args.start_step:
          logging.info(
              "Restart training from {}".format(args.start_step)
          )
          parameters_option(
              args.start_step, model, args.ckpt_dir, "Loading", ctx_l
          )
      else:
          model.initialize(ctx=ctx_l)
      model.hybridize()
  
      if args.raw:
          get_dataset_fn = functools.partial(
              get_pretrain_data_text,
              max_seq_length=args.max_seq_length,
              short_seq_prob=args.short_seq_prob,
              masked_lm_prob=args.masked_lm_prob,
              max_predictions_per_seq=args.max_predictions_per_seq,
              whole_word_mask=args.whole_word_mask,
              random_next_sentence=args.random_next_sentence,
              tokenizer=tokenizer,
              circle_length=args.circle_length,
              repeat=args.repeat,
              dataset_cached=args.dataset_cached,
              num_max_dataset_cached=args.num_max_dataset_cached,
          )
      else:
          get_dataset_fn = get_pretrain_data_npz
  
      if args.data_dir:
          tmp = []
          for dataset in args.data_dir.split(","):
              names = os.listdir(dataset)
              for i in range(len(names)):
                  names[i] = os.path.join(dataset, names[i])
              tmp.append(",".join(names))
  
          args.data = ",".join(tmp)
  
      data_train = get_dataset_fn(
          args.data,
          args.batch_size,
          shuffle=True,
          num_buckets=args.num_buckets,
          vocab=tokenizer.vocab,
          num_parts=num_workers,
          part_idx=rank,
          num_dataset_workers=args.num_dataset_workers,
          num_batch_workers=args.num_batch_workers,
      )
  
      param_dict = model.collect_params()
      # Do not apply weight decay to all the LayerNorm and bias
      for _, v in model.collect_params(".*beta|.*gamma|.*bias").items():
          v.wd_mult = 0.0
      # Set grad_req if gradient accumulation is required
      params = [p for p in param_dict.values() if p.grad_req != "null"]
      num_accumulated = args.num_accumulated
      if num_accumulated > 1:
          logging.info(
              "Using gradient accumulation. Effective global batch size = {}".format(
                  num_accumulated
                  * args.batch_size
                  * len(ctx_l)
                  * num_workers
              )
          )
          for p in params:
              p.grad_req = "add"
  
      num_steps = args.num_steps
      warmup_steps = int(num_steps * args.warmup_ratio)
      log_interval = args.log_interval
      save_interval = args.ckpt_interval
      logging.info(
          "#Total Training Steps={}, Warmup Steps={}, Save Interval={}".format(
              num_steps, warmup_steps, save_interval
          )
      )
      optimizer_params = {"learning_rate": args.lr, "wd": args.wd}
      if args.optimizer == "adamw":
          optimizer_params.update(
              {
                  "beta1": 0.9,
                  "beta2": 0.999,
                  "epsilon": 1e-6,
                  "correct_bias": False,
              }
          )
      if args.use_amp:
          optimizer_params.update({"multi_precision": True})
      if args.comm_backend == "horovod":
          trainer = hvd.DistributedTrainer(
              param_dict, args.optimizer, optimizer_params
          )
      elif args.comm_backend == "byteps":
          trainer = bps.DistributedTrainer(
              param_dict, args.optimizer, optimizer_params
          )
      else:
          trainer = mx.gluon.Trainer(
              param_dict,
              args.optimizer,
              optimizer_params,
              update_on_kvstore=False,
          )
      if args.start_step:
          logging.info(
              "Restart training from {}".format(args.start_step)
          )
          states_option(
              args.start_step,
              trainer,
              args.ckpt_dir,
              local_rank,
              "Loading",
          )
  
      # backend specific implementation
      if args.use_amp:
          amp.init_trainer(trainer)
      if args.comm_backend == "byteps":
          trainer._init_params()
      if args.comm_backend == "horovod":
          # Horovod: fetch and broadcast parameters
          hvd.broadcast_parameters(param_dict, root_rank=0)
  
      # prepare the loss function
      nsp_loss_fn = mx.gluon.loss.SoftmaxCELoss()
      mlm_loss_fn = mx.gluon.loss.SoftmaxCELoss()
      nsp_loss_fn.hybridize()
      mlm_loss_fn.hybridize()
  
      mlm_metric = MaskedAccuracy()
      nsp_metric = MaskedAccuracy()
      mlm_metric.reset()
      nsp_metric.reset()
  
      step_num = args.start_step
      if args.phase2:
          step_num -= args.phase1_num_steps
  
      running_mlm_loss, running_nsp_loss = 0.0, 0.0
      running_num_tks = 0
  
      train_start_time = time.time()
      tic = time.time()
      # start training
      train_loop_dataloader = grouper(repeat(data_train), len(ctx_l))
      while step_num < num_steps:
          step_num += 1
          for _ in range(num_accumulated):
              sample_l = next(train_loop_dataloader)
              mlm_loss_l = []
              nsp_loss_l = []
              loss_l = []
              ns_label_list, ns_pred_list = [], []
              mask_label_list, mask_pred_list, mask_weight_list = (
                  [],
                  [],
                  [],
              )
              for sample, ctx in zip(sample_l, ctx_l):
                  # prepare data
                  (
                      input_id,
                      masked_id,
                      masked_position,
                      masked_weight,
                      next_sentence_label,
                      segment_id,
                      valid_length,
                  ) = sample
                  input_id = input_id.as_in_ctx(ctx)
                  masked_id = masked_id.as_in_ctx(ctx)
                  masked_position = masked_position.as_in_ctx(ctx)
                  masked_weight = masked_weight.as_in_ctx(ctx)
                  next_sentence_label = next_sentence_label.as_in_ctx(
                      ctx
                  )
                  segment_id = segment_id.as_in_ctx(ctx)
                  valid_length = valid_length.as_in_ctx(ctx)
  
                  with mx.autograd.record():
                      _, _, nsp_score, mlm_scores = model(
                          input_id,
                          segment_id,
                          valid_length,
                          masked_position,
                      )
                      denominator = (
                          (masked_weight.sum() + 1e-8)
                          * num_accumulated
                          * len(ctx_l)
                      )
                      mlm_scores_r = mx.npx.reshape(
                          mlm_scores, (-5, -1)
                      )
                      masked_id_r = masked_id.reshape((-1,))
                      mlm_loss = (
                          mlm_loss_fn(
                              mlm_scores_r,
                              masked_id_r,
                              masked_weight.reshape((-1, 1)),
                          ).sum()
                          / denominator
                      )
                      denominator = num_accumulated * len(ctx_l)
                      nsp_loss = (
                          nsp_loss_fn(
                              nsp_score, next_sentence_label
                          ).mean()
                          / denominator
                      )
                      mlm_loss_l.append(mlm_loss)
                      nsp_loss_l.append(nsp_loss)
                      loss_l.append(mlm_loss + nsp_loss)
                      mask_label_list.append(masked_id_r)
                      mask_pred_list.append(mlm_scores_r)
                      mask_weight_list.append(
                          masked_weight.reshape((-1,))
                      )
                      ns_label_list.append(next_sentence_label)
                      ns_pred_list.append(nsp_score)
  
                  running_num_tks += valid_length.sum().as_in_ctx(
                      mx.cpu()
                  )
              if args.use_amp:
                  with mx.autograd.record():
                      with amp.scale_loss(loss_l, trainer) as loss_l:
                          for loss in loss_l:
                              loss.backward()
                  norm_clip_mult = num_workers * trainer.amp_loss_scale
              else:
                  with mx.autograd.record():
                      for loss in loss_l:
                          loss.backward()
                  norm_clip_mult = num_workers
              running_mlm_loss += (
                  sum([ele.as_in_ctx(mx.cpu()) for ele in mlm_loss_l])
                  .asnumpy()
                  .item()
              )
              running_nsp_loss += (
                  sum([ele.as_in_ctx(mx.cpu()) for ele in nsp_loss_l])
                  .asnumpy()
                  .item()
              )
              mlm_metric.update(
                  mask_label_list, mask_pred_list, mask_weight_list
              )
              nsp_metric.update(ns_label_list, ns_pred_list)
          # update
          trainer.allreduce_grads()
          total_norm, ratio, is_finite = clip_grad_global_norm(
              params, args.max_grad_norm * norm_clip_mult
          )
          total_norm = total_norm / norm_clip_mult
          # update learning rate
          scheduled_lr = args.lr
          if step_num <= warmup_steps:
              scheduled_lr *= step_num / warmup_steps
          else:
              offset = (num_steps - step_num) / (
                  num_steps - warmup_steps
              )
              scheduled_lr *= max(offset, 0)
          trainer.set_learning_rate(scheduled_lr)
  
          if (
              args.comm_backend == "horovod"
              or args.comm_backend == "byteps"
          ):
              # Note that horovod.trainer._scale is default to num_workers,
              # thus trainer.update(1) will scale the gradients by 1./num_workers.
              # *num_workers* of Horovod is the number of GPUs.
              trainer.update(1, ignore_stale_grad=True)
          else:
              # gluon.trainer._scale is default to 1.
              # *num_workers* of Trainer is the number of machines.
              trainer.update(num_workers, ignore_stale_grad=True)
  
          if num_accumulated > 1:
              # set grad to zero for gradient accumulation
              model.zero_grad()
  
          # saving
          if step_num % save_interval == 0 or step_num >= num_steps:
              states_option(
                  step_num, trainer, args.ckpt_dir, local_rank, "Saving"
              )
              if local_rank == 0:
                  parameters_option(
                      step_num, model, args.ckpt_dir, "Saving"
                  )
          # logging
          if step_num % log_interval == 0:
              running_mlm_loss /= log_interval
              running_nsp_loss /= log_interval
              toc = time.time()
              logging.info(
                  "[step {}], Loss mlm/nsp={:.5f}/{:.3f}, Acc mlm/nsp={:.3f}/{:.3f}, "
                  " LR={:.7f}, grad_norm={:.4f}. Time cost={:.2f} s,"
                  " Throughput={:.1f}K tks/s, ETA={:.2f} h".format(
                      step_num,
                      running_mlm_loss,
                      running_nsp_loss,
                      mlm_metric.get()[1],
                      nsp_metric.get()[1],
                      trainer.learning_rate,
                      total_norm,
                      toc - tic,
                      running_num_tks.asnumpy().item()
                      / (toc - tic)
                      / 1000,
                      (num_steps - step_num)
                      / (step_num / (toc - train_start_time))
                      / 3600,
                  )
              )
              mlm_metric.reset()
              nsp_metric.reset()
              tic = time.time()
  
              running_mlm_loss = 0
              running_nsp_loss = 0
              running_num_tks = 0
  
      logging.info("Finish training step: %d", step_num)
  
      mx.npx.waitall()
      train_end_time = time.time()
      logging.info(
          "Train cost={:.1f} s".format(
              train_end_time - train_start_time
          )
      )
  
      if local_rank == 0:
          model_name = args.model_name.replace("google", "gluon")
          save_dir = os.path.join(args.ckpt_dir, model_name)
          final_save(model, save_dir, tokenizer, cfg)
  
file_path: gluon-nlp/scripts/pretraining/bert/run_pretraining.py
imports: import os
  import time
  import shutil
  import logging
  import argparse
  import functools
  import collections
  import mxnet
  import numpy
  from pretraining_utils import get_pretrain_data_npz, get_pretrain_data_text, MaskedAccuracy
  from gluonnlp.utils.misc import repeat, grouper, set_seed, init_comm, logging_config, naming_convention
  from gluonnlp.initializer import TruncNorm
  from gluonnlp.models.bert import BertModel, BertForPretrain, get_pretrained_bert
  from gluonnlp.utils.parameter import clip_grad_global_norm
label_desc: Use image node_id 6 for calling a global function with example usage: train(args) without return types