Errors using torch.compile() on Megatron-LM GPT model 
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

enable torch2 on megatron-lm get model
![image](https://user-images.githubusercontent.com/42488903/226274021-4617261e-cd4d-4e9f-902c-1cda5c78a0d9.png)

```
def model_provider(pre_process=True, post_process=True):
    """Build the model."""
    print_rank_0('building GPT model ...')
    model = GPTModel(
        num_tokentypes=0,
        parallel_output=True,
        pre_process=pre_process,
        post_process=post_process
    )
    model = torch.compile(model)
    return model
```

training script like:
```
#! /bin/bash

# Runs the "345M" parameter model

GPUS_PER_NODE=4
# Change for multinode config
NNODES=2
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

DATA_PATH=/share/project/Megatron-DeepSpeed-data/data/meg-gpt2-oscar-en-10k_text_document
CHECKPOINT_PATH=checkpoints/gpt2

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
torchrun $DISTRIBUTED_ARGS \
       pretrain_gpt.py \
       --tensor-model-parallel-size 2 \
       --pipeline-model-parallel-size 2 \
       --no-async-tensor-model-parallel-allreduce \
       --num-layers 24 \
       --hidden-size 1024 \
       --num-attention-heads 16 \
       --micro-batch-size 4 \
       --global-batch-size 16 \
       --seq-length 1024 \
       --max-position-embeddings 1024 \
       --train-iters 500000 \
       --lr-decay-iters 320000 \
       --save $CHECKPOINT_PATH \
       --load $CHECKPOINT_PATH \
       --data-path $DATA_PATH \
       --vocab-file /share/project/Megatron-DeepSpeed-data/data/gpt2-vocab.json \
       --merge-file /share/project/Megatron-DeepSpeed-data/data/gpt2-merges.txt \
       --data-impl mmap \
       --split 949,50,1 \
       --distributed-backend nccl \
       --lr 0.00015 \
       --lr-decay-style cosine \
       --min-lr 1.0e-5 \
       --weight-decay 1e-2 \
       --clip-grad 1.0 \
       --lr-warmup-fraction .01 \
       --log-interval 100 \
       --save-interval 10000 \
       --eval-interval 1000 \
       --eval-iters 10 \
       --fp16
```


Got such error:
```
[after dataloaders are built] datetime: 2023-03-20 07:28:59
done with setup ...
training ...
[before the start of training step] datetime: 2023-03-20 07:28:59
[2023-03-20 07:29:05,653] torch._inductor.graph: [ERROR] Error from lowering
Traceback (most recent call last):
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/graph.py", line 333, in call_function
    out = lowerings[target](*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 2020, in index_put
    return index_put_(clone(x), indices, values, accumulate)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 2044, in index_put_
    return index_put_as_masked_fill(self, indices, values, accumulate)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 2028, in index_put_as_masked_fill
    return mutate_to(self, where(indices[0], value, self))
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 454, in where
    for i, x in zip(indices, broadcast_tensors(*[args[i] for i in indices])):
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 468, in broadcast_tensors
    target = functools.reduce(
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 280, in broadcast_symbolic_shapes
    V.graph.sizevars.guard_equals(a, b)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/sizevars.py", line 253, in guard_equals
    assert self.shape_env.evaluate_expr(sympy.Eq(left, right))
AssertionError
[2023-03-20 07:29:05,653] torch._inductor.graph: [ERROR] Error from lowering
Traceback (most recent call last):
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/graph.py", line 333, in call_function
    out = lowerings[target](*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 2020, in index_put
    return index_put_(clone(x), indices, values, accumulate)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 2044, in index_put_
    return index_put_as_masked_fill(self, indices, values, accumulate)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 2028, in index_put_as_masked_fill
    return mutate_to(self, where(indices[0], value, self))
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 454, in where
    for i, x in zip(indices, broadcast_tensors(*[args[i] for i in indices])):
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 225, in wrapped
    out = decomp_fn(*args, **kwargs)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 468, in broadcast_tensors
    target = functools.reduce(
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/lowering.py", line 280, in broadcast_symbolic_shapes
    V.graph.sizevars.guard_equals(a, b)
  File "/home/elrond/code/Megatron-LM/Megatron-LM/.env/lib/python3.8/site-packages/torch/_inductor/sizevars.py", line 253, in guard_equals
    assert self.shape_env.evaluate_expr(sympy.Eq(left, right))
AssertionError
```



### Versions

Collecting environment information...
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 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.26.0
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-113-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB

Nvidia driver version: 470.129.06
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

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          256
On-line CPU(s) list:             0-255
Thread(s) per core:              2
Core(s) per socket:              64
Socket(s):                       2
NUMA node(s):                    8
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD EPYC 7742 64-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         3279.251
CPU max MHz:                     2250.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        4491.31
Virtualization:                  AMD-V
L1d cache:                       4 MiB
L1i cache:                       4 MiB
L2 cache:                        64 MiB
L3 cache:                        512 MiB
NUMA node0 CPU(s):               0-15,128-143
NUMA node1 CPU(s):               16-31,144-159
NUMA node2 CPU(s):               32-47,160-175
NUMA node3 CPU(s):               48-63,176-191
NUMA node4 CPU(s):               64-79,192-207
NUMA node5 CPU(s):               80-95,208-223
NUMA node6 CPU(s):               96-111,224-239
NUMA node7 CPU(s):               112-127,240-255
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          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; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling
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 nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate sme ssbd mba sev ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.24.2
[pip3] torch==2.0.0
[pip3] torchvision==0.15.1
[pip3] triton==2.0.0
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.22.4           py38h99721a1_0    conda-forge
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.13.0a0+d321be6          pypi_0    pypi
[conda] torch-tensorrt            1.2.0a0                  pypi_0    pypi
[conda] torchtext                 0.11.0a0                 pypi_0    pypi
[conda] torchvision               0.14.0a0                 pypi_0    pypi



cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh