Bert cannot converge on toy dataset
Author: retepViolet
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.42.4
- Platform: Windows-10-10.0.22631-SP0
- Python version: 3.11.1
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cpu (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: Distributed Computing

### Who can help?

_No response_

### 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


[This is the training data](https://github.com/user-attachments/files/16276857/Genshin.Impact.txt)

Training method:
```
text = open('Genshin Impact.txt').read()
tokens = tokenizer(text)
x = tokens['input_ids']
mask = tokens['attention_mask']
y = tokenizer.onehot(x).flatten(0, 1)

bert = BertForMaskedLM(BertConfig(hidden_size = 512,
                  num_attention_heads = 1, num_hidden_layers = 1,
                  intermediate_size = 512*2,
                  hidden_dropout_prob = 0, attention_probs_dropout_prob = 0))

def train(epoch):
    bert.train()
    optimizer = torch.optim.Adam(bert.parameters(), lr = 0.003)
    loss_fn = nn.CrossEntropyLoss()
    for i in range(epoch):
        pred = bert(input_ids=x)[0].flatten(0, 1)
        loss = loss_fn(pred, y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        if (i+1) % 1 == 0:
            print(f'Epoch: {i+1}; Loss: {loss.item()}')

train(100)
```

### Expected behavior

I tried to train a BertModel (d_model=768, n_layer=1, n_head=1) from scratch using a small dataset (500-word passage). I tasked the model with simply repeating the input tokens, but it never converged, with the Cross Entropy Loss consistently around 4.

I then tried using a smaller d_model, changing it to 512, and the model converged quickly (loss < 0.001). This suggests that the model didn't converge initially because it was too large.

To verify my hypothesis, I wrote the code for a transformer encoder and used the same model parameters I had used for BERT. Surprisingly, it converged quickly. If the transformer encoder and BERT have the same structure, they should behave similarly, right? However, BertModel could not converge but mine could.

Comments:
Hi @retepViolet, thanks for raising an issue! 

This is a question best placed in our [forums](https://discuss.huggingface.co/). We try to reserve the github issues for feature requests and bug reports.

Here's a good guide on how to train a masked language model: https://huggingface.co/learn/nlp-course/chapter7/3?fw=pt

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.