Lightning-AI / Lightning-AI/lightning-thunder
thunderfx splitter fails to handle `SymInt` FX nodes
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
Instantiated from the following two:
- https://github.com/Lightning-AI/lightning-thunder/pull/2416#issuecomment-3177929542
- https://github.com/Lightning-AI/lightning-thunder/pull/2419
### To Reproduce
#### Code sample
```python
"""
This script is to reproduce the error below.
Env:
- container: pjnl-20250812
- thunder: https://github.com/Lightning-AI/lightning-thunder/pull/2416/commits/91289af7fcdc6ee9fafc497ec28bbd620900c819 but rebased on https://github.com/Lightning-AI/lightning-thunder/commit/84c76b44ef23dc3483ab1caea51dfbf80e93bc2c
[rank0]: Traceback (most recent call last):
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 959, in
[rank0]: main()
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 940, in main
[rank0]: metrics = run_semianalysis_benchmark(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 825, in run_semianalysis_benchmark
[rank0]: metrics = benchmark.run_benchmark()
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 639, in run_benchmark
[rank0]: iter_metrics = self.measure_inference_step(input_ids, past_key_values, self.config.output_length)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 583, in measure_inference_step
[rank0]: generation_result = self.generate(input_ids, max_new_tokens, past_key_values)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 565, in generate
[rank0]: next_token = self.decode_one_token(next_token, past_key_values)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 545, in decode_one_token
[rank0]: return self.get_next_token(input_ids, past_key_values)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_inference.py", line 523, in get_next_token
[rank0]: outputs = self.model(input_ids, past_key_values=past_key_values, use_cache=True)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/dynamo/compiler.py", line 241, in __call__
[rank0]: return self._func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 372, in __call__
[rank0]: return super().__call__(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 712, in compile_wrapper
[rank0]: raise e.remove_dynamo_frames() from None # see TORCHDYNAMO_VERBOSE=1
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1855, in _call_user_compiler
[rank0]: raise BackendCompilerFailed(
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1830, in _call_user_compiler
[rank0]: compiled_fn = compiler_fn(gm, example_inputs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/repro/after_dynamo.py", line 150, in __call__
[rank0]: compiled_gm = compiler_fn(gm, example_inputs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/__init__.py", line 2410, in __call__
[rank0]: return self.compiler_fn(model_, inputs_, **self.kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/dynamo/compiler.py", line 118, in __call__
[rank0]: split_module, subgraph_info = _splitter(gm, self._thunder_jit, self._torch_compile, sample_args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/pytorch/lightning-thunder/thunder/dynamo/splitter.py", line 175, in _splitter
[rank0]: if n.meta["example_value"].grad_fn is None:
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: torch._dynamo.exc.BackendCompilerFailed: backend='' raised:
[rank0]: AttributeError: 'SymInt' object has no attribute 'grad_fn'
The fix itself is quite simple as we can see in https://github.com/Lightning-AI/lightning-thunder/pull/2419,
but I want to reproduce with this standalone script.
"""
from __future__ import annotations
from sympy import Symbol
import torch
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from thunder.dynamo.compiler import thunderfx
class GraphModule(torch.nn.Module):
def forward(self, L_self_cumulative_length_1_: "Sym(s27)"):
l_self_cumulative_length_1_ = L_self_cumulative_length_1_
# File: /usr/local/lib/python3.12/dist-packages/transformers/cache_utils.py:1946 in _sliding_update, code: self.cumulative_length[layer_idx] += key_states.shape[-2]
add: "Sym(s27 + 1)" = l_self_cumulative_length_1_ + 1
l_self_cumulative_length_1_ = None
return (add,)
def main():
module = GraphModule().cuda()
jitted = thunderfx(module)
shape_env = ShapeEnv()
sym_expr = Symbol("n")
sym_int_with_hint = shape_env.create_symintnode(sym_expr, hint=27)
sample_args = (sym_int_with_hint,)
jitted(*sample_args)
if __name__ == "__main__":
main()
```
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with thunder/dynamo/splitter.py at the example_value handling shown in the traceback, then run the standalone reproduction from the issue with the SymInt input. Done means the thunderfx call no longer raises AttributeError for a SymInt FX node and the reproduction completes successfully.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100