During pre-training, using FA2 consumes more memory than using SDPA

As described in the title

When performing pre-training, using FA2 will consume more GPU memory than using SDPA.

I am using the trainer from transformers for training, and the simplified code is roughly as follows:

```python
model = AutoModelForMaskedLM.from_pretrained(
MODEL_PATH,
torch_dtype = torch.bfloat16 if is_torch_bf16_gpu_available() == True else torch.float16,
attn_implementation = "flash_attention_2" # spda
).to("cuda" if torch.cuda.is_available() else "cpu")

training_args = TrainingArguments(
bf16 = True,
optim = "paged_adamw_8bit",
warmup_ratio = 0.1,
weight_decay = 5e-5,
learning_rate = 5e-5,
num_train_epochs = 1,
per_device_eval_batch_size = 16,
per_device_train_batch_size = 8,
gradient_checkpointing = False,
)

trainer = Trainer(
args = training_args,
model = model,
data_collator = DataCollatorForLanguageModeling(
tokenizer = tokenizer,
mlm = True,
mlm_probability = 0.30,
pad_to_multiple_of = 8,
),
eval_dataset = eval_dataset,
train_dataset = train_dataset,
processing_class = tokenizer,
)
```

When all other parameters are kept consistent and only the `attn_implementation` is changed, the GPU memory usage rates are 48% and 88%, respectively.

When using FA2, the GPU memory usage is significantly higher than with SDPA and also much higher than with other traditional Bert-Like models, and there is no improvement in speed.

The same phenomenon has been observed on both Windows 11 24H2 and Ubuntu@WSL2.

ENVS：
PyTorch 2.5.1
Python 3.12.8
flash_attn v2.7.2.post1


No repository file or test is named. Start by reproducing the provided transformers Trainer setup and changing only attn_implementation between flash_attention_2 and SDPA under the stated environment, then compare GPU memory and throughput. Done means the discrepancy is explained and the resulting behavior is verified.

Written by the indexing model from the issue text.