[Activation Memory] remains unknown
Author: AaronZLT
Pull request: no
Labels: name: Performance

name: bug

name: Accelerate

### System Info


Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.44.0
- Platform: Linux-5.4.0-165-generic-x86_64-with-glibc2.31
- Python version: 3.10.14
- Huggingface_hub version: 0.24.5
- Safetensors version: 0.4.3
- Accelerate version: 0.31.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.4.0+cu121 (True)
- 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?: <NA>
- Using GPU in script?: <1x A800>
- GPU type: NVIDIA A800-SXM4-80GB



### Who can help?

## issue 1 - Brief
- Training GPT2-XL with 🤗 Transformers and 🤗 Accelerate
- Using _record_memory_history and _dump_snapshot from PyTorch for profiler
- The activation memory are always remain unknown

Training GPT2-XL with 🤗 Transformers, set optimizer to adamw_torch. However, whether mix-precision training or not, from the snapshot, the Activation are always unknown.

Reproduce

model: GPT2-XL
batch size: 4
sequence length: 1024
single GPU
optimizer: adamw_torch
bf16

```bash
# example command
accelerate launch benchmark.py --dataset alpaca --output_dir validate --logging_strategy steps --logging_steps 1 --save_strategy no --save_steps 100 --dataloader_num_workers 1 --remove_unused_columns False --do_train --ddp_find_unused_parameters False --overwrite_output_dir --profiler pytorch --profiler_warmup_step 3 --max_steps 5 --max_memory_MB 80000 --hard_padding True --save_total_limit 0 --num_train_epochs 2 --learning_rate 4e-4 --optim adamw_torch --per_device_train_batch_size 4 --source_max_len 512 --target_max_len 512
```

## issue 2 - brief
In the training above, Transformers v4.41.2 yields quite different memory overhead compared to v4.42.0+, the memory overhead with Transformers v4.41.2- are always higher. Is there any fix in it, what 'actiavtions' are erased in v4.42.0+?

![Trace_gpt2-xl-dataset_alpaca-gpus1-bs4-seq1024-lr0 0004-linear---------bf16----_step_3_to_5](https://github.com/user-attachments/assets/0d9580d9-e83b-4efc-b502-35aeb23e64b7)
![Trace_gpt2-xl-dataset_alpaca-gpus1-bs4-seq1024-lr0 0002-linear---------bf16----_step_3_to_5](https://github.com/user-attachments/assets/dda1e601-a228-4264-af4c-c793b47c4c66)

@ArthurZucker @muellerzr @SunMarc

### Information

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

### Tasks

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

### Reproduction

-

### Expected behavior

-

Comments:
Hey! Not sure this is a bug as the latest versions of `transformers` have less memory issue  AFAIU your comment no? 
Top of my head, this might correspond to SDPA support: 
<img width="752" alt="image" src="https://github.com/user-attachments/assets/db4851c7-5e7b-4ee8-882d-19c50dfa8765">



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.