Lightning-AI / Lightning-AI/lightning-thunder
autocast is incorrectly applied even if the requested device is different.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
From the example below, the autocast is applied only for device cuda, however thunder.jit still applies it to CPU inputs.
```python
import thunder
import torch
def foo(x, w):
return torch.nn.functional.linear(x, w)
device = torch.device("cpu")
with device:
x, w = torch.randn(16, 16), torch.randn(16, 16)
print(x.dtype, w.dtype)
jfoo = thunder.jit(foo)
# Autocast is applied to different device.
with torch.autocast("cuda", torch.bfloat16):
jit_out = jfoo(x, w)
print(thunder.last_traces(jfoo)[-1])
```
Output
```python
# Constructed by Delete Last Used (took 0 milliseconds)
from torch import Tensor
import torch
import torch.nn.functional
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(x, w):
# x: "cpu f32[16, 16]"
# w: "cpu f32[16, 16]"
t0 = Tensor.to(x, torch.bfloat16, copy=True) # t0: "cpu bf16[16, 16]"
# t0 = ltorch.to(x, torch.bfloat16, None, device=None, dtype=None, copy=True, memory_format=None) # t0: "cpu bf16[16, 16]"
# t0 = prims.convert_element_type(x, dtypes.thunder.dtypes.bfloat16) # t0: "cpu bf16[16, 16]"
del x
t1 = Tensor.to(w, torch.bfloat16, copy=True) # t1: "cpu bf16[16, 16]"
# t1 = ltorch.to(w, torch.bfloat16, None, device=None, dtype=None, copy=True, memory_format=None) # t1: "cpu bf16[16, 16]"
# t1 = prims.convert_element_type(w, dtypes.thunder.dtypes.bfloat16) # t1: "cpu bf16[16, 16]"
del w
t2 = torch.nn.functional.linear(t0, t1, None) # t2: "cpu bf16[16, 16]"
# t2 = ltorch.linear(t0, t1, None) # t2: "cpu bf16[16, 16]"
# t2 = prims.linear(t0, t1, None) # t2: "cpu bf16[16, 16]"
del t0, t1
return t2
```
cc @crcrpar
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the thunder.jit entry point and reproduce the example using CPU tensors inside a CUDA-only torch.autocast context. The fix is complete when CPU inputs are not converted to bfloat16 for an autocast device that differs from their device, while matching-device autocast behavior remains intact.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100