Lightning-AI / Lightning-AI/lightning-thunder
Thunder-torchcompile errors in llama RoPE modules
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
We observed errors in nvfuser CI when benchmarking rope modules from llama models using thunder-torchcompile.
Please see the standalone repro for the model definitions: [Llama RoPE Benchmarks](https://gist.github.com/Priya2698/e1e2146fde7024e6ab9109c8feec09ff).
Stack traces:
llama_3_8B:
```
> raise error_type(message_evaluated)
E torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function (*((FakeTensor(..., device='cuda:0', size=(2, 8, 4, 1024, 128),
E dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(2, 8, 1, 1024, 128),
E dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(2, 32, 1024, 128))), 2), **{}): got RuntimeError('Number of dimensions of tensors must match. Expected 5-D tensors, but got 4-D for tensor number 2 in the list')
E
E from user code:
E File "thunder.to_be_compiled_7", line 103, in to_be_compiled
E t113 = torch.cat((t112, t110, t73), 2) # t113: "cuda:0 bf16[2, 8, 6, 1024, 128]"
E
E 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"
/usr/local/lib/python3.12/dist-packages/torch/__init__.py:1650: TorchRuntimeError
---------------------------------------------------------------------------------------- Captured stderr call -----------------------------------------------------------------------------------------
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] failed while attempting to run meta for aten.cat.default
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] Traceback (most recent call last):
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2554, in _dispatch_impl
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] r = func(*args, **kwargs)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] ^^^^^^^^^^^^^^^^^^^^^
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 756, in __call__
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] return self._op(*args, **kwargs)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] ^^^^^^^^^^^^^^^^^^^^^^^^^
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_prims_common/wrappers.py", line 308, in _fn
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] result = fn(*args, **kwargs)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] ^^^^^^^^^^^^^^^^^^^
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 51, in inner
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] return disable_fn(*args, **kwargs)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 850, in _fn
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] return fn(*args, **kwargs)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] ^^^^^^^^^^^^^^^^^^^
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_prims_common/wrappers.py", line 149, in _fn
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] result = fn(**bound.arguments)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] ^^^^^^^^^^^^^^^^^^^^^
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/_refs/__init__.py", line 2795, in cat
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] torch._check(
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/__init__.py", line 1668, in _check
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] _check_with(RuntimeError, cond, message)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] File "/usr/local/lib/python3.12/dist-packages/torch/__init__.py", line 1650, in _check_with
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] raise error_type(message_evaluated)
E0611 16:28:19.380000 2241861 torch/_subclasses/fake_tensor.py:2558] [1/0] RuntimeError: Number of dimensions of tensors must match. Expected 5-D tensors, but got 4-D for tensor number 2 in the list
```
llama_2_7b:
```
> raise e.with_traceback(None) from e.__cause__
E torch._dynamo.exc.Unsupported: name 't71' is not defined
E
E from user code:
E File "thunder.to_be_compiled_7", line 100, in to_be_compiled
E t109 = torch.cat((t107, t106, t71), 2) # t109: "cuda:0 bf16[2, 32, 3, 1024, 128]"
E
E 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"
/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py:667: Unsupported
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the standalone Llama RoPE Benchmarks repro linked in the issue and run the llama_3_8B and llama_2_7b cases through thunder-torchcompile in the nvfuser CI setup. Investigate the reported torch.cat dimension mismatch and undefined-name errors; done means both benchmarks compile without these errors.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100