Repository: pytorch
Issue: `torch.compile` crash while trying to remove a node
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

I am trying to compile a model where part of the computed tensors are discarded. Eager mode runs successfully while the compiled model crashes.

Removing the accuracy calculation line from the `train_step` function avoids the crash, see the workaround below:

```python
def accuracy(out, labels):
    assert out.ndim == 2
    assert out.size(0) == labels.size(0)
    assert labels.ndim == 1 or (labels.ndim == 2 and labels.size(1) == 1)
    labels = labels.flatten()
    predictions = torch.argmax(out, 1)
    return (labels == predictions).sum(dtype=torch.float64) / labels.size(0)

@torch.compile
def train_step(minibatch, optimizer, model, loss_fn):
    category = "paper"
    node_features = {
        ntype: feat.float()
        for (ntype, name), feat in minibatch.node_features.items()
        if name == "feat"
    }
    labels = minibatch.labels[category].long()
    optimizer.zero_grad()
    out = model(minibatch.sampled_subgraphs, node_features)[category]
    loss = loss_fn(out, labels)
    # https://github.com/pytorch/pytorch/issues/133942
    # num_correct = accuracy(out, labels) * labels.size(0)
    num_correct = torch.zeros(1, dtype=torch.float64, device=out.device)
    loss.backward()
    optimizer.step()
    return loss.detach(), num_correct, labels.size(0)
```

### Error logs

```
root@a100cse:/localscratch/dgl-3/examples/graphbolt/pyg/labor# TORCHDYNAMO_VERBOSE=1 python ../hetero/node_classification.py 
Training in pinned-pinned-cuda mode.
The dataset is already preprocessed.
Dataset loaded
/localscratch/dgl-3/python/dgl/graphbolt/impl/torch_based_feature_store.py:524: GBWarning: `DiskBasedFeature.pin_memory_()` is not supported. Leaving unmodified.
  gb_warning(
Number of model parameters: 2412491
Training: 0it [00:06, ?it/s]
Traceback (most recent call last):
  File "/localscratch/dgl-3/examples/graphbolt/pyg/labor/../hetero/node_classification.py", line 430, in <module>
    main()
  File "/localscratch/dgl-3/examples/graphbolt/pyg/labor/../hetero/node_classification.py", line 373, in main
    train(train_dataloader, valid_dataloader, model, args.device)
  File "/localscratch/dgl-3/examples/graphbolt/pyg/labor/../hetero/node_classification.py", line 298, in train
    train_loss, train_acc, duration = train_helper(
  File "/localscratch/dgl-3/examples/graphbolt/pyg/labor/../hetero/node_classification.py", line 281, in train_helper
    loss, num_correct, num_samples = train_step(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 421, in _fn
    return fn(*args, **kwargs)
  File "/localscratch/dgl-3/examples/graphbolt/pyg/labor/../hetero/node_classification.py", line 265, in train_step
    optimizer.zero_grad()
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 1078, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 919, in _convert_frame
    result = inner_convert(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 456, in _convert_frame_assert
    return _compile(
  File "/usr/local/lib/python3.10/dist-packages/torch/_utils_internal.py", line 83, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/usr/local/lib/python3.10/dist-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 799, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 232, in time_wrapper
    r = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 618, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/bytecode_transformation.py", line 1184, in transform_code_object
    transformations(instructions, code_options)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 177, in _fn
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 564, in transform
    tracer.run()
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 2450, in run
    super().run()
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 892, in run
    while self.step():
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 804, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 496, in wrapper
    return handle_graph_break(self, inst, speculation.reason)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 565, in handle_graph_break
    self.output.compile_subgraph(self, reason=reason)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/output_graph.py", line 1122, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/output_graph.py", line 1314, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 232, in time_wrapper
    r = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/output_graph.py", line 1405, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/output_graph.py", line 1386, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/repro/after_dynamo.py", line 127, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/usr/local/lib/python3.10/dist-packages/torch/__init__.py", line 1944, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/compile_fx.py", line 1474, in compile_fx
    return aot_autograd(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/backends/common.py", line 65, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/aot_autograd.py", line 945, in aot_module_simplified
    compiled_fn, _ = create_aot_dispatcher_function(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 232, in time_wrapper
    r = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/aot_autograd.py", line 678, in create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 280, in aot_dispatch_autograd
    fw_module, bw_module = aot_config.partition_fn(
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/compile_fx.py", line 1406, in partition_fn
    _recursive_joint_graph_passes(graph)
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/compile_fx.py", line 241, in _recursive_joint_graph_passes
    joint_graph_passes(gm)
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/fx_passes/joint_graph.py", line 321, in joint_graph_passes
    count += patterns.apply(graph.graph)  # type: ignore[arg-type]
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/pattern_matcher.py", line 1696, in apply
    entry.apply(m, graph, node)  # type: ignore[arg-type]
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/pattern_matcher.py", line 1008, in apply
    self.handler(match, *match.args, **match.kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/fx_passes/joint_graph.py", line 465, in mul_softmax_pattern
    match.replace_by_example(repl, [inp, other])
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/pattern_matcher.py", line 239, in replace_by_example
    ReplacementPatternEntry.replace_with_graph(
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/pattern_matcher.py", line 1163, in replace_with_graph
    match.erase_nodes(graph)
  File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/pattern_matcher.py", line 204, in erase_nodes
    graph.erase_node(n)
  File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph.py", line 1021, in erase_node
    raise RuntimeError(f'Tried to erase Node {to_erase} but it still had {len(to_erase.users)} '
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: Tried to erase Node mul_9 but it still had 1 users in the graph: {argmax: None}!


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

### Minified repro

Run https://github.com/dmlc/dgl/pull/7722 with after installing pytorch geometric and dgl.

```python

from math import inf
import torch
from torch import tensor, device
import torch.fx as fx
import torch._dynamo
from torch._dynamo.testing import rand_strided
from torch._dynamo.debug_utils import run_fwd_maybe_bwd

import torch._dynamo.config
import torch._inductor.config
import torch._functorch.config
import torch.fx.experimental._config









from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.L__self___L__model___layers_2_dropout = Dropout(p=0.5, inplace=False)
        self.L__self___loss_fn = CrossEntropyLoss()



    def forward(self, L_L_labels_ : torch.Tensor, x_7, labels):
        l_l_labels_ = L_L_labels_
        out_30 = self.L__self___L__model___layers_2_dropout(x_7);  x_7 = None
        loss = self.L__self___loss_fn(out_30, l_l_labels_);  l_l_labels_ = None
        predictions = torch.argmax(out_30, 1);  out_30 = None
        eq = labels == predictions;  labels = predictions = None
        sum_1 = eq.sum(dtype = torch.float64);  eq = None
        truediv_15 = sum_1 / 1024;  sum_1 = None
        num_correct = truediv_15 * 1024;  truediv_15 = None
        return [loss, num_correct]


mod = Repro()

def load_args(reader):
    buf0 = reader.storage('a1b9334d88e7a05117e3540e1976a55e42a3f277', 8192, device=device(type='cuda', index=0), dtype_hint=torch.int64)
    reader.tensor(buf0, (1024,), dtype=torch.int64, is_leaf=True)  # L_L_labels_
    buf1 = reader.storage('3a396a0bd2d1e5fc6bb833c38aef3e0dd147a109', 626688, device=device(type='cuda', index=0))
    reader.tensor(buf1, (1024, 153), requires_grad=True)  # x_7
    reader.tensor(buf0, (1024,), dtype=torch.int64, is_leaf=True)  # labels
load_args._version = 0

if __name__ == '__main__':
    from torch._dynamo.repro.after_dynamo import run_repro
    run_repro(mod, load_args, accuracy=False, command='run',
        save_dir='/localscratch/dgl-3/examples/graphbolt/pyg/labor/checkpoints', autocast=False, backend='inductor')
```

### Versions

Collecting environment information...
PyTorch version: 2.4.0a0+3bcc3cddb5.nv24.07
Is debug build: False
CUDA used to build PyTorch: 12.5
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.30.0
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-117-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.5.82
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

Nvidia driver version: 550.54.15
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.2.1
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:                        43 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               256
On-line CPU(s) list:                  0-255
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7702 64-Core Processor
CPU family:                           23
Model:                                49
Thread(s) per core:                   2
Core(s) per socket:                   64
Socket(s):                            2
Stepping:                             0
Frequency boost:                      enabled
CPU max MHz:                          2183.5930
CPU min MHz:                          1500.0000
BogoMIPS:                             4000.13
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 rapl 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 ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es
Virtualization:                       AMD-V
L1d cache:                            4 MiB (128 instances)
L1i cache:                            4 MiB (128 instances)
L2 cache:                             64 MiB (128 instances)
L3 cache:                             512 MiB (32 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-63,128-191
NUMA node1 CPU(s):                    64-127,192-255
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 Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow:   Mitigation; safe RET
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; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] onnx==1.16.0
[pip3] optree==0.12.1
[pip3] pytorch-triton==3.0.0+989adb9a2
[pip3] torch==2.4.0a0+3bcc3cddb5.nv24.7
[pip3] torch_geometric==2.5.3
[pip3] torch-tensorrt==2.5.0a0
[pip3] torchdata==0.7.1
[pip3] torchmetrics==1.4.1
[pip3] torchvision==0.19.0a0
[conda] Could not collect

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire