Lightning-AI / Lightning-AI/lightning-thunder

thunderfx splitter fails to handle `SymInt` FX nodes

Open
#2,420 0 comments 0 reactions 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.