[Graph Partition] [Inductor] UnboundLocalError: cannot access local variable 'buf271' where it is not associated with a value

[wmhst7]

🐛 Describe the bug

Using "reduce-overhead" mode and "inductor backend for training, with torch._inductor.config.graph_partition = True. Run into inductor gen-code bug:

[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1130, in forward
[rank0]:     return compiled_fn(full_args)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 339, in runtime_wrapper
[rank0]:     all_outs = call_func_at_runtime_with_args(
[rank0]:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 129, in call_func_at_runtime_with_args
[rank0]:     out = normalize_as_list(f(args))
[rank0]:                             ^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 103, in g
[rank0]:     return f(*args)
[rank0]:            ^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/autograd/function.py", line 581, in apply
[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 2118, in forward
[rank0]:     fw_outs = call_func_at_runtime_with_args(
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 129, in call_func_at_runtime_with_args
[rank0]:     out = normalize_as_list(f(args))
[rank0]:                             ^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 526, in wrapper
[rank0]:     return compiled_fn(runtime_args)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 690, in inner_fn
[rank0]:     unwrapped_outs = compiled_fn(unwrapped_args)
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 724, in inner_fn
[rank0]:     outs = compiled_fn(args)
[rank0]:            ^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_inductor/output_code.py", line 613, in __call__
[rank0]:     return self.current_callable(inputs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/tiger/.local/lib/python3.11/site-packages/torch/_inductor/utils.py", line 3017, in run
[rank0]:     out = model(new_inputs)
[rank0]:           ^^^^^^^^^^^^^^^^^
[rank0]:   File "/tmp/torchinductor_tiger/tmpngii2htx/na/cnabkmabktacecyr75a7sgnkip7pjfcd672lse2ndmzilbphpxxh.py", line 5071, in call
[rank0]:     partition1_args = [buf301, buf305, buf306, primals_42, buf311, buf286, primals_45, buf294, buf271, s54, u0, u1]
[rank0]:                                                                                                ^^^^^^
[rank0]: UnboundLocalError: cannot access local variable 'buf271' where it is not associated with a value

Versions

Collecting environment information...
PyTorch version: 2.9.1+cu129
Is debug build: False
CUDA used to build PyTorch: 12.9
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 12 (bookworm) (x86_64)
GCC version: (Debian 12.2.0-14+deb12u1) 12.2.0
Clang version: Could not collect
CMake version: version 3.31.6
Libc version: glibc-2.36

Python version: 3.11.2 (main, Apr 28 2025, 14:11:48) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-5.15.152.bsk.10-amd64-x86_64-with-glibc2.36
Is CUDA available: True
CUDA runtime version: 12.9.86
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: NVIDIA H800
Nvidia driver version: 535.261.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.11.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.11.0
Is XPU available: False
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: 52 bits physical, 57 bits virtual
Byte Order: Little Endian
CPU(s): 180
On-line CPU(s) list: 0-179
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) Platinum 8457C
CPU family: 6
Model: 143
Thread(s) per core: 2
Core(s) per socket: 45
Socket(s): 2
Stepping: 8
BogoMIPS: 5200.00
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd arat avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid cldemote movdiri movdir64b fsrm md_clear serialize tsxldtrk arch_lbr amx_bf16 avx512_fp16 amx_tile amx_int8 arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 4.2 MiB (90 instances)
L1i cache: 2.8 MiB (90 instances)
L2 cache: 180 MiB (90 instances)
L3 cache: 195 MiB (2 instances)
NUMA node(s): 2
NUMA node0 CPU(s): 0-89
NUMA node1 CPU(s): 90-179
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: Unknown: No mitigations
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: 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; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.9.1.4
[pip3] nvidia-cuda-cupti-cu12==12.9.79
[pip3] nvidia-cuda-nvrtc-cu12==12.9.86
[pip3] nvidia-cuda-runtime-cu12==12.9.79
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cufft-cu12==11.4.1.4
[pip3] nvidia-curand-cu12==10.3.10.19
[pip3] nvidia-cusolver-cu12==11.7.5.82
[pip3] nvidia-cusparse-cu12==12.5.10.65
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-nccl-cu12==2.27.5
[pip3] nvidia-nvjitlink-cu12==12.9.86
[pip3] nvidia-nvtx-cu12==12.9.79
[pip3] optree==0.17.0
[pip3] torch==2.9.1+cu129
[pip3] torchaudio==2.9.1+cu129
[pip3] torchvision==0.24.1+cu129
[pip3] triton==3.5.1
[conda] Could not collect

cc @mcarilli @ezyang @eellison @penguinwu @BoyuanFeng @chauhang @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo

[wmhst7]

In output_code.py, buf271 occurs 3 times. It seems buf271 is introduced with no initialization, which is the cause:

def partition_1(args):
    buf301, buf305, buf306, primals_42, buf311, buf286, primals_45, buf294, buf271, s54, u0, u1 = args
    args.clear()
    ...
...
class Runner:
    def __init__(self, partitions):
        self.partitions = partitions

    def recursively_apply_fns(self, fns):
        new_callables = []
        for fn, c in zip(fns, self.partitions):
            new_callables.append(fn(c))
        self.partitions = new_callables

    def call(self, args):
        primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9, primals_10, primals_11, primals_12, primals_13, primals_14, primals_15, primals_16, primals_17, primals_18, primals_19, primals_20, primals_21, primals_22, primals_23, primals_24, primals_25, primals_26, primals_27, primals_28, primals_29, primals_30, primals_31, primals_32, primals_33, primals_34, primals_35, primals_36, primals_37, primals_38, primals_39, primals_40, primals_41, primals_42, primals_43, primals_44, primals_45 = args
        args.clear()
        s9 = primals_4
        s54 = primals_44
        u1 = reinterpret_tensor(primals_3, (), (), 8).item()
        buf324 = None
        partition0_args = [primals_3, primals_5, primals_42, primals_45, primals_21, primals_2, primals_7, primals_6, primals_8, primals_9, primals_10, primals_11, primals_12, primals_13, primals_14, primals_15, primals_16, primals_18, primals_17, primals_22, primals_23, primals_24, primals_25, primals_26, primals_28, primals_29, primals_30, primals_31, primals_32, primals_33, primals_34, primals_36, primals_35, primals_39, primals_40, primals_41, primals_37, primals_19, primals_20, primals_27, primals_38, primals_43, s54, s9]
        del primals_9, primals_20, primals_27, primals_38
        (buf306, buf299, buf300, buf294, buf305, buf297, buf298, buf339, buf286, buf311, buf317, buf318, buf292, buf293, primals_3, primals_5, primals_42, primals_45, primals_21, primals_2, primals_7, primals_6, primals_8, primals_10, primals_11, primals_12, primals_13, primals_14, primals_15, primals_16, primals_18, primals_17, primals_22, primals_23, primals_24, primals_25, primals_26, primals_28, primals_29, primals_30, primals_31, primals_32, primals_33, primals_34, primals_36, primals_35, primals_39, primals_40, primals_41, primals_37, primals_19, primals_43) = self.partitions[0](partition0_args)
        del partition0_args
        with torch.cuda._DeviceGuard(3):
            torch.cuda.set_device(3)
            # Topologically Sorted Source Nodes: [getitem_116, idx_lst, lt_2, valid_idx_lst], Original ATen: [aten.slice, aten.sub, aten.lt, aten.index]
            buf301 = torch.ops.aten.index.Tensor(buf299, [buf300])
            u0 = buf301.size(0)
            del buf299
            del buf300
        partition1_args = [buf301, buf305, buf306, primals_42, buf311, buf286, primals_45, buf294, buf271, s54, u0, u1]
        del buf306, buf311, buf286, buf294, buf271

[wmhst7]

Update: found the root cause. It is a cuda graph partition bug when handling custom ops with side effects.

[williamwen42]

cc @BoyuanFeng for cudagraphs issue