Lightning-AI / Lightning-AI/lightning-thunder

more general inplace support (index_copy_ in litgpt fails to trace)

Open
#1,743 1 comment 0 reactions 1 assignee View on GitHub

@ali-alshaar7 is already working on this.

Since Feb 4, 2025.

enhancement in-place
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.

## 🐛 Bug

### To Reproduce

Steps to reproduce the behavior:

1. Go to '...'
1. Run '....'
1. Scroll down to '....'
1. See error

#### Code sample

```
import thunder, torch, litgpt

with torch.device('cuda'):
m = litgpt.GPT.from_name('Llama-3.2-1B').bfloat16().requires_grad_(False)
m.max_seq_length = 1024
m.set_kv_cache(16)
inp1 = torch.ones(2, 16, device="cuda", dtype=torch.int32)
inp_pos1 = torch.stack([torch.arange(16, device="cuda") for _ in range(2)])
jm = thunder.jit(m)
jm(inp1, inp_pos1)
```

fail with

```
File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 547, in get_computation_and_inputs
orig_to_view_swap_map = check_inplace_to_views(computation_trc)
File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/functionalization.py", line 100, in check_inplace_to_views
check(
File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/baseutils.py", line 146, in check
raise exception_type(s())
NotImplementedError: in-place op of `torch.Tensor.index_copy_` to ``, a view tensor of `` is not supported because 524288 != 1048576
```

### Expected behavior

it compiles

### Environment

- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:

### Additional context

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.