(DDP) RoBERTa_large training with `torch.compile` results in OOM and other issues
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

Hi,

I'm working with Jan 28 nightly build of PyTorch (nightly branch)
https://github.com/pytorch/pytorch/commit/5d6a4f697cac34d15262aad8afab096170d29ce1

RoBERTa architecture here: https://arxiv.org/pdf/1907.11692.pdf

The model definition && training scripts come from Fairseq
1. Trainer: https://github.com/facebookresearch/fairseq/blob/main/fairseq/trainer.py
2. Model def (RoBERTa Large): https://github.com/facebookresearch/fairseq/blob/main/fairseq/models/roberta/model.py#L34

The script was adapted by wrapping the `wrapped_model` in `trainer.py`

```
self._wrapped_model = torch.compile(self._wrapped_model)
```

I'm using the wiki text dataset: https://huggingface.co/datasets/wikitext

I'm working on an EC2 setup with `p4d.24xlarge` instance. 
Find instance specification here: https://aws.amazon.com/ec2/instance-types/p4/

These are the hyperparameter settings:

```bash
fairseq-train wikitext --adam-eps 1e-06 --arch roberta_large --attention-dropout 0.1 --clip-norm 0.0 --criterion masked_lm --distributed-backend nccl --distributed-no-spawn --dropout 0.1 --encoder-embed-dim 2048 --encoder-ffn-embed-dim 8192 --encoder-layers 24 --log-format simple --log-interval 10 --lr 0.0001 --lr-scheduler polynomial_decay --max-sentences 8 --max-update 500 --optimizer adam --sample-break-mode complete --skip-invalid-size-inputs-valid-test --task masked_lm --tokens-per-sample 512 --total-num-update 100 --update-freq 1 --weight-decay 0.01 --no-save --memory-efficient-fp16 --skip-invalid-size-inputs-valid-test --no-last-checkpoints
```

Note that to replicate this setup, you will have to rebuild Apex against the latest version of PT, I couldn't do this in the PyTorch nightly images because of CUDA missing dependencies on that container. Perhaps the PyTorch nightly devel container will have these dependencies (will validate later today).

Once that is done, you will have to rebuild Fairseq. 
Fairseq uses quite a few torchscript annotation and conditional code branching such as `if torch.jit.tracing: do_f`, and I've observed that dynamo symbolic converter throws this warning

```
[2023-01-28 17:07:39,262] torch._dynamo.symbolic_convert: [DEBUG] FAILED INLINING <code object forward_scriptable at 0x7fe3ac862f50, file "/fairseq/fairseq/models/transformer/transformer_encoder.py", line 173>
``` 
https://github.com/facebookresearch/fairseq/blob/main/fairseq/models/transformer/transformer_encoder.py#L173

> What is the guidance on using torchscript in conjunction to torch.compile? Is it recommended to not do that? 


However, this won't cause compilation to stop but down the road inductor will throw an error

```
  File "/opt/conda/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 187, in run
    result = launcher(
  File "<string>", line 6, in launcher
ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
```

training works fine runs fine without torch.compile wrap. 

So, I'm not certain if this is an issue specifically with my setup but I can confirm that I can run a few of the dynamo benchmarks successfully. I'll try to get fairseq working on the nightly container

--

Anyways after removing the jit branches in code, the training continues but runs into OOM issues. The /fsx/pytorch-nightly/pytorch-2.0/pytorch/aten/src/ATen/cuda/CUDAGraph.cpp fails to allocate memory.


More questions. 

These warnings: 
/opt/conda/lib/python3.9/site-packages/torch/nn/functional.py:4872: UserWarning: This function is deprecated please rebuild your models with the public version of sdpa.
  warnings.warn("This function is deprecated please rebuild your models with the public version of sdpa.")

[2023-01-31 16:46:36,214] torch._inductor.lowering: [WARNING] using triton random, expect difference from eager

-> What is actionable for these? Specifically, for spda, what is the recommendation? Any particular library that should be used. What is the private (as opposed to public) version of SPDA?

Moreover, I've attached the minified scripts, but they don't reproduce the issue. Do they really help you as a domain specialist if the problem is not reproducible?

I've not tried running this without the DDP wrapper, but eventually I need to train this model/its variants on a large cluster, so even if the non DDP version works, thats not helpful. 

Moreover, all the artifacts pasted here are collected with `mpirun -np 1 <command>`, so use DDP with 1 rank. This is not relevant, but giving you the full picture. 

### Error logs

```
/opt/conda/lib/python3.9/site-packages/torch/cuda/graphs.py:82: UserWarning: The CUDA Graph is empty. This ususally means that the graph was attempted to be captured on wrong device or stream. (Triggered internally at /fsx/pytorch-nightly/pytorch-2.0/pytorch/aten/src/ATen/cuda/CUDAGraph.cpp:191.)
  super(CUDAGraph, self).capture_end()
2023-01-31 16:46:57 | WARNING | fairseq.trainer | |===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 1            |        cudaMalloc retries: 1         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |  32019 MiB |  32019 MiB |  66253 MiB |  34234 MiB |
|       from large pool |  32007 MiB |  32007 MiB |  65929 MiB |  33922 MiB |
|       from small pool |     12 MiB |     12 MiB |    323 MiB |    311 MiB |
|---------------------------------------------------------------------------|
| Active memory         |  32019 MiB |  32019 MiB |  66253 MiB |  34234 MiB |
|       from large pool |  32007 MiB |  32007 MiB |  65929 MiB |  33922 MiB |
|       from small pool |     12 MiB |     12 MiB |    323 MiB |    311 MiB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |  39034 MiB |  39034 MiB |  57400 MiB |  18366 MiB |
|       from large pool |  38856 MiB |  38856 MiB |  57160 MiB |  18304 MiB |
|       from small pool |    178 MiB |    178 MiB |    240 MiB |     62 MiB |
|---------------------------------------------------------------------------|
| Non-releasable memory |   4116 MiB |   4129 MiB |  15251 MiB |  11135 MiB |
|       from large pool |   3950 MiB |   3962 MiB |  14611 MiB |  10660 MiB |
|       from small pool |    165 MiB |    167 MiB |    640 MiB |    474 MiB |
|---------------------------------------------------------------------------|
| Allocations           |    2847    |    2847    |    4943    |    2096    |
|       from large pool |    1275    |    1275    |    2281    |    1006    |
|       from small pool |    1572    |    1572    |    2662    |    1090    |
|---------------------------------------------------------------------------|
| Active allocs         |    2847    |    2847    |    4943    |    2096    |
|       from large pool |    1275    |    1275    |    2281    |    1006    |
|       from small pool |    1572    |    1572    |    2662    |    1090    |
|---------------------------------------------------------------------------|
| GPU reserved segments |    1211    |    1211    |    1705    |     494    |
|       from large pool |    1122    |    1122    |    1585    |     463    |
|       from small pool |      89    |      89    |     120    |      31    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |     549    |     549    |    1403    |     854    |
|       from large pool |     410    |     410    |     851    |     441    |
|       from small pool |     139    |     139    |     552    |     413    |
|---------------------------------------------------------------------------|
| Oversize allocations  |       0    |       0    |       0    |       0    |
|---------------------------------------------------------------------------|
| Oversize GPU segments |       0    |       0    |       0    |       0    |
|===========================================================================|
2023-01-31 16:46:57 | ERROR | fairseq.trainer | OOM during optimization, irrecoverable
Traceback (most recent call last):
  File "/opt/conda/bin/fairseq-train", line 8, in <module>
    sys.exit(cli_main())
  File "/fsx/roberta/fairseq/fairseq_cli/train.py", line 574, in cli_main
    distributed_utils.call_main(cfg, main)
  File "/fsx/roberta/fairseq/fairseq/distributed/utils.py", line 393, in call_main
    distributed_main(cfg.distributed_training.device_id, main, cfg, kwargs)
  File "/fsx/roberta/fairseq/fairseq/distributed/utils.py", line 366, in distributed_main
    main(cfg, **kwargs)
  File "/fsx/roberta/fairseq/fairseq_cli/train.py", line 205, in main
    valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
  File "/opt/conda/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/fsx/roberta/fairseq/fairseq_cli/train.py", line 331, in train
    log_output = trainer.train_step(samples)
  File "/opt/conda/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/fsx/roberta/fairseq/fairseq/trainer.py", line 1034, in train_step
    raise e
  File "/fsx/roberta/fairseq/fairseq/trainer.py", line 979, in train_step
    self.task.optimizer_step(
  File "/fsx/roberta/fairseq/fairseq/tasks/fairseq_task.py", line 545, in optimizer_step
    optimizer.step()
  File "/fsx/roberta/fairseq/fairseq/optim/fp16_optimizer.py", line 450, in step
    self.wrapped_optimizer.step(closure, groups=groups)
  File "/fsx/roberta/fairseq/fairseq/optim/fairseq_optimizer.py", line 135, in step
    self.optimizer.step(closure)
  File "/opt/conda/lib/python3.9/site-packages/torch/optim/optimizer.py", line 253, in wrapper
    out = func(*args, **kwargs)
  File "/fsx/roberta/fairseq/fairseq/optim/fused_adam.py", line 329, in step
    state["exp_avg_sq"] = torch.zeros_like(
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 39.59 GiB total capacity; 31.27 GiB already allocated; 66.19 MiB free; 38.12 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

### Minified repro

Details too long to add in this comment. Repro details after dynamo/AOT are here: https://gist.github.com/0x6b64/966c2900fc7609d8bb9eaec22f8d6cc0

### Versions

```bash
Collecting environment information...
PyTorch version: 2.0.0a0+git5876d91
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.24.3
Libc version: glibc-2.31

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:58:50)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-4.14.296-222.539.amzn2.x86_64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: NVIDIA A100-SXM4-40GB 

Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] bert-pytorch==0.0.1a4
[pip3] clip-anytorch==2.5.0
[pip3] CoCa-pytorch==0.0.7
[pip3] dalle2-pytorch==1.10.5
[pip3] ema-pytorch==0.1.4
[pip3] functorch==1.14.0a0+408bcf1
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.5
[pip3] pytorch-transformers==1.2.0
[pip3] pytorch-warmup==0.1.1
[pip3] rotary-embedding-torch==0.2.1
[pip3] torch==2.0.0a0+git5876d91
[pip3] torch-fidelity==0.3.0
[pip3] torch-struct==0.5
[pip3] torchaudio==2.0.0a0+4699ef2
[pip3] torchdata==0.6.0a0+a1612ee
[pip3] torchmetrics==0.11.0
[pip3] torchrec-nightly==2023.1.29
[pip3] torchtext==0.15.0a0+f653dac
[pip3] torchvision==0.15.0a0+c35e8d5
[pip3] vector-quantize-pytorch==0.10.15
[conda] bert-pytorch              0.0.1a4                   dev_0    <develop>
[conda] clip-anytorch             2.5.0                    pypi_0    pypi
[conda] coca-pytorch              0.0.7                    pypi_0    pypi
[conda] dalle2-pytorch            1.10.5                   pypi_0    pypi
[conda] ema-pytorch               0.1.4                    pypi_0    pypi
[conda] functorch                 1.14.0a0+408bcf1          pypi_0    pypi
[conda] magma-cuda117             2.6.1                         1    pytorch
[conda] mkl                       2022.2.1         h84fe81f_16997    conda-forge
[conda] mkl-include               2023.0.0         h84fe81f_25396    conda-forge
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] pytorch-transformers      1.2.0                    pypi_0    pypi
[conda] pytorch-warmup            0.1.1                    pypi_0    pypi
[conda] rotary-embedding-torch    0.2.1                    pypi_0    pypi
[conda] torch                     2.0.0a0+git5876d91          pypi_0    pypi
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torch-struct              0.5                      pypi_0    pypi
[conda] torchaudio                2.0.0a0+4699ef2          pypi_0    pypi
[conda] torchdata                 0.6.0a0+a1612ee          pypi_0    pypi
[conda] torchmetrics              0.11.0                   pypi_0    pypi
[conda] torchrec-nightly          2023.1.29                pypi_0    pypi
[conda] torchtext                 0.15.0a0+f653dac          pypi_0    pypi
[conda] torchvision               0.15.0a0+c35e8d5          pypi_0    pypi
[conda] vector-quantize-pytorch   0.10.15                  pypi_0    pypi
```

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh