def train(
        checkpoint_dir: str = None,
        load_checkpoint_dir: str = None,
        # Dataset Parameters
        mask_prob: float = 0.15,
        random_replace_prob: float = 0.1,
        unmask_replace_prob: float = 0.1,
        max_seq_length: int = 512,
        tokenizer: str = "roberta-base",
        # Model Parameters
        num_layers: int = 6,
        num_heads: int = 8,
        ff_dim: int = 512,
        h_dim: int = 256,
        dropout: float = 0.1,
        # Training Parameters
        batch_size: int = 8,
        num_iterations: int = 10000,
        checkpoint_every: int = 1000,
        log_every: int = 10,
        local_rank: int = -1,
) -> pathlib.Path:
    """Trains a [Bert style](https://arxiv.org/pdf/1810.04805.pdf)
    (transformer encoder only) model for MLM Task

    Args:
        checkpoint_dir (str):
            The base experiment directory to save experiments to
        mask_prob (float, optional):
            The fraction of tokens to mask. Defaults to 0.15.
        random_replace_prob (float, optional):
            The fraction of masked tokens to replace with random token.
            Defaults to 0.1.
        unmask_replace_prob (float, optional):
            The fraction of masked tokens to leave unchanged.
            Defaults to 0.1.
        max_seq_length (int, optional):
            The maximum sequence length of the examples. Defaults to 512.
        tokenizer (str, optional):
            The tokenizer to use. Defaults to "roberta-base".
        num_layers (int, optional):
            The number of layers in the Bert model. Defaults to 6.
        num_heads (int, optional):
            Number of attention heads to use. Defaults to 8.
        ff_dim (int, optional):
            Size of the intermediate dimension in the FF layer.
            Defaults to 512.
        h_dim (int, optional):
            Size of intermediate representations.
            Defaults to 256.
        dropout (float, optional):
            Amout of Dropout to use. Defaults to 0.1.
        batch_size (int, optional):
            The minibatch size. Defaults to 8.
        num_iterations (int, optional):
            Total number of iterations to run the model for.
            Defaults to 10000.
        checkpoint_every (int, optional):
            Save checkpoint after these many steps.

            ..note ::

                You want this to be frequent enough that you can
                resume training in case it crashes, but not so much
                that you fill up your entire storage !

            Defaults to 1000.
        log_every (int, optional):
            Print logs after these many steps. Defaults to 10.
        local_rank (int, optional):
            Which GPU to run on (-1 for CPU). Defaults to -1.

    Returns:
        pathlib.Path: The final experiment directory

    """
    device = (torch.device(get_accelerator().device_name(), local_rank) if (local_rank > -1)
              and get_accelerator().is_available() else torch.device("cpu"))
    ################################
    ###### Create Exp. Dir #########
    ################################
    if checkpoint_dir is None and load_checkpoint_dir is None:
        logger.error("Need to specify one of checkpoint_dir"
                     " or load_checkpoint_dir")
        return
    if checkpoint_dir is not None and load_checkpoint_dir is not None:
        logger.error("Cannot specify both checkpoint_dir"
                     " and load_checkpoint_dir")
        return
    if checkpoint_dir:
        logger.info("Creating Experiment Directory")
        checkpoint_dir = pathlib.Path(checkpoint_dir)
        checkpoint_dir.mkdir(exist_ok=True)
        all_arguments = {
            # Dataset Params
            "mask_prob": mask_prob,
            "random_replace_prob": random_replace_prob,
            "unmask_replace_prob": unmask_replace_prob,
            "max_seq_length": max_seq_length,
            "tokenizer": tokenizer,
            # Model Params
            "num_layers": num_layers,
            "num_heads": num_heads,
            "ff_dim": ff_dim,
            "h_dim": h_dim,
            "dropout": dropout,
            # Training Params
            "batch_size": batch_size,
            "num_iterations": num_iterations,
            "checkpoint_every": checkpoint_every,
        }
        exp_dir = create_experiment_dir(checkpoint_dir, all_arguments)
        logger.info(f"Experiment Directory created at {exp_dir}")
    else:
        logger.info("Loading from Experiment Directory")
        load_checkpoint_dir = pathlib.Path(load_checkpoint_dir)
        assert load_checkpoint_dir.exists()
        with (load_checkpoint_dir / "hparams.json").open("r") as handle:
            hparams = json.load(handle)
        # Set the hparams
        # Dataset Params
        mask_prob = hparams.get("mask_prob", mask_prob)
        tokenizer = hparams.get("tokenizer", tokenizer)
        random_replace_prob = hparams.get("random_replace_prob",
                                          random_replace_prob)
        unmask_replace_prob = hparams.get("unmask_replace_prob",
                                          unmask_replace_prob)
        max_seq_length = hparams.get("max_seq_length", max_seq_length)
        # Model Params
        ff_dim = hparams.get("ff_dim", ff_dim)
        h_dim = hparams.get("h_dim", h_dim)
        dropout = hparams.get("dropout", dropout)
        num_layers = hparams.get("num_layers", num_layers)
        num_heads = hparams.get("num_heads", num_heads)
        # Training Params
        batch_size = hparams.get("batch_size", batch_size)
        _num_iterations = hparams.get("num_iterations", num_iterations)
        num_iterations = max(num_iterations, _num_iterations)
        checkpoint_every = hparams.get("checkpoint_every", checkpoint_every)
        exp_dir = load_checkpoint_dir
    # Tensorboard writer
    tb_dir = exp_dir / "tb_dir"
    assert tb_dir.exists()
    summary_writer = SummaryWriter(log_dir=tb_dir)
    ################################
    ###### Create Datasets #########
    ################################
    logger.info("Creating Datasets")
    data_iterator = create_data_iterator(
        mask_prob=mask_prob,
        random_replace_prob=random_replace_prob,
        unmask_replace_prob=unmask_replace_prob,
        tokenizer=tokenizer,
        max_seq_length=max_seq_length,
        batch_size=batch_size,
    )
    logger.info("Dataset Creation Done")
    ################################
    ###### Create Model ############
    ################################
    logger.info("Creating Model")
    model = create_model(
        num_layers=num_layers,
        num_heads=num_heads,
        ff_dim=ff_dim,
        h_dim=h_dim,
        dropout=dropout,
    )
    model = model.to(device)
    logger.info("Model Creation Done")
    ################################
    ###### Create Optimizer #######
    ################################
    logger.info("Creating Optimizer")
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    logger.info("Optimizer Creation Done")
    ################################
    #### Load Model checkpoint #####
    ################################
    start_step = 1
    if load_checkpoint_dir is not None:
        checkpoint_step, model, optimizer = load_model_checkpoint(
            load_checkpoint_dir, model, optimizer)
        start_step = checkpoint_step + 1

    ################################
    ####### The Training Loop ######
    ################################
    logger.info(
        f"Total number of model parameters: {sum([p.numel() for p in model.parameters()]):,d}"
    )
    model.train()
    losses = []
    for step, batch in enumerate(data_iterator, start=start_step):
        if step >= num_iterations:
            break
        optimizer.zero_grad()
        # Move the tensors to device
        for key, value in batch.items():
            batch[key] = value.to(device)
        # Forward pass
        loss = model(**batch)
        # Backward pass
        loss.backward()
        # Optimizer Step
        optimizer.step()
        losses.append(loss.item())
        if step % log_every == 0:
            logger.info("Loss: {0:.4f}".format(np.mean(losses)))
            summary_writer.add_scalar(f"Train/loss", np.mean(losses), step)
        if step % checkpoint_every == 0:
            state_dict = {
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
            }
            torch.save(obj=state_dict,
                       f=str(exp_dir / f"checkpoint.iter_{step}.pt"))
            logger.info("Saved model to {0}".format(
                (exp_dir / f"checkpoint.iter_{step}.pt")))
    # Save the last checkpoint if not saved yet
    if step % checkpoint_every != 0:
        state_dict = {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
        }
        torch.save(obj=state_dict,
                   f=str(exp_dir / f"checkpoint.iter_{step}.pt"))
        logger.info("Saved model to {0}".format(
            (exp_dir / f"checkpoint.iter_{step}.pt")))

    return exp_dir
Trains a [Bert style](https://arxiv.org/pdf/1810.04805.pdf)
    (transformer encoder only) model for MLM Task

    Args:
        checkpoint_dir (str):
            The base experiment directory to save experiments to
        mask_prob (float, optional):
            The fraction of tokens to mask. Defaults to 0.15.
        random_replace_prob (float, optional):
            The fraction of masked tokens to replace with random token.
            Defaults to 0.1.
        unmask_replace_prob (float, optional):
            The fraction of masked tokens to leave unchanged.
            Defaults to 0.1.
        max_seq_length (int, optional):
            The maximum sequence length of the examples. Defaults to 512.
        tokenizer (str, optional):
            The tokenizer to use. Defaults to "roberta-base".
        num_layers (int, optional):
            The number of layers in the Bert model. Defaults to 6.
        num_heads (int, optional):
            Number of attention heads to use. Defaults to 8.
        ff_dim (int, optional):
            Size of the intermediate dimension in the FF layer.
            Defaults to 512.
        h_dim (int, optional):
            Size of intermediate representations.
            Defaults to 256.
        dropout (float, optional):
            Amout of Dropout to use. Defaults to 0.1.
        batch_size (int, optional):
            The minibatch size. Defaults to 8.
        num_iterations (int, optional):
            Total number of iterations to run the model for.
            Defaults to 10000.
        checkpoint_every (int, optional):
            Save checkpoint after these many steps.

            ..note ::

                You want this to be frequent enough that you can
                resume training in case it crashes, but not so much
                that you fill up your entire storage !

            Defaults to 1000.
        log_every (int, optional):
            Print logs after these many steps. Defaults to 10.
        local_rank (int, optional):
            Which GPU to run on (-1 for CPU). Defaults to -1.

    Returns:
        pathlib.Path: The final experiment directory

    
training/HelloDeepSpeed/train_bert.py