Lightning-AI / Lightning-AI/lightning-thunder

Incompatibility with HF Model Qwen2-1.5B - Tensor Indexing Error (1-D vs 2-D)

Open
#1,135 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

huggingface mixology program-coverage triage review
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

  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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.