Repository: pytorch
Issue: storing std of tensor leads to inplace backwards error
Labels: module: autograd,triaged,oncall: pt2,module: pt2-dispatcher

### 🐛 Describe the bug

This is an extremely simple model that:
1. takes a `[S D]` FloatTensor and an `[L]` Int Tensor index
2. does a linear, stores the std of the result, and then does another linear on the specified indices
3. returns an `[S D]` FloatTensor and does MSE backward with the results (dummy task)

It works correctly in eager and fails on compile.

```
import torch
import torch.nn.functional as F
from torch import Tensor, nn

class L(nn.Module):
    def __init__(self, d: int):
        super().__init__()
        self.a = nn.Linear(d, d, bias=False)
        self.b = nn.Linear(d, d, bias=False)

    def forward(self, x: Tensor, i: Tensor) -> Tensor:
        x = self.a(x)
        self.std = x.std()
        x[i] = self.b(x[i])
        return x


def init_and_backward(compile: bool = False):
    S, D = 5, 8
    # inputs
    x, y = torch.randn(S, D).half(), torch.randn(S, D).half()
    idx = torch.tensor([1, 2], dtype=torch.int64).cuda()

    # model
    m = L(D).half().cuda()
    if compile: m = torch.compile(m)

    # fwd-bck; ensure all grads exist
    F.mse_loss(m(x, idx), y).backward()
    assert all(p.grad is not None for p in m.parameters())

torch.set_default_device("cuda")
init_and_backward()
print('eager worked fine.')
init_and_backward(compile=True)
```

### Error logs

```
$ python3 rep.py
eager worked fine.
/.venv/lib/python3.10/site-packages/torch/autograd/graph.py:768: UserWarning: Error detected in StdBackward0. Traceback of forward call that caused the error:
  File "/rep.py", line 13, in forward
    self.std = x.std()
 (Triggered internally at ../torch/csrc/autograd/python_anomaly_mode.cpp:111.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
Traceback (most recent call last):
  File "/rep.py", line 35, in <module>
    init_and_backward(compile=True)
  File "/rep.py", line 29, in init_and_backward
    F.mse_loss(m(x, idx), y).backward()
  File "/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 433, in _fn
    return fn(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1116, in __call__
    return self._torchdynamo_orig_callable(
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 948, in __call__
    result = self._inner_convert(
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
  File "/.venv/lib/python3.10/site-packages/torch/_utils_internal.py", line 84, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/.venv/lib/python3.10/site-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 "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1185, in transform_code_object
    transformations(instructions, code_options)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 582, in transform
    tracer.run()
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2451, in run
    super().run()
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 893, in run
    while self.step():
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 805, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2642, in RETURN_VALUE
    self._return(inst)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2627, in _return
    self.output.compile_subgraph(
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1123, 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 "/.venv/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1318, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1409, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1390, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/repro/after_dynamo.py", line 129, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/.venv/lib/python3.10/site-packages/torch/__init__.py", line 1951, 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 "/.venv/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1505, in compile_fx
    return aot_autograd(
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/backends/common.py", line 69, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 954, in aot_module_simplified
    compiled_fn, _ = create_aot_dispatcher_function(
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 687, in create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 242, in aot_dispatch_autograd
    fx_g, joint_inputs, maybe_subclass_meta = aot_dispatch_autograd_graph(
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 276, in aot_dispatch_autograd_graph
    fx_g = _create_graph(joint_fn_to_trace, updated_joint_inputs, aot_config=aot_config)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 46, in _create_graph
    fx_g = make_fx(
  File "/.venv/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1421, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/.venv/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1367, in trace
    return self._trace_inner(f, *args)
  File "/.venv/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1354, in _trace_inner
    t = dispatch_trace(
  File "/.venv/lib/python3.10/site-packages/torch/_compile.py", line 31, in inner
    return disable_fn(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 600, in _fn
    return fn(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 642, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 600, in _fn
    return fn(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 822, in trace
    (self.create_arg(fn(*args)),),
  File "/.venv/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 675, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/.venv/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 660, in wrapped
    out = f(*tensors)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 599, in joint_helper
    return _functionalized_f_helper(primals, tangents)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 388, in _functionalized_f_helper
    f_outs = fn(*f_args)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 256, in inner_fn_with_anomaly
    return inner_fn(*args)
  File "/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 241, in inner_fn
    backward_out = torch.autograd.grad(
  File "/.venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 385, in grad
    return handle_torch_function(
  File "/.venv/lib/python3.10/site-packages/torch/overrides.py", line 1630, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 705, in __torch_function__
    return func(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 385, in grad
    return handle_torch_function(
  File "/.venv/lib/python3.10/site-packages/torch/overrides.py", line 1630, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/utils/_device.py", line 79, in __torch_function__
    return func(*args, **kwargs)
  File "/.venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 436, in grad
    result = _engine_run_backward(
  File "/.venv/lib/python3.10/site-packages/torch/autograd/graph.py", line 768, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.HalfTensor [5, 8]], which is output 0 of IndexPutBackward0, is at version 1; expected version 0 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```
I search-and-replaced my home folder name from the logs. These are otherwise completely unmodified.

### Minified repro

the minified script produced by `TORCHDYNAMO_REPRO_AFTER="dynamo" python3 rep.py` produces an invalid file, because the script generated does not know that the model should be casted to `.half()`, and runs into a trivial Float!=Half error.

```
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___a = Linear(in_features=8, out_features=8, bias=False).cuda()
        self.L__self___b = Linear(in_features=8, out_features=8, bias=False).cuda()



    def forward(self, L_x_ : torch.Tensor, L_i_ : torch.Tensor):
        l_x_ = L_x_
        l_i_ = L_i_
        x = self.L__self___a(l_x_);  l_x_ = None
        std = x.std()
        getitem = x[l_i_]
        l__self___b = self.L__self___b(getitem);  getitem = None
        x[l_i_] = l__self___b;  setitem = x;  l_i_ = l__self___b = None
        return (x, std)


mod = Repro()

def load_args(reader):
    buf0 = reader.storage('c0b101c34b70d81b6b0558c6f54708c77417b3ae', 80, device=device(type='cuda', index=0), dtype_hint=torch.float16)
    reader.tensor(buf0, (5, 8), dtype=torch.float16, is_leaf=True)  # L_x_
    buf1 = reader.storage('bfc7e924eea68b337b7599cc5e3dcbda910e662d', 16, device=device(type='cuda', index=0), dtype_hint=torch.int64)
    reader.tensor(buf1, (2,), dtype=torch.int64, is_leaf=True)  # L_i_
load_args._version = 0

if __name__ == '__main__':
    from torch._dynamo.repro.after_dynamo import run_repro
    run_repro(mod, load_args, accuracy=False, command='minify',
        save_dir='/home/a/torch_compile_debug/run_2024_08_04_11_10_18_838436-pid_2039424/minifier/checkpoints', autocast=False, backend='inductor')
```

I have nonetheless included it here to follow the issue creation instructions.

### Versions

```
Collecting environment information...
PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
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.29.3
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: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090
GPU 3: NVIDIA GeForce RTX 3090
GPU 4: NVIDIA GeForce RTX 3090
GPU 5: NVIDIA GeForce RTX 3090
GPU 6: NVIDIA GeForce RTX 3090

Nvidia driver version: 550.54.15
cuDNN version: Could not collect
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):                               32
On-line CPU(s) list:                  0-31
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7282 16-Core Processor
CPU family:                           23
Model:                                49
Thread(s) per core:                   2
Core(s) per socket:                   16
Socket(s):                            1
Stepping:                             0
Frequency boost:                      enabled
CPU max MHz:                          2800.0000
CPU min MHz:                          1500.0000
BogoMIPS:                             5599.72
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 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:                            512 KiB (16 instances)
L1i cache:                            512 KiB (16 instances)
L2 cache:                             8 MiB (16 instances)
L3 cache:                             64 MiB (4 instances)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-31
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] numpy==2.0.1
[pip3] torch==2.4.0
[pip3] torchaudio==2.4.0
[pip3] torchvision==0.19.0
[pip3] triton==3.0.0
[conda] Could not collect
```

The above venv is on Ubuntu with 3090s on a torch 2.4 stable.

However, I also tested the script with H100s on `torch-2.5.0.dev20240803+cu124`, and received the same result.

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @chauhang @penguinwu @zou3519 @bdhirsh