def train(*, rng: jnp.ndarray, config: ml_collections.ConfigDict,
          dataset: dataset_utils.Dataset, workdir: str,
          writer: metric_writers.MetricWriter):
  """Main training loop lives in this function."""
  lead_host = jax.process_index() == 0
  dtype = train_utils.get_dtype(config)

  tokenizer_fn = maskgvt_train_lib.load_vq_tokenizer(
      config, tokenizer_only=True)['tokenizer']
  writer.write_hparams(eval_utils.flatten_config(config))

  # Build the flax_model and the optimizers.
  rng, init_rng, task_rng = jax.random.split(rng, 3)
  input_spec = [(
      dataset.meta_data['input_shape'],  # bs, t, h, w, 3
      dataset.meta_data.get('input_dtype', dtype))]

  # Get training tasks
  task_reg = task_registry.get_task_registry(config, True)
  task_ids = config.get('tasks', ('full_generation',))
  logging.info('Training tasks %s', task_reg)

  model, model_config, train_state, learning_rate_fn = create_train_state(
      input_spec, task_reg[task_ids[0]], tokenizer_fn, config, init_rng)
  del model_config

  # Restores the transformer model if needed.
  if config.logging.enable_checkpoint:
    train_state = checkpoints.restore_checkpoint(
        ckpt_dir=workdir, target=train_state)
  start_step = int(train_state.global_step)

  # Log model size
  if start_step == 0:
    param_size = parameter_overview.count_parameters(
        train_state.params)
    writer.write_scalars(0, {'param_size': param_size})

  if (start_step == 0  # Which means "no" checkpoint is restored!
      and config.get('init_from') is not None):
    init_checkpoint_path = config.init_from.get('checkpoint_path')
    train_state = maskgvt_train_lib.init_from_pretrained_checkpoint(
        init_checkpoint_path, train_state,
        config.init_from.get('expand_embedding', False), config)

  # Get learning rate scheduler.
  total_steps, steps_per_epoch = scenic_train_utils.get_num_training_steps(
      config, dataset.meta_data)
  train_utils.check_training_step(steps_per_epoch, config)

  # pmap
  tokenizer_pmapped = jax.pmap(
      functools.partial(
          maskgvt_train_lib.tokenize_with_cond,
          tokenizer_fn=tokenizer_fn),
      axis_name='device',
      static_broadcasted_argnums=(1,),
      donate_argnums=(0),)

  train_state = jax_utils.replicate(train_state)
  train_step_pmapped = jax.pmap(
      functools.partial(train_step, flax_model=model, config=config),
      axis_name='device',
      static_broadcasted_argnums=(2,),
      # We can donate the buffer of train_state.
      donate_argnums=(0),
  )

  sample_step_pmapped = jax.pmap(
      functools.partial(sample_step, flax_model=model, config=config),
      static_broadcasted_argnums=(2,),
      axis_name='device',
  )

  checkpoint_steps = config.logging.get('checkpoint_steps', steps_per_epoch)
  log_metric_steps = config.logging.get('log_metric_steps', checkpoint_steps)
  log_sample_size = config.logging.get('log_sample_size', 8)

  logging.info('Starting training loop at step %d of total_steps=%d.',
               start_step, total_steps)

  report_progress = periodic_actions.ReportProgress(
      num_train_steps=total_steps, writer=writer)
  hooks = [report_progress]
  if config.get('xprof', True) and lead_host:
    hooks.append(periodic_actions.Profile(num_profile_steps=5, logdir=workdir))

  train_metrics = None

  with metric_writers.ensure_flushes(writer):
    for step in range(start_step + 1, total_steps + 1):
      with jax.profiler.StepTraceAnnotation('train', step_num=step):
        train_batch = next(dataset.train_iter)
        cur_task_rng = jax.random.fold_in(task_rng, step)
        # TODO(Lijun-Yu): use probabilistic sampling to favor stronger tasks.
        task_idx = jax.random.choice(cur_task_rng, len(task_ids))
        task = task_reg[task_ids[task_idx]]
        train_batch.update(
            tokenizer_pmapped(train_batch.pop('inputs'), task.cond_fn))
        train_state, metrics_update = train_step_pmapped(
            train_state, train_batch, task)
        metric_update = jax_utils.unreplicate(metrics_update)
        train_metrics = (
            metric_update
            if train_metrics is None else train_metrics.merge(metric_update))
      # Quick indication that training is happening.
      logging.log_first_n(logging.INFO, 'Finished training step %d.', 5, step)
      for h in hooks:
        h(step)

      ############### LOG TRAIN METRICS ###############
      if (step % log_metric_steps == 1) or (step == total_steps):
        train_metrics = train_metrics.compute()
        train_metrics = {f'train/{k}': v for k, v in train_metrics.items()}
        writer.write_scalars(step, train_metrics)
        writer.write_scalars(step, {'train/lr': learning_rate_fn(step-1)})

        # Reset metric accumulation for next evaluation cycle.
        if step < total_steps:
          train_metrics = None

      ##################### CHECKPOINTING ############################
      if ((step % checkpoint_steps == 1) or (step == total_steps)):
        with report_progress.timed('checkpoint'):
          # Sync model state across replicas.
          train_state = maskgvt_train_lib.sync_model_state_across_replicas(
              train_state)
          if config.logging.enable_checkpoint:
            train_utils.save_checkpoint(workdir, train_state,
                                        config.logging.checkpoint_kept)

        ############### LOG TRAIN SAMPLES ###############
        with report_progress.timed('sample'):
          sample_batch = jax.tree_util.tree_map(
              lambda x: x[:, :log_sample_size], train_batch)
          batch_outputs, _ = sample_step_pmapped(train_state, sample_batch,
                                                 task)
          batch_outputs = jax.device_get(jax_utils.unreplicate(batch_outputs))
          samples = {}
          key = 'predicted_tokens'
          samples[key] = batch_outputs
          samples = {f'train_sample/{k}': v for k, v in samples.items()}
          writer.write_images(step, samples)
          del sample_batch, batch_outputs, samples

  # Tell evaluation job to stop.
  task_manager = eval_utils.TaskManager(workdir)
  task_manager.mark_training_done()
  # Wait until computations are done before exiting.
  jax.random.normal(jax.random.PRNGKey(0), ()).block_until_ready()
  logging.info('Finishing training at step %d', total_steps)
  # Return the train and eval summary after last step for regression testing.
  return train_state, train_metrics
Main training loop lives in this function.
videogvt/trainers/lmgvt_trainer.py