Function _train

linear_moe/finetune_utils.py:159–265

Train the model. Args: model (nn.Module): The model to train. optimizer (Optimizer): The optimizer to use for gradient updates. opt_param_scheduler (Optional): The optimizer parameter scheduler. forward_step (callable): The forward step function for the mode

(model, optimizer, opt_param_scheduler, forward_step,
           train_dataloader, valid_dataloader, end_of_epoch_callback, config)

def _train(model, optimizer, opt_param_scheduler, forward_step,
           train_dataloader, valid_dataloader, end_of_epoch_callback, config):
    """
    Train the model.

    Args:
        model (nn.Module): The model to train.
        optimizer (Optimizer): The optimizer to use for gradient updates.
        opt_param_scheduler (Optional): The optimizer parameter scheduler.
        forward_step (callable): The forward step function for the model.
        train_dataloader (DataLoader): The dataloader for training data.
        valid_dataloader (DataLoader): The dataloader for validation data.
        end_of_epoch_callback (Optional[callable]): The callback function to call at the end of each epoch.
    """

    args = get_args()
    timers = get_timers()

    assert get_num_microbatches(
    ) == 1, "finetuning with gradient accumulation doesn't currently work"

    # Turn on training mode which enables dropout.
    for m in model:
        m.train()

    # Tracking loss.
    losses_dict_sum = {}

    # Starting epoch and iteration
    start_epoch = args.iteration // args.train_iters_per_epoch
    start_iteration = args.iteration % args.train_iters_per_epoch
    iteration = args.iteration

    # Memory reporting flag.
    report_memory_flag = True
    # For each remaining epoch
    timers('interval-time', log_level=0).start(barrier=True)
    for epoch in range(start_epoch, args.epochs):
        print_rank_0('working on epoch {} ...'.format(epoch + 1))

        # Set the data loader epoch to shuffle the index iterator.
        train_dataloader.sampler.set_epoch(args.seed + epoch)

        # For all the batches in the dataset.
        for iteration_, batch in enumerate(train_dataloader):

            # Ignore the iterations before starting value
            if iteration_ < start_iteration:
                continue
            # Set to zero so the next epoch does not skip any batches.
            start_iteration = 0

            # Train for one step.
            out = train_step(forward_step, batch, model, optimizer,
                             opt_param_scheduler, config)

            losses_dict, skipped_iter, grad_norm, num_zeros_in_grad = out
            iteration += 1