Lightning-AI / Lightning-AI/lightning-thunder
enter/exit_autocast of torch.amp.autocast_mode
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'm trying to get a fuller picture of what we need to support NeVA. As such I'm using:
```python
def thunder_backend(gm, args):
gm.real_recompile()
from thunder.examine import examine
try: # Examine may raise an exception
thunder.examine.examine(gm, *args)
except Exception as e:
print(f"Hit problem with examine:\n{e}")
# Don't really use Thunder just return the original graph
return gm
...
#model.model = thunder.jit(model.model)
model.model = torch.compile(backend=thunder_backend)(model.model)
```
(thanks Ivan for the great idea!)
And one of the issues that gets reported is e.g.
```
Found 2 distinct operations, of which 0 (0.0%) are supported
Please file an issue requesting the following operators here: https://github.com/Lightning-AI/lightning-thunder/issues/
new
_enter_autocast of torch.amp.autocast_mode
_exit_autocast of torch.amp.autocast_mode
```
### Pitch
This looks like it is going to be important for #343.
### Alternatives / Potential work-arounds
### Minimal Repro
cc @apaz-cli @crcrpar @tfogal
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 by reproducing the reported torch.compile case with the thunder_backend example and inspect how torch.amp.autocast_mode operations are handled. The work is complete when _enter_autocast and _exit_autocast are supported and no longer appear in the unsupported-operations report.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100