code: def train(args):
      use_amp = args.dtype == "float16"
      (
          store,
          num_workers,
          rank,
          local_rank,
          is_master_node,
          ctx_l,
      ) = init_comm(args.comm_backend, args.gpus)
      setup_logging(args, local_rank)
      cfg, tokenizer, qa_net, use_segmentation = get_network(
          args.model_name,
          ctx_l,
          args.classifier_dropout,
          args.param_checkpoint,
          args.backbone_path,
      )
  
      logging.info("Prepare training data")
      train_features = get_squad_features(
          args, tokenizer, segment="train"
      )
      dataset_processor = SquadDatasetProcessor(
          tokenizer=tokenizer,
          doc_stride=args.doc_stride,
          max_seq_length=args.max_seq_length,
          max_query_length=args.max_query_length,
      )
      logging.info("Processing the Training data:")
      (
          train_dataset,
          num_answer_mismatch,
          num_unreliable,
      ) = dataset_processor.get_train(
          train_features, skip_unreliable=True
      )
      logging.info(
          "Done! #Unreliable Span={} / #Mismatched Answer={} / #Total={}".format(
              num_unreliable, num_answer_mismatch, len(train_features)
          )
      )
  
      # Get dataset statistics
      num_impossible = 0
      for sample in train_dataset:
          num_impossible += sample.is_impossible
      logging.info(
          "Before Chunking, #Train/Is Impossible = {}/{}".format(
              len(train_features),
              sum([ele.is_impossible for ele in train_features]),
          )
      )
      logging.info(
          "After Chunking, #Train Sample/Is Impossible = {}/{}".format(
              len(train_dataset), num_impossible
          )
      )
  
      # Shuffle the dataset using a fixed seed across all workers
      rs = np.random.RandomState(args.pre_shuffle_seed)
      rs.shuffle(train_dataset)
      sampler = SplitSampler(
          len(train_dataset),
          num_parts=num_workers,
          part_index=rank,
          even_size=True,
      )
      train_dataloader = mx.gluon.data.DataLoader(
          train_dataset,
          batchify_fn=dataset_processor.BatchifyFunction,
          batch_size=args.batch_size,
          num_workers=0,
          sampler=sampler,
      )
      if "electra" in args.model_name:
          # Froze parameters, does not work for albert model since parameters in all layers are shared
          if args.untunable_depth > 0:
              qa_net.backbone.frozen_params(args.untunable_depth)
          if args.layerwise_decay > 0:
              qa_net.backbone.apply_layerwise_decay(
                  args.layerwise_decay
              )
  
      logging.info("Creating distributed trainer...")
      # Collect differentiable parameters
      param_dict = deduplicate_param_dict(qa_net.collect_params())
      # Do not apply weight decay to all the LayerNorm and bias
      for _, v in qa_net.collect_params(
          ".*beta|.*gamma|.*bias"
      ).items():
          v.wd_mult = 0.0
      params = [p for p in param_dict.values() if p.grad_req != "null"]
      # Set grad_req if gradient accumulation is required
      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"
      # backend specific implementation
      if args.comm_backend == "horovod":
          # Horovod: fetch and broadcast parameters
          hvd.broadcast_parameters(param_dict, root_rank=0)
  
      epoch_size = (len(train_dataloader) + len(ctx_l) - 1) // len(
          ctx_l
      )
      if args.num_train_steps is not None:
          num_train_steps = args.num_train_steps
      else:
          num_train_steps = int(
              args.epochs * epoch_size / args.num_accumulated
          )
      if args.warmup_steps is not None:
          warmup_steps = args.warmup_steps
      else:
          warmup_steps = int(num_train_steps * args.warmup_ratio)
      assert (
          warmup_steps is not None
      ), "Must specify either warmup_steps or warmup_ratio"
      log_interval = args.log_interval
      save_interval = (
          args.save_interval
          if args.save_interval is not None
          else epoch_size // args.num_accumulated
      )
      logging.info(
          "#Total Training Steps={}, Warmup={}, Save Interval={}".format(
              num_train_steps, warmup_steps, save_interval
          )
      )
  
      # set up optimization
      lr_scheduler = PolyScheduler(
          max_update=num_train_steps,
          base_lr=args.lr,
          warmup_begin_lr=0,
          pwr=1,
          final_lr=0,
          warmup_steps=warmup_steps,
          warmup_mode="linear",
      )
      optimizer_params = {
          "learning_rate": args.lr,
          "wd": args.wd,
          "lr_scheduler": lr_scheduler,
      }
      adam_betas = ast.literal_eval(args.adam_betas)
      if args.optimizer == "adamw":
          optimizer_params.update(
              {
                  "beta1": adam_betas[0],
                  "beta2": adam_betas[1],
                  "epsilon": args.adam_epsilon,
                  "correct_bias": False,
              }
          )
      elif args.optimizer == "adam":
          optimizer_params.update(
              {
                  "beta1": adam_betas[0],
                  "beta2": adam_betas[1],
                  "epsilon": args.adam_epsilon,
              }
          )
      if use_amp:
          optimizer_params.update({"multi_precision": True})
      if args.comm_backend == "horovod":
          trainer = hvd.DistributedTrainer(
              param_dict, args.optimizer, optimizer_params
          )
      else:
          trainer = mx.gluon.Trainer(
              param_dict,
              args.optimizer,
              optimizer_params,
              update_on_kvstore=False,
          )
      if use_amp:
          amp.init_trainer(trainer)
      log_span_loss = 0
      log_answerable_loss = 0
      log_total_loss = 0
      log_sample_num = 0
  
      global_tic = time.time()
      tic = time.time()
      for step_num, batch_data in enumerate(
          grouper(
              repeat(train_dataloader), len(ctx_l) * num_accumulated
          )
      ):
          for sample_l in grouper(batch_data, len(ctx_l)):
              loss_l = []
              span_loss_l = []
              answerable_loss_l = []
              for sample, ctx in zip(sample_l, ctx_l):
                  if sample is None:
                      continue
                  # Copy the data to device
                  tokens = sample.data.as_in_ctx(ctx)
                  log_sample_num += len(tokens)
                  segment_ids = (
                      sample.segment_ids.as_in_ctx(ctx)
                      if use_segmentation
                      else None
                  )
                  valid_length = sample.valid_length.as_in_ctx(ctx)
                  p_mask = sample.masks.as_in_ctx(ctx)
                  gt_start = sample.gt_start.as_in_ctx(ctx).astype(
                      np.int32
                  )
                  gt_end = sample.gt_end.as_in_ctx(ctx).astype(np.int32)
                  is_impossible = sample.is_impossible.as_in_ctx(
                      ctx
                  ).astype(np.int32)
                  batch_idx = mx.np.arange(
                      tokens.shape[0], dtype=np.int32, ctx=ctx
                  )
                  p_mask = (
                      1 - p_mask
                  )  # In the network, we use 1 --> no_mask, 0 --> mask
                  with mx.autograd.record():
                      (
                          start_logits,
                          end_logits,
                          answerable_logits,
                      ) = qa_net(
                          tokens,
                          segment_ids,
                          valid_length,
                          p_mask,
                          gt_start,
                      )
                      sel_start_logits = start_logits[
                          batch_idx, gt_start
                      ]
                      sel_end_logits = end_logits[batch_idx, gt_end]
                      sel_answerable_logits = answerable_logits[
                          batch_idx, is_impossible
                      ]
                      span_loss = (
                          -0.5
                          * (sel_start_logits + sel_end_logits).mean()
                      )
                      answerable_loss = (
                          -0.5 * sel_answerable_logits.mean()
                      )
                      loss = (span_loss + answerable_loss) / (
                          len(ctx_l) * num_accumulated
                      )
                      loss_l.append(loss)
                      span_loss_l.append(span_loss)
                      answerable_loss_l.append(answerable_loss)
              if use_amp:
                  with mx.autograd.record():
                      with amp.scale_loss(
                          loss_l, trainer
                      ) as amp_loss_l:
                          for loss in amp_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
  
              # All Reduce the Step Loss
              log_span_loss += sum(
                  [ele.as_in_ctx(ctx_l[0]) for ele in span_loss_l]
              ).asnumpy()
              log_total_loss += sum(
                  [ele.as_in_ctx(ctx_l[0]) for ele in loss_l]
              ).asnumpy()
              log_answerable_loss += sum(
                  [ele.as_in_ctx(ctx_l[0]) for ele in answerable_loss_l]
              ).asnumpy()
          # update
          trainer.allreduce_grads()
  
          if args.max_grad_norm > 0:
              total_norm, ratio, is_finite = clip_grad_global_norm(
                  params, args.max_grad_norm * norm_clip_mult
              )
          else:
              total_norm = grad_global_norm(params)
  
          if args.comm_backend == "horovod":
              # Note that horovod.trainer._scale is default to num_workers,
              # thus trainer.update(1) will scale the gradients by 1./num_workers
              trainer.update(1, ignore_stale_grad=True)
          else:
              # gluon.trainer._scale is default to 1
              trainer.update(num_workers, ignore_stale_grad=True)
  
          total_norm = total_norm / norm_clip_mult
          if args.num_accumulated > 1:
              # set grad to zero for gradient accumulation
              qa_net.zero_grad()
  
          # saving
          if (
              local_rank == 0
              and (step_num + 1) % save_interval == 0
              or (step_num + 1) >= num_train_steps
          ):
              version_prefix = "squad" + args.version
              ckpt_name = "{}_{}_{}.params".format(
                  args.model_name, version_prefix, (step_num + 1)
              )
              params_saved = os.path.join(args.output_dir, ckpt_name)
              qa_net.save_parameters(params_saved)
              ckpt_candidates = [
                  f
                  for f in os.listdir(args.output_dir)
                  if f.endswith(".params")
              ]
              # keep last `max_saved_ckpt` checkpoints
              if len(ckpt_candidates) > args.max_saved_ckpt:
                  ckpt_candidates.sort(key=lambda ele: (len(ele), ele))
                  os.remove(
                      os.path.join(args.output_dir, ckpt_candidates[0])
                  )
              logging.info("Params saved in: {}".format(params_saved))
  
          # logging
          if (step_num + 1) % log_interval == 0:
              log_span_loss /= log_sample_num
              log_answerable_loss /= log_sample_num
              log_total_loss /= log_sample_num
              toc = time.time()
              logging.info(
                  "Step: {}/{}, Loss span/answer/total={:.4f}/{:.4f}/{:.4f},"
                  " LR={:.8f}, grad_norm={:.4f}. Time cost={:.2f}, Throughput={:.2f} samples/s"
                  " ETA={:.2f}h".format(
                      (step_num + 1),
                      num_train_steps,
                      log_span_loss,
                      log_answerable_loss,
                      log_total_loss,
                      trainer.learning_rate,
                      total_norm,
                      toc - tic,
                      log_sample_num / (toc - tic),
                      (num_train_steps - (step_num + 1))
                      / ((step_num + 1) / (toc - global_tic))
                      / 3600,
                  )
              )
              tic = time.time()
              log_span_loss = 0
              log_answerable_loss = 0
              log_total_loss = 0
              log_sample_num = 0
  
          if (step_num + 1) >= num_train_steps:
              toc = time.time()
              logging.info(
                  "Finish training step: {} within {} hours".format(
                      step_num + 1, (toc - global_tic) / 3600
                  )
              )
              break
  
      return params_saved
  
imports: import os
  import json
  import time
  import logging
  import argparse
  import ast
  import functools
  import collections
  import dataclasses
  from dataclasses import dataclass
  from multiprocessing import Pool, cpu_count
  import mxnet
  import numpy
  from mxnet.lr_scheduler import PolyScheduler
  import gluonnlp.data.batchify
  from models import ModelForQABasic, ModelForQAConditionalV1
  from eval_utils import squad_eval
  from squad_utils import SquadFeature, get_squad_examples, convert_squad_example_to_feature
  from gluonnlp.models import get_backbone
  from gluonnlp.utils.misc import repeat, grouper, set_seed, init_comm, logging_config, parse_ctx
  from gluonnlp.initializer import TruncNorm
  from gluonnlp.data.sampler import SplitSampler
  from gluonnlp.utils.parameter import grad_global_norm, clip_grad_global_norm, count_parameters, deduplicate_param_dict
label_desc: Use image node_id 5 for calling a global function with example usage: train(args) and returns: params_saved