Lightning-AI / Lightning-AI/lightning-thunder
jit: `torch.cuda.stream` and other related functionality are silently ignored when jitting.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
As found by @syed-ahmed, `thunder.jit` silently ignores the use of `torch.cuda.stream` in the jitted function. There should a warning or an error. As suggested by @IvanYashchuk, we could also make this a SharpEdge.
```python
import torch
import thunder
def func(a):
s = torch.cuda.Stream()
with torch.cuda.stream(s):
for i in range(3):
a = torch.matmul(a, a)
a = torch.nn.functional.relu(a)
return a
a = torch.randn(256, 256, device=torch.device("cuda"))
# Thunder: ignores streams
jfunc = thunder.jit(func, sharp_edges=thunder.SHARP_EDGES_OPTIONS.ERROR)
jfunc(a)
# Eager: uses streams
func(a)
```
cc @apaz-cli
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
Start with the reproduction in the issue and compare eager execution with thunder.jit when using torch.cuda.Stream and torch.cuda.stream. Trace how sharp_edges=thunder.SHARP_EDGES_OPTIONS.ERROR handles unsupported functionality. Done means stream usage is no longer silently ignored and the chosen warning, error, or SharpEdge behavior is covered by a regression test.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100