def save_checkpoint(workdir: str,
                    train_state: AllTrainState,
                    max_to_keep: int = 3,
                    overwrite: bool = False):
  """Saves a checkpoint.

  First syncs the model state across replicas, then it unreplicates it by taking
  the train state of the first replica and saves it as a checkpoint.

  Args:
    workdir: Experiment directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    max_to_keep: The number of checkpoints to keep.
    overwrite: Overwrite existing checkpoint  if a checkpoint
      at the current or a later step already exits (default: False).
  """
  # Get train state from the first replica.
  checkpoint_state = jax.device_get(jax_utils.unreplicate(train_state))
  checkpoints.save_checkpoint_multiprocess(
      workdir,
      checkpoint_state,
      int(checkpoint_state.global_step),
      overwrite=overwrite,
      keep=max_to_keep)
Saves a checkpoint.

  First syncs the model state across replicas, then it unreplicates it by taking
  the train state of the first replica and saves it as a checkpoint.

  Args:
    workdir: Experiment directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    max_to_keep: The number of checkpoints to keep.
    overwrite: Overwrite existing checkpoint  if a checkpoint
      at the current or a later step already exits (default: False).
  
videogvt/train_lib/train_utils.py