Repository: pytorch
Issue: `torch.jit.trace` memory usage increase although forward is constant, and gets much slower than forward with model depth increase
Labels: oncall: jit

### 🐛 Describe the bug

Using `traced_model = torch.jit.trace(model, example_inputs)`, memory usage is increasing over the model depth although with `result = model(example_inputs)` the memory usage is constant.

The model in question has a loop construct where a submodule is called iteratively.

Note that for both running `forward` and `torch.jit.trace`, the decorator `torch.inference_mode()` is used.

The submodule inside the loop is a unet, and it is already a `ScriptModule` when calling `torch.jit.trace` on the whole model. Not tracing it in advance does not change the results.

I notice that `torch.jit.trace` alternates using all the physical CPU cores, and using only one CPU core. Not sure why.

The results are:

<details>
  <summary>Slowdown of `torch.jit.trace`</summary>

![image](https://user-images.githubusercontent.com/9808326/216381415-13a92b88-c71c-4479-88bd-bc858a47414a.png)

</details>


<details>
  <summary>Memory usage of `torch.jit.trace`, with the inner unet pretraced</summary>

* n_loop = 30: max allocation `5270 MB`, `max_mem/n_loop = 175.66`
* n_loop = 40: max allocation `6592 MB`, `max_mem/n_loop = 164.8`
* n_loop = 50: max allocation `8084 MB`, `max_mem/n_loop = 161.68`
* n_loop = 60: max allocation `9532 MB`, `max_mem/n_loop = 158.86`
* n_loop = 120: max allocation `18181 MB`, `max_mem/n_loop = 151.50`
</details>

The loop looks like this: https://github.com/fxmarty/optimum/blob/9ef352b0f1775fe5e0048490722ea931aa6f4fd5/research/stable_diffusion_end_to_end_onnx/scriptable_pipeline_stable_diffusion.py#L298-L318

The script to run the forward is: https://github.com/fxmarty/optimum/blob/9ef352b0f1775fe5e0048490722ea931aa6f4fd5/research/stable_diffusion_end_to_end_onnx/run_custom_pytorch_pipeline.py

The script to run `torch.jit.trace` is: https://github.com/fxmarty/optimum/blob/9ef352b0f1775fe5e0048490722ea931aa6f4fd5/research/stable_diffusion_end_to_end_onnx/create_scriptmodule_by_tracing.py

### Versions

PyTorch version: 1.13.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: 14.0.0-1ubuntu1
CMake version: version 3.25.0
Libc version: glibc-2.35

Python version: 3.9.12 (main, Apr  5 2022, 06:56:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.15.0-58-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060 Laptop GPU
Nvidia driver version: 515.86.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn.so.8.7.0
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.7.0
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.7.0
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.7.0
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.7.0
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.7.0
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.7.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] torch==1.13.1
[pip3] torch-model-archiver==0.6.1
[pip3] torch-workflow-archiver==0.2.5
[pip3] torchinfo==1.7.0
[pip3] torchserve==0.6.1
[pip3] torchtriton==2.0.0+0d7e753227
[pip3] torchvision==0.14.1
[conda] cudatoolkit               11.3.1               h2bc3f7f_2    anaconda
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] torch                     1.13.1                   pypi_0    pypi
[conda] torch-model-archiver      0.6.1                    pypi_0    pypi
[conda] torch-workflow-archiver   0.2.5                    pypi_0    pypi
[conda] torchinfo                 1.7.0                    pypi_0    pypi
[conda] torchserve                0.6.1                    pypi_0    pypi
[conda] torchtriton               2.0.0+0d7e753227          pypi_0    pypi
[conda] torchvision               0.14.1                   pypi_0    pypi


cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel