Repository: pytorch
Issue: Onnx export memory leak with qlora+deepspeed stage 2+ORT
Labels: module: onnx,triaged

### 🐛 Describe the bug

In Qlora+Deepspeed stage2+ORT profile, observe memory not released after export
![image](https://github.com/pytorch/pytorch/assets/47403504/1f1cd7ec-def4-4ed7-aa0d-88133725f0a2)
Before export 8.49G
After export 12.48G

Note that I did not see memory leak with lora+deepspeed stage2+ORT. Pytorch and ORT memory usage is on par for lora+deepspeed stage2+ORT.

I observed this with Gemma but it's probably universal for all models

Repro info
[run_clm.zip](https://github.com/pytorch/pytorch/files/14675818/run_clm.zip)
export ORTMODULE_USE_EFFICIENT_ATTENTION=1
export ORTMODULE_ONNX_OPSET_VERSION=15
export ORTMODULE_FALLBACK_POLICY=FALLBACK_DISABLE
export APPLY_4BIT=true
export PYTORCH_NO_CUDA_MEMORY_CACHING=1
export ORTMODULE_DEEPCOPY_BEFORE_MODEL_EXPORT=0
export APPLY_ORT=true
nsys profile --trace cuda,nvtx,osrt,cublas,cudnn --force-overwrite true -s cpu -b fp --cudabacktrace all --cuda-memory-usage true --stats false  -o profile_clm_qlora_mem_ort_nomc_nodc_1%p --capture-range=cudaProfilerApi --capture-range-end=stop-shutdown \
python launcher.py gmon torchrun --nproc_per_node 8 run_clm.py --model_name_or_path google/gemma-7b-it --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --do_train --save_strategy 'no' --fp16 --block_size 256 --gradient_accumulation_steps 1 --per_device_train_batch_size 1 --num_train_epochs 5 --output_dir output_dir --overwrite_output_dir --deepspeed zero_stage_2.json \


### Versions

Collecting environment information...
PyTorch version: 2.3.0.dev20240215+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.26.0
Libc version: glibc-2.31

Python version: 3.8.18 (default, Sep 11 2023, 13:40:15)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.138.1-4.cm2-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 525.85.12
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.0
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
Byte Order:                         Little Endian
Address sizes:                      48 bits physical, 48 bits virtual
CPU(s):                             96
On-line CPU(s) list:                0-95
Thread(s) per core:                 1
Core(s) per socket:                 48
Socket(s):                          2
NUMA node(s):                       4
Vendor ID:                          AuthenticAMD
CPU family:                         23
Model:                              49
Model name:                         AMD EPYC 7V12 64-Core Processor
Stepping:                           0
CPU MHz:                            2445.434
BogoMIPS:                           4890.86
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          3 MiB
L1i cache:                          3 MiB
L2 cache:                           48 MiB
L3 cache:                           384 MiB
NUMA node0 CPU(s):                  0-23
NUMA node1 CPU(s):                  24-47
NUMA node2 CPU(s):                  48-71
NUMA node3 CPU(s):                  72-95
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; untrained return thunk; SMT disabled
Vulnerability Spec rstack overflow: Mitigation; safe RET, no microcode
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; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru arat umip rdpid

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.22.2
[pip3] onnx==1.11.0
[pip3] onnxruntime-training==1.18.0.dev20240215001+cu118
[pip3] pytorch-lightning==1.9.3
[pip3] pytorch-triton==3.0.0+901819d2b6
[pip3] torch==2.3.0.dev20240215+cu118
[pip3] torch-nebula==0.16.7
[pip3] torch-ort==1.16.0
[pip3] torchaudio==2.2.0.dev20240215+cu118
[pip3] torchmetrics==0.11.3
[pip3] torchsnapshot==0.1.0
[pip3] torchvision==0.18.0.dev20240215+cu118
[pip3] triton==2.2.0
[conda] magma-cuda118             2.6.1                         1    pytorch
[conda] mkl                       2021.4.0                 pypi_0    pypi
[conda] mkl-include               2021.4.0                 pypi_0    pypi
[conda] numpy                     1.22.2                   pypi_0    pypi
[conda] pytorch-lightning         1.9.3                    pypi_0    pypi
[conda] pytorch-triton            3.0.0+901819d2b6          pypi_0    pypi
[conda] torch                     2.3.0.dev20240215+cu118          pypi_0    pypi
[conda] torch-nebula              0.16.7                   pypi_0    pypi
[conda] torch-ort                 1.16.0                   pypi_0    pypi
[conda] torchaudio                2.2.0.dev20240215+cu118          pypi_0    pypi
[conda] torchmetrics              0.11.3                   pypi_0    pypi
[conda] torchsnapshot             0.1.0                    pypi_0    pypi
[conda] torchvision               0.18.0.dev20240215+cu118          pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi