Lightning-AI / Lightning-AI/lightning-thunder

autocast is incorrectly applied even if the requested device is different.

Open
#709 5 comments 0 reactions 0 assignees View on GitHub
amp
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.