resume_from_checkpoint may still fail with auto_find_batch_size
Author: aspears-premierinc
Pull request: no
Labels: name: trainer

name: bug

### System Info

(Possibly related to https://github.com/huggingface/transformers/issues/25956 and https://stackoverflow.com/questions/77404935/peft-lora-trainer-no-executable-batch-size-found)

# Summary
I'm running a finetuning script based off of [this example](https://github.com/huggingface/notebooks/blob/main/sagemaker/05_spot_instances/scripts/train.py), except with `auto_find_batch_size` set to True and with fixed train / eval batch sizes removed. I'm seeing the following (sporadic) behavior when using `resume_from_checkpoint` and `auto_find_batch_size` together in Sagemaker training jobs:

1. Model completes 1 epoch
2. Checkpoint fails to propagate to S3
3. Training restarts from epoch 0 with number of expected steps doubled
4. Behavior continues until we encounter a `No executable batch size found` error and training job fails. 

Any help is much appreciated!

### Who can help?

@muellerz, @pacman100

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Example below for Gemma 7B, but has occurred with Flan T5 and UL2 models as well.

## Requirements File

```
accelerate==0.26.1
bitsandbytes==0.42.0
transformers==4.38.1
datasets==2.16.1
peft==0.6.2
evaluate==0.4.1
nltk==3.8.1
diffusers==0.25.1
trl==0.7.10
flash-attn==2.5.0
huggingface_hub==0.20.3
aws-encryption-sdk>=3.1.0
python-json-logger>=2.0.2
pycryptodome>=3.17
```
## Calling Script - Relevant Snippets
### Args
```
model_output_dir = '/opt/ml/model'
output_dir = '/opt/ml/checkpoints'

hyperparameters = {
    'model_id': base_model_id,  # pre-trained model
    'dataset_id': dataset_id,  # dataset id
    'train_dataset_path': input_paths['train'],  # path to training dataset
    'eval_dataset_path': input_paths.get('fine_tuning_holdout', input_paths.get('validation')),
    'epochs': 2,  # number of training epochs
    'learning_rate': 2e-4,  # learning rate used during training
    'max_seq_length': max_seq_length,
    'quant_type': quant_type,    
    'lora_r': 8,
    'lora_alpha': 16,
    'lora_dropout': 0.1,
    'model_output_dir': model_output_dir,
    'output_dir': output_dir,
}
```

### HuggingFace Estimator
```
huggingface_estimator = HuggingFace(
    entry_point='entry_point.py',  # train script
    source_dir=str(pathlib.Path(__file__).parent.absolute() / 'qlora_causal_lm'),
    # directory which includes all the files needed for training
    instance_type=instance_type,  # instances type used for the training job
    instance_count=1,  # the number of instances used for training
    base_job_name=job_name,  # the name of the training job
    role=AWS_ROLE_MODELING,  # Iam role used in training job to access AWS ressources, e.g. S3
    volume_size=300,  # the size of the EBS volume in GB
    transformers_version='4.28.1',  # the transformers version used in the training job
    pytorch_version='2.0.0',  # the pytorch_version version used in the training job
    py_version='py310',  # the python version used in the training job
    output_dir=model_output_dir,  # the output path in which the trained model will be saved
    metric_definitions=metric_definitions,
    hyperparameters=hyperparameters,  # the hyperparameters passed to the training job
    environment=env,
    use_spot_instances=True,
    max_wait=90000,
    max_run=86400,
    checkpoint_s3_uri=f's3://{SESS.default_bucket()}/checkpoints/{job_name}',
    checkpoint_local_path=output_dir,
)
```

## Entry Point Script - Relevant Snippets
```
# Define LoRA Config
peft_config = LoraConfig(
    r=args.lora_r,  # the dimension of the low-rank matrices - the higher r the more trainable parameters
    lora_alpha=args.lora_alpha,  # The alpha parameter for Lora scaling.
    lora_dropout=args.lora_dropout,  # The dropout probability for Lora layers.
    target_modules=target_modules,
    bias="none",  # Bias type for Lora. Can be 'none', 'all' or 'lora_only'
    task_type=TaskType.CAUSAL_LM
)

# prepare model for training
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)

model.print_trainable_parameters()

train_args = TrainingArguments(
    output_dir=args.output_dir,
    auto_find_batch_size=True,
    bf16=args.bf16,  # Use BF16 if available
    learning_rate=args.learning_rate,
    num_train_epochs=args.epochs,
    gradient_checkpointing=args.gradient_checkpointing,
    logging_dir=f"{args.output_data_dir}/logs",
    logging_strategy="steps",
    logging_steps=50,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    load_best_model_at_end=True,
    push_to_hub=False,
)

```
```
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    max_seq_length=args.max_seq_length,
    formatting_func=prompt_formatter_with_decryption,
    data_collator=collator,
    args=train_args,
)

# train model
if get_last_checkpoint(args.output_dir) is not None:
    logger.info("***** Continuing training from last checkpoint *****")
    last_checkpoint = get_last_checkpoint(args.output_dir)
    trainer.train(resume_from_checkpoint=last_checkpoint)
else:
    logger.info("***** Training new model *****")
    trainer.train()
```

### Expected behavior

- Checkpoint should propagate to S3 after each epoch
- Training should continue after epoch 1, rather than starting over
- Number of expected steps should not continue to double indefinitely

Comments:
Thanks for raising @aspears-premierinc! 

I think you've tagged the wrong Mueller cc @muellerzr  

> Thanks for raising @aspears-premierinc!
> 
> I think you've tagged the wrong Mueller cc @muellerzr

Ahh, my bad! Thanks @amyeroberts 👍 

> Checkpoint fails to propagate to S3

I would like to believe that this is the sole issue here, however the fact it fails on T5 is very, very odd. Can you explain more here? I'm unfamiliar with sagemaker but if we're relying on a checkpoint to resume from and one does not exist... shouldn't it just break from the start this way? 😅 

What would be helpful is if you could provide the output of the `trainer_state.json` here please!



> > Checkpoint fails to propagate to S3
> 
> I would like to believe that this is the sole issue here, however the fact it fails on T5 is very, very odd. Can you explain more here? I'm unfamiliar with sagemaker but if we're relying on a checkpoint to resume from and one does not exist... shouldn't it just break from the start this way? 😅
> 
> What would be helpful is if you could provide the output of the `trainer_state.json` here please!


Thanks for taking the time to reply. Yes, it's very odd and I assume it's a very strange edge case.

So, there's a failsafe in my code to ensure it starts from epoch 0 if no checkpoint is present 👍 That's why it doesn't break immediately.

```
# train model
if get_last_checkpoint(args.output_dir) is not None:
    logger.info("***** Continuing training from last checkpoint *****")
    last_checkpoint = get_last_checkpoint(args.output_dir)
    trainer.train(resume_from_checkpoint=last_checkpoint)
else:
    logger.info("***** Training new model *****")
    trainer.train()
```

I can't hand over the `trainer_state.json` for any training jobs that encountered this issue because the trainer state is lost when the Sagemaker job fails. But here's a description of a job that failed in the same way for T5 XXL:

**Data**
- 100 samples for train (purposefully very small, to try to reproduce the issue)
- 100 for validation

**Params**
```
...
per_device_eval_batch_size=32,  # Extremely large, to try to reproduce the issue
per_device_train_batch_size=32,
auto_find_batch_size=True
...
```
Note this is slightly different than my original example, as I manually set the `per_device` batch sizes in addition to enabling `auto_find_batch_size`.

Makes me think that `auto_find_batch_size` may be picking a batch that's too large in relation to the number of remaining steps? 



Just chiming in to say looks like other folks have run into similar issues within the past month:

- https://stackoverflow.com/questions/77334292/perform-peft-with-lora-on-flan-t5-model-causing-no-executable-batch-size-error
- https://community.deeplearning.ai/t/no-executable-batch-size-found/482838

Seems like PEFT + LORA + `auto_find_batch_size=True` is the common denominator

Gentle ping @muellerzr @pacman100 

Another ping @muellerzr @pacman100 

cc @muellerzr @SunMarc 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.