Lightning-AI / Lightning-AI/lightning-thunder
sdpa_ex - Incorrect device in trace vs from actual computation
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
`sdpa_ex` implementation of `torch.nn.functional.scaled_dot_product_attention` returns all output tensor proxy in trace to be on `cuda` but at runtime some outputs are on `cpu`.
Repro
```python
import torch
import thunder
torch.backends.cuda.enable_cudnn_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(True)
def fn(q, k, v):
return torch.nn.functional.scaled_dot_product_attention(q, k, v, None, 0.0, is_causal=True, scale=0.08838834764831843)
q = torch.randn(8, 32, 1024, 128, device='cuda', requires_grad=True)
k = torch.randn(8, 32, 1024, 128, device='cuda', requires_grad=True)
t = torch.randn(8, 1024, 12288, device='cuda', requires_grad=False)
v = torch.as_strided(t, (8, 32, 1024, 128), ((12582912, 384, 12288, 1)))
jfn = thunder.jit(fn)
o = jfn(q, k, v)
extrace = thunder.last_traces(jfn)[-1]
print(extrace)
saved_tensors_trace = extrace.bound_symbols[4].args[1][0]
print("TRACE OUTPUT DEVICES", list(t.device.device_str() for t in saved_tensors_trace))
cache_entry, inps, pro_to_epi = jfn._lc_cd._get_computation_and_inputs(q, k, v)
output = cache_entry.computation_fn(*inps)
data_for_autograd, (saved_tensors, saved_other) = output
print("ACTUAL SAVED DEVICES", list(str(o.device) for o in saved_tensors))
```
Output
```python
TRACE OUTPUT DEVICES ['cuda:0', 'cuda:0', 'cuda:0', 'cuda:0', 'cuda:0', 'cuda:0', 'cuda:0']
ACTUAL SAVED DEVICES ['cuda:0', 'cuda:0', 'cuda:0', 'cuda:0', 'cpu', 'cpu', 'cuda:0']
```
Trace
```python
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def augmented_forward_fn(q, k, v):
# q: "cuda:0 f32[8, 32, 1024, 128]"
# k: "cuda:0 f32[8, 32, 1024, 128]"
# v: "cuda:0 f32[8, 32, 1024, 128]"
(t0, t1, t2, t3) = sdpaex_grad_forward_scaled_dot_product_efficient_attention(q, k, v, None, 0.0, True, 0.08838834764831843)
return {'output': t0, 'flat_args': [q, k, v], 'flat_output': (t0,)}, ((k, q, t0, t1, t2, t3, v), ())
```
cc @carmocca
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.