docs: def main():
      """
      Launches data-parallel multi-gpu training.
      """
      mlperf_log.ROOT_DIR_GNMT = os.path.dirname(os.path.abspath(__file__))
      mlperf_log.LOGGER.propagate = False
  
      args = parse_args()
      device = utils.set_device(args.cuda, args.local_rank)
      distributed = utils.init_distributed(args.cuda)
      gnmt_print(key=mlperf_log.RUN_START, sync=True)
      args.rank = utils.get_rank()
  
      if not args.cudnn:
          torch.backends.cudnn.enabled = False
  
      # create directory for results
      save_path = os.path.join(args.results_dir, args.save)
      args.save_path = save_path
      os.makedirs(save_path, exist_ok=True)
  
      # setup logging
      log_filename = f'log_rank_{utils.get_rank()}.log'
      utils.setup_logging(os.path.join(save_path, log_filename))
  
      if args.env:
          utils.log_env_info()
  
      logging.info(f'Saving results to: {save_path}')
      logging.info(f'Run arguments: {args}')
  
      # automatically set train_iter_size based on train_global_batch_size,
      # world_size and per-worker train_batch_size
      if args.train_global_batch_size is not None:
          global_bs = args.train_global_batch_size
          bs = args.train_batch_size
          world_size = utils.get_world_size()
          assert global_bs % (bs * world_size) == 0
          args.train_iter_size = global_bs // (bs * world_size)
          logging.info(f'Global batch size was set in the config, '
                       f'Setting train_iter_size to {args.train_iter_size}')
  
      worker_seeds, shuffling_seeds = utils.setup_seeds(args.seed, args.epochs,
                                                        device)
      worker_seed = worker_seeds[args.rank]
      logging.info(f'Worker {args.rank} is using worker seed: {worker_seed}')
      torch.manual_seed(worker_seed)
  
      # build tokenizer
      pad_vocab = utils.pad_vocabulary(args.math)
      tokenizer = Tokenizer(os.path.join(args.dataset_dir, config.VOCAB_FNAME),
                            pad_vocab)
  
      # build datasets
      gnmt_print(key=mlperf_log.PREPROC_TOKENIZE_TRAINING, sync=False)
      gnmt_print(key=mlperf_log.TRAIN_HP_MAX_SEQ_LEN,
                 value=args.max_length_train, sync=False)
  
      train_data = LazyParallelDataset(
          src_fname=os.path.join(args.dataset_dir, config.SRC_TRAIN_FNAME),
          tgt_fname=os.path.join(args.dataset_dir, config.TGT_TRAIN_FNAME),
          tokenizer=tokenizer,
          min_len=args.min_length_train,
          max_len=args.max_length_train,
          sort=False,
          max_size=args.max_size)
  
      gnmt_print(key=mlperf_log.PREPROC_NUM_TRAIN_EXAMPLES,
                 value=len(train_data), sync=False)
  
      val_data = ParallelDataset(
          src_fname=os.path.join(args.dataset_dir, config.SRC_VAL_FNAME),
          tgt_fname=os.path.join(args.dataset_dir, config.TGT_VAL_FNAME),
          tokenizer=tokenizer,
          min_len=args.min_length_val,
          max_len=args.max_length_val,
          sort=True)
  
      gnmt_print(key=mlperf_log.PREPROC_TOKENIZE_EVAL, sync=False)
  
      test_data = TextDataset(
          src_fname=os.path.join(args.dataset_dir, config.SRC_TEST_FNAME),
          tokenizer=tokenizer,
          min_len=args.min_length_test,
          max_len=args.max_length_test,
          sort=True)
  
      gnmt_print(key=mlperf_log.PREPROC_NUM_EVAL_EXAMPLES,
                 value=len(test_data), sync=False)
  
      vocab_size = tokenizer.vocab_size
      gnmt_print(key=mlperf_log.PREPROC_VOCAB_SIZE,
                 value=vocab_size, sync=False)
  
      # build GNMT model
      model_config = {'hidden_size': args.hidden_size,
                      'num_layers': args.num_layers,
                      'dropout': args.dropout, 'batch_first': False,
                      'share_embedding': args.share_embedding}
      model = GNMT(vocab_size=vocab_size, **model_config)
      logging.info(model)
  
      batch_first = model.batch_first
  
      # define loss function (criterion) and optimizer
      criterion = build_criterion(vocab_size, config.PAD, args.smoothing)
  
      opt_config = {'optimizer': args.optimizer, 'lr': args.lr}
      opt_config.update(literal_eval(args.optimizer_extra))
      logging.info(f'Training optimizer config: {opt_config}')
  
      scheduler_config = {'warmup_steps': args.warmup_steps,
                          'remain_steps': args.remain_steps,
                          'decay_interval': args.decay_interval,
                          'decay_steps': args.decay_steps,
                          'decay_factor': args.decay_factor}
  
      logging.info(f'Training LR schedule config: {scheduler_config}')
  
      num_parameters = sum([l.nelement() for l in model.parameters()])
      logging.info(f'Number of parameters: {num_parameters}')
  
      batching_opt = {'shard_size': args.shard_size,
                      'num_buckets': args.num_buckets}
      # get data loaders
      train_loader = train_data.get_loader(batch_size=args.train_batch_size,
                                           seeds=shuffling_seeds,
                                           batch_first=batch_first,
                                           shuffle=True,
                                           batching=args.batching,
                                           batching_opt=batching_opt,
                                           num_workers=args.train_loader_workers)
  
      gnmt_print(key=mlperf_log.INPUT_BATCH_SIZE,
                 value=args.train_batch_size * utils.get_world_size(),
                 sync=False)
      gnmt_print(key=mlperf_log.INPUT_SIZE,
                 value=train_loader.sampler.num_samples, sync=False)
  
      val_loader = val_data.get_loader(batch_size=args.val_batch_size,
                                       batch_first=batch_first,
                                       shuffle=False,
                                       num_workers=args.val_loader_workers)
  
      test_loader = test_data.get_loader(batch_size=args.test_batch_size,
                                         batch_first=batch_first,
                                         shuffle=False,
                                         pad=True,
                                         num_workers=args.test_loader_workers)
  
      gnmt_print(key=mlperf_log.EVAL_SIZE,
                 value=len(test_loader.dataset), sync=False)
  
      translator = Translator(model=model,
                              tokenizer=tokenizer,
                              loader=test_loader,
                              beam_size=args.beam_size,
                              max_seq_len=args.max_length_test,
                              len_norm_factor=args.len_norm_factor,
                              len_norm_const=args.len_norm_const,
                              cov_penalty_factor=args.cov_penalty_factor,
                              cuda=args.cuda,
                              print_freq=args.print_freq,
                              dataset_dir=args.dataset_dir,
                              target_bleu=args.target_bleu,
                              save_path=args.save_path)
  
      # create trainer
      total_train_iters = len(train_loader) // args.train_iter_size * args.epochs
      save_info = {'model_config': model_config, 'config': args, 'tokenizer':
                   tokenizer.get_state()}
      trainer_options = dict(
          criterion=criterion,
          grad_clip=args.grad_clip,
          iter_size=args.train_iter_size,
          save_path=save_path,
          save_freq=args.save_freq,
          save_info=save_info,
          opt_config=opt_config,
          scheduler_config=scheduler_config,
          train_iterations=total_train_iters,
          batch_first=batch_first,
          keep_checkpoints=args.keep_checkpoints,
          math=args.math,
          print_freq=args.print_freq,
          cuda=args.cuda,
          distributed=distributed,
          intra_epoch_eval=args.intra_epoch_eval,
          translator=translator)
  
      trainer_options['model'] = model
      trainer = trainers.Seq2SeqTrainer(**trainer_options)
  
      # optionally resume from a checkpoint
      if args.resume:
          checkpoint_file = args.resume
          if os.path.isdir(checkpoint_file):
              checkpoint_file = os.path.join(
                  checkpoint_file, 'model_best.pth')
          if os.path.isfile(checkpoint_file):
              trainer.load(checkpoint_file)
          else:
              logging.error(f'No checkpoint found at {args.resume}')
  
      # training loop
      best_loss = float('inf')
      break_training = False
      test_bleu = None
      gnmt_print(key=mlperf_log.TRAIN_LOOP, sync=True)
      for epoch in range(args.start_epoch, args.epochs):
          logging.info(f'Starting epoch {epoch}')
          gnmt_print(key=mlperf_log.TRAIN_EPOCH,
                     value=epoch, sync=True)
  
          train_loader.sampler.set_epoch(epoch)
  
          trainer.epoch = epoch
          train_loss, train_perf = trainer.optimize(train_loader)
  
          # evaluate on validation set
          if args.eval:
              logging.info(f'Running validation on dev set')
              val_loss, val_perf = trainer.evaluate(val_loader)
  
              # remember best prec@1 and save checkpoint
              gnmt_print(key=mlperf_log.TRAIN_CHECKPOINT, sync=False)
              if args.rank == 0:
                  is_best = val_loss < best_loss
                  best_loss = min(val_loss, best_loss)
                  trainer.save(save_all=args.save_all, is_best=is_best)
  
          if args.eval:
              gnmt_print(key=mlperf_log.EVAL_START, value=epoch, sync=True)
              test_bleu, break_training = translator.run(calc_bleu=True,
                                                         epoch=epoch)
              gnmt_print(key=mlperf_log.EVAL_ACCURACY,
                         value={"epoch": epoch, "value": round(test_bleu, 2)},
                         sync=False)
              gnmt_print(key=mlperf_log.EVAL_TARGET,
                         value=args.target_bleu, sync=False)
              gnmt_print(key=mlperf_log.EVAL_STOP, sync=True)
  
          acc_log = []
          acc_log += [f'Summary: Epoch: {epoch}']
          acc_log += [f'Training Loss: {train_loss:.4f}']
          if args.eval:
              acc_log += [f'Validation Loss: {val_loss:.4f}']
              acc_log += [f'Test BLEU: {test_bleu:.2f}']
  
          perf_log = []
          perf_log += [f'Performance: Epoch: {epoch}']
          perf_log += [f'Training: {train_perf:.0f} Tok/s']
          if args.eval:
              perf_log += [f'Validation: {val_perf:.0f} Tok/s']
  
          if args.rank == 0:
              logging.info('\t'.join(acc_log))
              logging.info('\t'.join(perf_log))
  
          logging.info(f'Finished epoch {epoch}')
          if break_training:
              break
  
      gnmt_print(key=mlperf_log.RUN_STOP,
                 value={"success": bool(break_training)}, sync=True)
      gnmt_print(key=mlperf_log.RUN_FINAL, sync=False)