CUDA Graphs with FSDP 2 and flex attention?

zaptrem: <p>Hello,</p>
<p>Is torch.compile with cuda graphs (either via reduce-overhead or max-autotune) compatible with FSDP 2 and flex attention? When both cuda graphs and FSDP 2 are on I get seemingly unending messages like:</p>
<p>``<br>
torch/_inductor/cudagraph_trees.py:2213] [__cudagraphs] Recording function 6 of graph recording id 355<br>
``</p>
<p>with recording id increasing infinitely and function N increasing to ~9. Steps take much longer than usual and never seem to speed up (waited ~15 minutes). When I disable cuda graphs training works as expected, and when I disable FSDP (and use no parallelism) CUDA graphs work as expected.</p>
<p>I also tried it with SimpleFSDP from torchtitan and got this error:</p>
<p>``<br>
[rank0]:   File “/home/ubuntu/sky_workdir/model_stuff/models.py”, line 1586, in forward<br>
[rank0]:     def forward(self, x, flex_attention_masks=None, *args, **kwargs):<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py”, line 929, in _fn<br>
[rank0]:     return fn(*args, **kwargs)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py”, line 1241, in forward<br>
[rank0]:     return compiled_fn(full_args)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py”, line 370, in runtime_wrapper<br>
[rank0]:     all_outs = call_func_at_runtime_with_args(<br>
[rank0]:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py”, line 126, in call_func_at_runtime_with_args<br>
[rank0]:     out = normalize_as_list(f(args))<br>
[rank0]:                             ^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py”, line 100, in g<br>
[rank0]:     return f(*args)<br>
[rank0]:            ^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/autograd/function.py”, line 576, in apply<br>
[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py”, line 2074, in forward<br>
[rank0]:     fw_outs = call_func_at_runtime_with_args(<br>
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py”, line 126, in call_func_at_runtime_with_args<br>
[rank0]:     out = normalize_as_list(f(args))<br>
[rank0]:                             ^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py”, line 556, in wrapper<br>
[rank0]:     return compiled_fn(runtime_args)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py”, line 716, in inner_fn<br>
[rank0]:     unwrapped_outs = compiled_fn(unwrapped_args)<br>
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py”, line 750, in inner_fn<br>
[rank0]:     outs = compiled_fn(args)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/output_code.py”, line 584, in <strong>call</strong><br>
[rank0]:     return self.current_callable(inputs)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/compile_fx.py”, line 1655, in run<br>
[rank0]:     return compiled_fn(new_inputs)  # type: ignore[arg-type]<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/cudagraph_trees.py”, line 387, in deferred_cudagraphify<br>
[rank0]:     return fn(inputs)<br>
[rank0]:            ^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/utils.py”, line 2716, in run<br>
[rank0]:     out = model(new_inputs)<br>
[rank0]:           ^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/cudagraph_trees.py”, line 2012, in run<br>
[rank0]:     out = self._run(new_inputs, function_id)<br>
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/cudagraph_trees.py”, line 2182, in _run<br>
[rank0]:     return self.record_function(new_inputs, function_id)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/cudagraph_trees.py”, line 2219, in record_function<br>
[rank0]:     node = CUDAGraphNode(<br>
[rank0]:            ^^^^^^^^^^^^^^<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/_inductor/cudagraph_trees.py”, line 1054, in <strong>init</strong><br>
[rank0]:     self.graph.replay()<br>
[rank0]:   File “/home/ubuntu/.cache/pypoetry/virtualenvs/the-model-GtEv2spP-py3.11/lib/python3.11/site-packages/torch/cuda/graphs.py”, line 117, in replay<br>
[rank0]:     super().replay()<br>
[rank0]: torch.AcceleratorError: CUDA error: an illegal memory access was encountered<br>
``</p>
<p>Both tried with 2.8 and latest 2.9 RC build as of 10/02/2025 at 7pm PT.</p>