Lightning-AI / Lightning-AI/lightning-thunder
Consistency in the order of applying a transform and using `thunder.jit`
Open
Nobody has claimed this yet.
bug
transforms
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
grad transform has to be applied after thunder.jit while autocast has to be applied before thunder.jit
import torch
import torch.nn as nn
import thunder
from thunder.core.transforms import grad, autocast
from thunder.examine import examine
def foo(x):
return x
x = torch.randn(3, device='cpu')
jfoo = thunder.jit(foo)
# RuntimeError: Can only transform compiled thunder functions
# o = thunder.jit(grad(foo))(x)
# Works
o = grad(thunder.jit(foo))(x)
# Works
o = thunder.jit(autocast(foo, dtype=thunder.dtypes.bfloat16))(x)
# NotImplementedError: Attempting to execute outside of a tracing context, which is not supported
# o = autocast(thunder.jit(foo), dtype=thunder.dtypes.bfloat16)(x)
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 grad and autocast transform entry points and the thunder.jit behavior shown in the issue's Python reproducer. Run the three working and failing examples to trace their ordering and tracing-context differences; done means transform application has consistent, documented behavior without the reported runtime errors.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100