Python version: 3.9.6


Transformers version: 4.11.3


Ray version: 1.7.0


Model Type: BART-Large


Task: Summarization


Using PyTorch Trainer


Hi @patrickvonplaten,
First off, thank you for helping creating the fantastic huggingface library. My issue is that I’m trying to finetune a bart-large model on a summarization task, and I am using the population-based training (PBT) scheduler from Ray Tune for this purpose. However, there are some cases (not all the time) when a trial is stuck, “perturbing” the same trial without creating any more checkpoints. The “perturb” is different every time (e.g., a new learning_rate is used). This constant “perturbing” of the same trial continues without end, until the disk quota is exceeded and the project fails.
Here is the most relevant part of the code:

SEED = 4


# Define Training Arguments

training_args = Seq2SeqTrainingArguments(
    do_train = True,
    evaluation_strategy = "epoch", 
    eval_accumulation_steps=4,
    fp16 = True,
    gradient_accumulation_steps = 4,
    learning_rate = 2.232e-5,
    load_best_model_at_end=True,
    logging_strategy = 'epoch',
    metric_for_best_model = 'eval_rouge1',
    num_train_epochs = 5,
    output_dir = 'experimental_summarization_runs',
    overwrite_output_dir = True,
    per_device_eval_batch_size = 1,
    per_device_train_batch_size = 1,
    predict_with_generate = True,
    remove_unused_columns=True,
    report_to="wandb",
    save_strategy = 'epoch',
    save_total_limit = 1,
    seed = SEED,
    warmup_ratio =0.06,
    weight_decay = 0.01,
)


# In[82]:


# Define Trainer for Hyperparameter Tuning only on portion of train and eval datasets
trainer = Seq2SeqTrainer(
    args = training_args,
    tokenizer = tokenizer,
    train_dataset = encoded_train_dataset.shuffle(seed=SEED).select(range(700)),
    eval_dataset = encoded_eval_dataset.shuffle(seed=SEED).select(range(100)),
    model_init = model_init,
    compute_metrics=compute_metrics,
)

# Create Scheduler to use for Hyperparameter Search
scheduler = PopulationBasedTraining(
        time_attr="training_iteration",
        metric='eval_rouge1',
        mode="max",
        perturbation_interval=1,
        hyperparam_mutations={
            "learning_rate": tune.loguniform(1e-6, 1e-2),
            "per_device_train_batch_size": [1],
            "per_device_eval_batch_size": [1],
            "gradient_accumulation_steps": [4],
            "eval_accumulation_steps": [4],
        })

tune_config = {
    "seed": tune.choice(list(range(1, 42))),
    "num_train_epochs": tune.grid_search([5]),
}


# ## Run Hyperparameter Search

# In[ ]:


# Run Hyperparameter Search and save best trial
best_trial = trainer.hyperparameter_search(
    hp_space = lambda _: tune_config,
    direction="maximize",
    backend="ray",
    keep_checkpoints_num=1,
    checkpoint_score_attr='eval_rouge1',
    local_dir = 'hyperparam_results',
    raise_on_failed_trial = False,
    resources_per_trial = {"cpu": 56, "gpu": 4}, 
    n_trials=5,  
    scheduler=scheduler,
)

Here are hopefully relevant snipits from the stdout file:

image1808×551 21.6 KB


image1880×547 22 KB


image1897×538 22 KB

…

image1781×535 21.4 KB

Along with some snipits from the .err file:

image1918×697 65.7 KB


image1892×468 43.7 KB


image1879×460 43.6 KB

Just as reference, I was able to run the exact same code with the BART-base model successfully. I’m worried if this has something to do with memory requirements.