pytorch / pytorch/pytorch

[HOP] Backward pass for scan HOP can result in invalid tensor metadata

Open
#182,381 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: autograd module: fakeTensor module: higher order operators module: pt2-dispatcher oncall: pt2 triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

This can be reproduced by running any one of the tests from `test/inductor/test_control_flow.py` in `ScanTests.test_scan_in_cond` with `autograd==True`. This reproduces only once https://github.com/pytorch/pytorch/pull/159523 merges, since it exposed the underlying issue. A [tlparse log](https://github.com/user-attachments/files/27378119/dedicated_log_torch_trace_xmqc1m8u.log) is attached for `test_scan_in_cond_device_cpu_dynamic_False_reverse_False_dim_0_pred_False_scan_length_1_autograd_True` as a representative example.

The underlying issue appears to be triggered by the following sequence:

1. The forward pass for the model is a scan followed by a `sum`.
2. When doing autograd on the model, the `sum` gradient is a single value with `torch.expand` called on it, so all strides are 0.
3. This gets passed into a new scan on the backward/autograd pass as an init tensor (in `ScanAutogradImpl.call_backward`). The corresponding `xs` tensor does not have all-0 strides.
4. The backwards pass scan gets converted to a `while_loop` by Inductor in a post-grad pass, and `FakeTensorUpdater` ultimately calls the `while_loop` FakeTensorMode impl, which does checks on tensor metadata consistency that fail (due to the mismatch in strides between the `init` tensor and the `xs` tensor slices).

This error can be triggered before the conversion to `while_loop` by patching `torch._higher_order_ops.scan.trace_scan` to not use cloned init tensors at [this line](https://github.com/pytorch/pytorch/blob/af17dfed76e67624a5e0fb856ef784a24ba7a81b/torch/_higher_order_ops/scan.py#L411) when running the same metadata consistency checks, since `clone()` copies overlapping strided tensors into non-overlapping strided tensors.

I _suspect_ the fix may look something like a two-pass step in `ScanAutogradImpl.call_backward`, where we determine what the output strides for the scan returns are, then optionally do a clone on the init tensors to match and regenerate the subgraph.

### Error logs

```
Traceback (most recent call last):
File "/home/bglass/miniforge3/envs/pytorch-dev/lib/python3.10/unittest/case.py", line 59, in testPartExecutor
yield
File "/home/bglass/miniforge3/envs/pytorch-dev/lib/python3.10/unittest/case.py", line 591, in run
self._callTestMethod(testMethod)
File "/home/bglass/miniforge3/envs/pytorch-dev/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
method()
File "/home/bglass/git/pytorch/torch/testing/_internal/common_utils.py", line 3528, in wrapper
method(*args, **kwargs)
File "/home/bglass/git/pytorch/torch/testing/_internal/common_utils.py", line 603, in instantiated_test
test(self, **param_kwargs)
File "/home/bglass/miniforge3/envs/pytorch-dev/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/home/bglass/git/pytorch/test/inductor/test_control_flow.py", line 2087, in test_scan_in_cond
self._run_test(
File "/home/bglass/git/pytorch/test/inductor/test_control_flow.py", line 1977, in _run_test
result_compiled = _run_model(model3, [scan] + inputs)
File "/home/bglass/git/pytorch/test/inductor/test_control_flow.py", line 1961, in _run_model
loss.backward()
File "/home/bglass/git/pytorch/torch/_tensor.py", line 631, in backward
torch.autograd.backward(
File "/home/bglass/git/pytorch/torch/autograd/__init__.py", line 381, in backward
_engine_run_backward(
File "/home/bglass/git/pytorch/torch/autograd/graph.py", line 913, in _engine_run_backward
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
File "/home/bglass/git/pytorch/torch/autograd/function.py", line 333, in apply_boxed
return self._get_user_fn()(self, *args)
File "/home/bglass/git/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 3423, in backward
return impl_fn()
File "/home/bglass/git/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 3403, in impl_fn
out = CompiledFunction._backward_impl(ctx, all_args)
File "/home/bglass/git/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 3463, in _backward_impl
compiled_bw = backward_compiler.get_or_compile(
File "/home/bglass/git/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 3100, in get_or_compile
self.compiled_bw = self.aot_config.bw_compiler(
File "/home/bglass/git/pytorch/torch/_functorch/_aot_autograd/schemas.py", line 1421, in __call__
output_code = self.compiler_fn(gm, example_inputs)
File "/home/bglass/git/pytorch/torch/_dynamo/backends/common.py", line 83, in _wrapped_bw_compiler
disable(
File "/home/bglass/git/pytorch/torch/_dynamo/eval_frame.py", line 1353, in _fn
return fn(*args, **kwargs)
File "/home/bglass/git/pytorch/torch/_utils_internal.py", line 96, in wrapper_function
return function(*args, **kwargs)
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 2923, in bw_compiler
return compile_fx_backward(
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 2585, in compile_fx_backward
return inner_compile(
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 836, in compile_fx_inner
return wrap_compiler_debug(_compile_fx_inner, compiler_name="inductor")(
File "/home/bglass/git/pytorch/torch/_dynamo/repro/after_aot.py", line 314, in debug_wrapper
inner_compiled_fn = compiler_fn(gm, example_inputs)
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 1047, in _compile_fx_inner
raise InductorError(e, currentframe()).with_traceback(
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 1039, in _compile_fx_inner
mb_compiled_graph = fx_codegen_and_compile(
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 1845, in fx_codegen_and_compile
return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs)
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 1378, in codegen_and_compile
_recursive_post_grad_passes(gm, is_inference=is_inference)
File "/home/bglass/git/pytorch/torch/_inductor/compile_fx.py", line 581, in _recursive_post_grad_passes
post_grad_passes(gm, is_inference)
File "/home/bglass/git/pytorch/torch/_inductor/fx_passes/post_grad.py", line 234, in post_grad_passes
fake_tensor_updater.incremental_update()
File "/home/bglass/git/pytorch/torch/_inductor/fx_utils.py", line 485, in incremental_update
].incremental_update()
File "/home/bglass/git/pytorch/torch/_inductor/fx_utils.py", line 509, in incremental_update
new_fake_tensor = node.target(*args, **kwargs)
File "/home/bglass/git/pytorch/torch/_higher_order_ops/while_loop.py", line 64, in __call__
return super().__call__(
File "/home/bglass/git/pytorch/torch/_ops.py", line 539, in __call__
return self.dispatch(dispatch_key_set.highestPriorityTypeId(), *args, **kwargs)
File "/home/bglass/git/pytorch/torch/_ops.py", line 386, in dispatch
return kernel(*args, **kwargs)
File "/home/bglass/git/pytorch/torch/_ops.py", line 336, in maybe_run_autograd
return self(*args, **kwargs)
File "/home/bglass/git/pytorch/torch/_higher_order_ops/while_loop.py", line 64, in __call__
return super().__call__(
File "/home/bglass/git/pytorch/torch/_ops.py", line 539, in __call__
return self.dispatch(dispatch_key_set.highestPriorityTypeId(), *args, **kwargs)
File "/home/bglass/git/pytorch/torch/_ops.py", line 422, in dispatch
result = handler(mode, *args, **kwargs)
File "/home/bglass/git/pytorch/torch/_higher_order_ops/while_loop.py", line 566, in while_loop_fake_tensor_mode
check_meta_consistency(
File "/home/bglass/git/pytorch/torch/_higher_order_ops/utils.py", line 256, in check_meta_consistency
raise torch._dynamo.exc.UncapturedHigherOrderOpError(
torch._inductor.exc.InductorError: UncapturedHigherOrderOpError: Expected carried_inputs and body_output to have same metadata but found:
pair[2] differ in 'stride: (0, 0, 0) vs (16, 4, 1)', where lhs is FakeTensor(..., size=(4, 4, 4), dtype=torch.float64) and rhs is FakeTensor(..., size=(4, 4, 4), dtype=torch.float64)

To execute this test, run the following from the base repo dir:
python test/inductor/test_control_flow.py ScanTests.test_scan_in_cond_device_cpu_dynamic_False_reverse_False_dim_0_pred_False_scan_length_1_autograd_True
```

### Versions

```
Collecting environment information...
PyTorch version: 2.13.0a0+gita6972ad
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (conda-forge gcc 15.2.0-18) 15.2.0
Clang version: 22.1.0 (https://github.com/conda-forge/clangdev-feedstock cb6d0bfa739b2869ba0512a73fc5c0f2ede71d42)
CMake version: version 4.2.3
Libc version: glibc-2.39

Python version: 3.10.20 | packaged by conda-forge | (main, Mar 5 2026, 16:42:22) [GCC 14.3.0] (64-bit runtime)
Python platform: Linux-6.8.0-110-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.0.88
CUDA_MODULE_LOADING set to:
GPU models and configuration:
GPU 0: Quadro RTX 8000
GPU 1: Quadro RTX 8000

Nvidia driver version: 590.44.01
cuDNN version: Probably one of the following:
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_adv.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_cnn.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_engines_precompiled.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_engines_runtime_compiled.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_graph.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_heuristic.so.9
/space/usr/local/cuda-12.6.3/targets/x86_64-linux/lib/libcudnn_ops.so.9
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False
Caching allocator config: N/A

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): 48
On-line CPU(s) list: 0-47
Vendor ID: AuthenticAMD
Model name: AMD Ryzen Threadripper 2970WX 24-Core Processor
CPU family: 23
Model: 8
Thread(s) per core: 2
Core(s) per socket: 24
Socket(s): 1
Stepping: 2
Frequency boost: enabled
CPU(s) scaling MHz: 76%
CPU max MHz: 3000.0000
CPU min MHz: 2200.0000
BogoMIPS: 5988.02
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 amd_dcm 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 skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb hw_pstate ssbd ibpb vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif overflow_recov succor smca sev sev_es ibpb_exit_to_user
Virtualization: AMD-V
L1d cache: 768 KiB (24 instances)
L1i cache: 1.5 MiB (24 instances)
L2 cache: 12 MiB (24 instances)
L3 cache: 64 MiB (8 instances)
NUMA node(s): 4
NUMA node0 CPU(s): 0-5,24-29
NUMA node1 CPU(s): 12-17,36-41
NUMA node2 CPU(s): 6-11,30-35
NUMA node3 CPU(s): 18-23,42-47
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: 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 vulnerable
Vulnerability Spec rstack overflow: Mitigation; Safe RET
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Retpolines; IBPB conditional; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Mitigation; IBPB before exit to userspace

Versions of relevant libraries:
[pip3] bert_pytorch==0.0.1a4
[pip3] flake8==7.3.0
[pip3] functorch==1.14.0a0+b71aa0b
[pip3] mypy==1.20.0
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.2.6
[pip3] onnx==1.20.0
[pip3] onnx-ir==0.1.12
[pip3] optree==0.19.0
[pip3] pytorch-labs-segment-anything-fast==0.2
[pip3] torch==2.13.0a0+gita6972ad
[pip3] torch_geometric==2.4.0
[pip3] torchao==0.17.0
[pip3] torchaudio==2.11.0a0+c0cbdb9
[pip3] torchdata==0.12.0a0+e640e6f
[pip3] torchmultimodal==0.1.0b0
[pip3] torchvision==0.27.0a0+499ca51
[pip3] triton==3.7.0+git88b227e2
[conda] bert-pytorch 0.0.1a4 pypi_0 pypi
[conda] cuda-cudart 13.0.96 hecca717_0 conda-forge
[conda] cuda-cudart-dev 13.0.96 hecca717_0 conda-forge
[conda] cuda-cudart-dev_linux-64 13.0.96 h376f20c_0 conda-forge
[conda] cuda-cudart-static 13.0.96 hecca717_0 conda-forge
[conda] cuda-cudart-static_linux-64 13.0.96 h376f20c_0 conda-forge
[conda] cuda-cudart_linux-64 13.0.96 h376f20c_0 conda-forge
[conda] cuda-cupti 13.0.85 h676940d_0 conda-forge
[conda] cuda-cupti-dev 13.0.85 h676940d_0 conda-forge
[conda] cuda-libraries-dev 13.0.2 ha770c72_0 conda-forge
[conda] cuda-nvrtc 13.0.88 hecca717_0 conda-forge
[conda] cuda-nvrtc-dev 13.0.88 hecca717_0 conda-forge
[conda] cuda-nvtx 13.0.85 hecca717_0 conda-forge
[conda] cuda-nvtx-dev 13.0.85 ha770c72_0 conda-forge
[conda] cuda-opencl 13.0.85 hecca717_0 conda-forge
[conda] cuda-opencl-dev 13.0.85 hecca717_0 conda-forge
[conda] cudnn 9.20.0.48 h886f0b6_0 conda-forge
[conda] functorch 1.14.0a0+b71aa0b pypi_0 pypi
[conda] libcublas 13.1.0.3 h676940d_0 conda-forge
[conda] libcublas-dev 13.1.0.3 h676940d_0 conda-forge
[conda] libcudnn 9.20.0.48 ha4b6413_0 conda-forge
[conda] libcudnn-dev 9.20.0.48 h7bcfba5_0 conda-forge
[conda] libcufft 12.0.0.61 hecca717_0 conda-forge
[conda] libcufft-dev 12.0.0.61 hecca717_0 conda-forge
[conda] libcurand 10.4.0.35 h676940d_1 conda-forge
[conda] libcurand-dev 10.4.0.35 h676940d_1 conda-forge
[conda] libcusolver 12.0.4.66 h676940d_1 conda-forge
[conda] libcusolver-dev 12.0.4.66 h676940d_1 conda-forge
[conda] libcusparse 12.6.3.3 hecca717_0 conda-forge
[conda] libcusparse-dev 12.6.3.3 hecca717_0 conda-forge
[conda] libmagma 2.9.0 hd93470c_6 conda-forge
[conda] libmagma-devel 2.9.0 h5dfe813_6 conda-forge
[conda] libmagma_sparse 2.9.0 h69b1545_6 conda-forge
[conda] libnvjitlink 13.0.88 hecca717_0 conda-forge
[conda] libnvjitlink-dev 13.0.88 hecca717_0 conda-forge
[conda] mkl 2025.3.1 h0e700b2_10 conda-forge
[conda] mkl-include 2025.3.1 hf2ce2f3_10 conda-forge
[conda] nccl 2.29.3.1 h8340e53_0 conda-forge
[conda] numpy 2.2.6 py310hefbff90_0 conda-forge
[conda] optree 0.19.0 py310h03d9f68_0 conda-forge
[conda] pytorch-labs-segment-anything-fast 0.2 pypi_0 pypi
[conda] tbb 2022.3.0 hb700be7_2 conda-forge
[conda] torch 2.13.0a0+gita6972ad pypi_0 pypi
[conda] torch-geometric 2.4.0 pypi_0 pypi
[conda] torchao 0.17.0 pypi_0 pypi
[conda] torchaudio 2.11.0a0+c0cbdb9 pypi_0 pypi
[conda] torchdata 0.12.0a0+e640e6f pypi_0 pypi
[conda] torchmultimodal 0.1.0b0 pypi_0 pypi
[conda] torchvision 0.27.0a0+499ca51 pypi_0 pypi
[conda] triton 3.7.0+git88b227e2 pypi_0 pypi
```

cc @ezyang @albanD @gqchen @nikitaved @soulitzer @Varal7 @bobrenjc93 @chauhang @penguinwu @eellison @aorenste @ydwu4 @bdhirsh

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.