Lightning-AI / Lightning-AI/lightning-thunder

enter/exit_autocast of torch.amp.autocast_mode

Open
#824 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

amp dynamo nemo 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'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

  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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.