Lightning-AI / Lightning-AI/lightning-thunder
Incompatibility with HF Model Qwen2-1.5B - Tensor Indexing Error (1-D vs 2-D)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Model / language coverage
I encountered an issue while attempting to use `thunder.jit` with models outside of the lit-gpt universe, specifically the Hugging Face model [Qwen2-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2-1.5B-Instruct). The following error is thrown:
```
RuntimeError: Advanced indexing currently only supports zero or one-dimensional integer tensors, but found a tensor with dtype int64 and 2 dimensions.
```
The shape of the tensor in question is actually (1, 1024), which could potentially be handled with squeeze().
### Pitch
Supporting this case could enable compatibility with Qwen2-1.5B-Instruct and possibly with other models from Qwen family.
### Alternatives / Potential work-arounds
Adding the following code to the function `_advanced_indexing` in `thunder/clang/__init__.py` resolves the issue temporarily:
```py
if isinstance(x, TensorLike):
dims_to_squeeze = tuple([i for i, d in enumerate(x.shape) if d == 1])
if len(dims_to_squeeze) > 0:
x = prims.squeeze(x, dims_to_squeeze)
```
However, the same issue re-emerges in the prims.take_meta function:
```
RuntimeError: Expected index to be a 1-D or 0-D tensor, but index.ndim=2!
```
### Minimal Repro
```py
from transformers import AutoModelForCausalLM, AutoTokenizer
import thunder
import torch
# Define device and model
DEVICE = torch.device('cuda', 0)
model = AutoModelForCausalLM.from_pretrained('Qwen/Qwen2-1.5B-Instruct', torch_dtype=torch.bfloat16, device_map="cuda")
# Compilation with thunder.jit
model = thunder.jit(model)
# Tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
# Sample input with shape (1, 1024)
inputs = tokenizer("Hello world!", return_tensors="pt").input_ids.to(DEVICE)
inputs = inputs.repeat(1, 1024 // inputs.shape[1])
# Forward pass
output = model(inputs)
```
results in
```
File "/home/mmikulski/miniconda3/envs/report2/lib/python3.10/site-packages/thunder/core/prims.py", line 2882, in take_meta
utils.check(index.ndim <= 1, lambda: f"Expected index to a 1-D or 0-D tensor, but index.ndim={index.ndim}!")
File "/home/mmikulski/miniconda3/envs/report2/lib/python3.10/site-packages/thunder/core/baseutils.py", line 103, in check
raise exception_type(s())
RuntimeError: Expected index to a 1-D or 0-D tensor, but index.ndim=2!
```
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
Reproduce the failure with the Qwen2-1.5B-Instruct example, then inspect _advanced_indexing in thunder/clang/__init__.py and take_meta in thunder/core/prims.py. Trace how the (1, 1024) index is handled in both locations; done means the reported model compilation and forward pass no longer raise the tensor-dimension errors.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- huggingface, python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100