Training Loss Sudden Spike After 8 Hours of pre-training a BERT Model
Author: M98M
Pull request: no

### System Info

os: windows server 2019
gpu: A100 40Gb
RAM: 128 GB

python 3.10.11
transformers 4.31.0
pytorch 2.0.1+cu118



### Who can help?

 @muellerz 
 @pacman100

### Information

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

### Tasks

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

### Reproduction

```
config_large = BertConfig(vocab_size=50_000,hidden_size= 1024,  intermediate_size= 4096,
num_attention_heads= 16,  num_hidden_layers= 20,)
model_large = BertForMaskedLM(config=config_large)

training_args = TrainingArguments(
    output_dir="./mymodel",
    overwrite_output_dir=False, # changed
    num_train_epochs=20,
    per_gpu_train_batch_size=16,
    #
    logging_steps = 0.00001,
    save_strategy = 'steps',
    save_steps= 0.008,
    save_total_limit=20,
    evaluation_strategy='steps',
    eval_steps = 0.008,
    #
    learning_rate=5e-5,
    warmup_steps= 20000,
    #
    tf32 = True,
    optim = "adamw_torch_fused",
    group_by_length = True,
    #
    prediction_loss_only=True,
    #
    hub_model_id = 'mymodel',
    push_to_hub = True,
    hub_strategy = 'every_save',
    hub_private_repo = True,
)

trainer = Trainer(
    model=model_large,
    args=training_args,
    data_collator=data_collator_wwm,
    train_dataset=data['train'],
    eval_dataset=data['test']
)
```

### Expected behavior

I am pretraining a bert model. I did it for bert base and results were **very good** and train loss got to less that 2.8 for my dataset (after 5 epochs in 10 days):
![image](https://github.com/huggingface/transformers/assets/29666386/a409c467-93d4-420b-8425-14bbb51a4051)

but when I tried for bert Large the loss gets stuck around 8.0 in less than an hour and didnt change even after 10 hours. (same train args as bert base but batch size is halved to fit)
I tried different things like warmup and different Learning Rates to no avail. The GPU memory usage for bert large was around 39GB (compared to 33GB for bert base with 32 batch size)
![image](https://github.com/huggingface/transformers/assets/29666386/68cf8e1b-83c2-4359-825a-58a2a8648fa0)

So I thought maybe if I make the model a bit smaller it would work. I changed the number of layers to 20 (from 24) and it worked fine for a while but not even midway through the first epoch it spiked back to the loss of 8.0. I have never seen such behavior in transformer models.

![image](https://github.com/huggingface/transformers/assets/29666386/4d349270-f71b-48e4-b3f4-6253e38acab9)

the learning rate graph looks completely normal. this just looks like an infamous case of exploding gradient i suppose, but I haven't heard of them happening in transformer models. 

Comments:
Hey 🤗  thanks for opening an issue! We try to keep the github issues for bugs/feature requests.
Could you ask your question on the [forum](https://discuss.huggingface.co/) instead? I'm sure the community will be of help!

Thanks!

Hello, could you try a lower learning rate as you have decreased the batch size, along with warmup and linear delay for learning rate?

> Hello, could you try a lower learning rate as you have decreased the batch size, along with warmup and linear delay for learning rate?
after some investigation, it seems that it's vanishing gradient problem. (plot of min/avg gradients for each layer after loss jumps)
![image](https://github.com/huggingface/transformers/assets/29666386/1c322d8a-0381-4439-862c-b6afbe896bf4)

while before the model collapses it looks like this:
![image](https://github.com/huggingface/transformers/assets/29666386/b2f4ba73-e690-491e-9fc4-bd893eab6a36)

So I have tried increasing warmup to a much higher steps (10%) and lowered LR to 4e-5. the model is working fine for last couple days, but it's still early to conclude.

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.

@M98M 

> > Hello, could you try a lower learning rate as you have decreased the batch size, along with warmup and linear delay for learning rate?
> > after some investigation, it seems that it's vanishing gradient problem. (plot of min/avg gradients for each layer after loss jumps)
> > ![image](https://user-images.githubusercontent.com/29666386/271579835-1c322d8a-0381-4439-862c-b6afbe896bf4.png)
> 
> while before the model collapses it looks like this: ![image](https://user-images.githubusercontent.com/29666386/271580617-b2f4ba73-e690-491e-9fc4-bd893eab6a36.png)
> 
> So I have tried increasing warmup to a much higher steps (10%) and lowered LR to 4e-5. the model is working fine for last couple days, but it's still early to conclude.

hi @M98M, have u solved the problem after longered the warmup phase and lowered the LR?

Hi @keyu-tian. Yes, the longer warmup seems to have fixed it. The training was completed without further issues.

@M98M many thanks for your quick reply and helpful experiences!

Loss spikes can come from high LR, batch variance, optimizer instability, or bad data. I use ZClip, an adaptive gradient clipping method that adjusts clipping bounds dynamically — super effective at smoothing out spikes. https://github.com/bluorion-com/ZClip