Hello,
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:
``
torch/_inductor/cudagraph_trees.py:2213] [__cudagraphs] Recording function 6 of graph recording id 355
``
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.
I also tried it with SimpleFSDP from torchtitan and got this error:
``
[rank0]: File “/home/ubuntu/sky_workdir/model_stuff/models.py”, line 1586, in forward
[rank0]: def forward(self, x, flex_attention_masks=None, *args, **kwargs):
[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
[rank0]: return fn(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^
[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
[rank0]: return compiled_fn(full_args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: all_outs = call_func_at_runtime_with_args(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: out = normalize_as_list(f(args))
[rank0]: ^^^^^^^
[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
[rank0]: return f(*args)
[rank0]: ^^^^^^^^
[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
[rank0]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: fw_outs = call_func_at_runtime_with_args(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: out = normalize_as_list(f(args))
[rank0]: ^^^^^^^
[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
[rank0]: return compiled_fn(runtime_args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: unwrapped_outs = compiled_fn(unwrapped_args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: outs = compiled_fn(args)
[rank0]: ^^^^^^^^^^^^^^^^^
[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 call
[rank0]: return self.current_callable(inputs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: return compiled_fn(new_inputs) # type: ignore[arg-type]
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: return fn(inputs)
[rank0]: ^^^^^^^^^^
[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
[rank0]: out = model(new_inputs)
[rank0]: ^^^^^^^^^^^^^^^^^
[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
[rank0]: out = self._run(new_inputs, function_id)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: return self.record_function(new_inputs, function_id)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[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
[rank0]: node = CUDAGraphNode(
[rank0]: ^^^^^^^^^^^^^^
[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 init
[rank0]: self.graph.replay()
[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
[rank0]: super().replay()
[rank0]: torch.AcceleratorError: CUDA error: an illegal memory access was encountered
``
Both tried with 2.8 and latest 2.9 RC build as of 10/02/2025 at 7pm PT.