Lightning-AI / Lightning-AI/lightning-thunder

Thunder-torchcompile errors in llama RoPE modules

Open
#2,297 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.