pytorch - 💡(How to fix) Fix DISABLED test_flex_attention_with_dynamic_max_autotune_graph_partition_cuda (__main__.TestLearnableBiasesCUDA) [1 comments, 1 participants]

Traceback (most recent call last): File "/__w/pytorch/pytorch/test/inductor/test_flex_attention.py", line 8168, in test_flex_attention_with_dynamic_max_autotune_graph_partition self._test_flex_attention_with_dynamic_max_autotune(device) File "/__w/pytorch/pytorch/test/inductor/test_flex_attention.py", line 8146, in _test_flex_attention_with_dynamic_max_autotune out = compiled_sdpa( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1062, in compile_wrapper raise e.remove_dynamo_frames() from None # see TORCHDYNAMO_VERBOSE=1 File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1078, in _compile_fx_inner raise InductorError(e, currentframe()).with_traceback( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1058, in _compile_fx_inner mb_compiled_graph = fx_codegen_and_compile( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/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 "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1521, in codegen_and_compile graph.run(example_inputs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1059, in run return super().run(args) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/fx/interpreter.py", line 197, in run self.env[node] = self.run_node(node) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1910, in run_node result = super().run_node(n) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/fx/interpreter.py", line 294, in run_node return getattr(self, n.op)(n.target, args, kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1482, in call_function raise LoweringException( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1459, in call_function out = lowerings[target](args, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/lowering.py", line 518, in wrapped out = decomp_fn(args, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/kernel/flex/flex_attention.py", line 495, in flex_attention out, _ = autotune_select_algorithm( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 5635, in autotune_select_algorithm return cache(args, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 3931, in call timings = self.do_autotuning( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 4275, in do_autotuning self.log_results( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 5470, in log_results AlgorithmSelectorCache.maybe_log_flex_attention_results( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_inductor/select_algorithm.py", line 5374, in maybe_log_flex_attention_results "B": int(B), File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/sympy/core/expr.py", line 307, in int raise TypeError("Cannot convert symbols to int") torch._inductor.exc.InductorError: LoweringException: TypeError: Cannot convert symbols to int target: flex_attention args[0]: TensorBox(StorageBox( InputBuffer(name='primals_3', layout=FixedLayout('cuda:0', torch.float32, size=[s5, 16, s37, 64], stride=[1024s37, 64s37, 64, 1])) )) args[1]: TensorBox(StorageBox( InputBuffer(name='primals_4', layout=FixedLayout('cuda:0', torch.float32, size=[s5, 16, s37, 64], stride=[1024s37, 64s37, 64, 1])) )) args[2]: TensorBox(StorageBox( InputBuffer(name='primals_5', layout=FixedLayout('cuda:0', torch.float32, size=[s5, 16, s37, 64], stride=[1024s37, 64s37, 64, 1])) )) args[3]: Subgraph(name='sdpa_score0', graph_module=<lambda>(), graph=None) args[4]: (s37, s71, TensorBox(StorageBox( InputBuffer(name='primals_11', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_8', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_13', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_14', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_15', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_16', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_17', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_18', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), s59, s30, Subgraph(name='sdpa_mask0', graph_module=<lambda>(), graph=None)) args[5]: 0.125 args[6]: {'BACKEND': 'AUTO', 'PRESCALE_QK': False, 'ROWS_GUARANTEED_SAFE': False, 'BLOCKS_ARE_CONTIGUOUS': False, 'WRITE_DQ': True, 'OUTPUT_LOGSUMEXP': True, 'OUTPUT_MAX': False} args[7]: (s39,) args[8]: ()TypeError: Cannot convert symbols to int target: flex_attention args[0]: TensorBox(StorageBox( InputBuffer(name='primals_3', layout=FixedLayout('cuda:0', torch.float32, size=[s5, 16, s37, 64], stride=[1024s37, 64s37, 64, 1])) )) args[1]: TensorBox(StorageBox( InputBuffer(name='primals_4', layout=FixedLayout('cuda:0', torch.float32, size=[s5, 16, s37, 64], stride=[1024s37, 64s37, 64, 1])) )) args[2]: TensorBox(StorageBox( InputBuffer(name='primals_5', layout=FixedLayout('cuda:0', torch.float32, size=[s5, 16, s37, 64], stride=[1024s37, 64*s37, 64, 1])) )) args[3]: Subgraph(name='sdpa_score0', graph_module=<lambda>(), graph=None) args[4]: (s37, s71, TensorBox(StorageBox( InputBuffer(name='primals_11', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_8', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_13', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_14', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_15', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_16', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_17', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22], stride=[s22, s22, 1])) )), TensorBox(StorageBox( InputBuffer(name='primals_18', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, s22, s22], stride=[s222, s222, s22, 1])) )), s59, s30, Subgraph(name='sdpa_mask0', graph_module=<lambda>(), graph=None)) args[5]: 0.125 args[6]: {'BACKEND': 'AUTO', 'PRESCALE_QK': False, 'ROWS_GUARANTEED_SAFE': False, 'BLOCKS_ARE_CONTIGUOUS': False, 'WRITE_DQ': True, 'OUTPUT_LOGSUMEXP': True, 'OUTPUT_MAX': False} args[7]: (s39,) args[8]: () Found from : File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 2040, in flex_attention out, lse, max_scores = flex_attention_hop(

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

To execute this test, run the following from the base repo dir: python test/inductor/test_flex_attention.py TestLearnableBiasesCUDA.test_flex_attention_with_dynamic_max_autotune_graph_partition_cuda

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

TL;DR

The most likely fix for the failing test is to investigate and resolve the TypeError: Cannot convert symbols to int error in the flex_attention function.

Guidance

• Investigate the flex_attention function in torch/nn/attention/flex_attention.py to understand why it's trying to convert symbols to integers.
• Check the input arguments to the flex_attention function, specifically args[4], which contains several TensorBox objects with symbolic sizes (e.g., s37, s71).
• Verify that the input tensors have the correct sizes and data types.
• Consider setting TORCHDYNAMO_VERBOSE=1 and TORCH_LOGS="+dynamo" to get more detailed error messages and internal stack traces.

Example

No specific code example is provided, as the issue requires further investigation into the flex_attention function and its input arguments.

Notes

The error message suggests that there's an issue with converting symbolic sizes to integers in the flex_attention function. This might be related to the dynamic autotuning feature in PyTorch. Further investigation is needed to determine the root cause of the issue.

Recommendation

Apply a workaround by investigating and resolving the TypeError: Cannot convert symbols to int error in the flex_attention function, rather than upgrading to a fixed version (as no specific version is mentioned in the issue).