Repository: pytorch
Issue: Mistral works on (2.0.1+cu117) but "CUDA out of memory" on (2.1.2+cu121)
Labels: high priority,module: cuda,module: memory usage,triaged,module: regression

### 🐛 Describe the bug

I pip upgraded torch from 2.0.1 to 2.1.2 and without any code changes I now run out of CUDA memory loading Mistral7B on NVIDIA GeForce RTX 3060 TI.
- From 2.0.1 to 2.1.2. Did memory allocators or max_split_size_mb change? Does torch reserves more memory? 
- Any steps I can do to further debug this issue?

From Traceback:
I see failing to allocate 20MB and 205MB reserved by unused by PyTorch. Does 2.1.2 reserves more memory? "this process has 17179869184.00 GiB memory" is this GiB or bytes? 

```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.set_default_device("cuda:0")

tokenizer = AutoTokenizer.from_pretrained(
    'mistralai/Mistral-7B-v0.1',
    trust_remote_code=True,
)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    'mistralai/Mistral-7B-v0.1',
    trust_remote_code=True,
    torch_dtype=torch.float16,
)

generated_tokens = tokenizer("Tell me a joke", return_tensors="pt")

model_params = {
    'do_sample': True,
    'max_new_tokens': 64,
    'input_ids': generated_tokens.input_ids,
    'attention_mask': generated_tokens.attention_mask,
}

output = model.generate(**model_params)
print(tokenizer.batch_decode(output)[0])

```

Traceback (most recent call last):
  File "/home/bpevangelista/projects/kfastml/test.py", line 12, in <module>
    model = AutoModelForCausalLM.from_pretrained(
  File "/home/bpevangelista/.local/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
    return model_class.from_pretrained(
  File "/home/bpevangelista/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3706, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/bpevangelista/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4091, in _load_pretrained_model
    state_dict = load_state_dict(shard_file)
  File "/home/bpevangelista/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 510, in load_state_dict
    return safe_load_file(checkpoint_file)
  File "/home/bpevangelista/.local/lib/python3.10/site-packages/safetensors/torch.py", line 310, in load_file
    result[k] = f.get_tensor(k)
  File "/home/bpevangelista/.local/lib/python3.10/site-packages/torch/utils/_device.py", line 77, in __torch_function__
    return func(*args, **kwargs)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB. GPU 0 has a total capacty of 8.00 GiB of which 0 bytes is free. Including non-PyTorch memory, this process has 17179869184.00 GiB memory in use. Of the allocated memory 22.36 GiB is allocated by PyTorch, and 205.16 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

### Versions

Collecting environment information...
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.27.2
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060 Ti
Nvidia driver version: 546.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      39 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             16
On-line CPU(s) list:                0-15
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Core(TM) i7-10700KF CPU @ 3.80GHz
CPU family:                         6
Model:                              165
Thread(s) per core:                 2
Core(s) per socket:                 8
Socket(s):                          1
Stepping:                           5
BogoMIPS:                           7583.99
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology cpuid pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi ept vpid ept_ad fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt xsaveopt xsavec xgetbv1 xsaves flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          256 KiB (8 instances)
L1i cache:                          256 KiB (8 instances)
L2 cache:                           2 MiB (8 instances)
L3 cache:                           16 MiB (1 instance)
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Unknown: Dependent on hypervisor status
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] torch==2.1.2
[pip3] torchinfo==1.8.0
[pip3] triton==2.1.0
[conda] Could not collect

cc @ezyang @gchanan @zou3519 @kadeng @ptrblck