Function finetune

linear_moe/finetune_utils.py:268–350

Main finetune function used across all tasks. Args: model (nn.Module): The model to fine-tune. optimizer (Optimizer): The optimizer to use for gradient updates. opt_param_scheduler (Optional): The optimizer parameter scheduler. forward_step (callable): The fo

(train_valid_datasets_provider,
             model_provider,
             model_type=ModelType.encoder_or_decoder,
             forward_step=_cross_entropy_forward_step,
             end_of_epoch_callback_provider=None,
             task_collate_fn=None)

def finetune(train_valid_datasets_provider,
             model_provider,
             model_type=ModelType.encoder_or_decoder,
             forward_step=_cross_entropy_forward_step,
             end_of_epoch_callback_provider=None,
             task_collate_fn=None):
    """
    Main finetune function used across all tasks.
    Args:
        model (nn.Module): The model to fine-tune.
        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 args.rampup_batch_size is None, \
        'batch size scaling is not supported for finetuning'

    # Train and validation data loaders.
    timers('train/valid/test dataset/dataloder', log_level=0).start()
    if args.epochs > 0:
        train_dataset, valid_dataset = train_valid_datasets_provider()
        train_dataloader, valid_dataloader = _build_train_valid_dataloaders(
            train_dataset, valid_dataset, task_collate_fn)
    else:
        args.train_iters = 0
    timers('train/valid/test dataset/dataloder').stop()

    # Build calback function.
    timers('callback function', log_level=0).start()
    end_of_epoch_callback = None
    if end_of_epoch_callback_provider is not None:
        end_of_epoch_callback = end_of_epoch_callback_provider()
    timers('callback function').stop()

    # Build model, optimizer and learning rate scheduler.
    timers('model and optimizer', log_level=0).start()
    model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
        model_provider, model_type)
    timers('model and optimizer').stop()

    # If pretrained checkpoint is provided and we have not trained for
    # any iteration (i.e., iteration is zero), then load the pretrained
    # checkpoint.
    timers('pretrained checkpoint', log_level=0).start(barrier=True)
    if args.iteration == 0 and args.pretrained_checkpoint is not None:
        original_load = args.load
        args.load = args.pretrained_checkpoint
        original_rng = args.no_load_rng
        args.no_load_rng = True
        _ = load_checkpoint(model, None, None)
        args.load = original_load
        args.no_load_rng = original_rng
        # This is critical when only model is loaded. We should make sure