Lightning-AI / Lightning-AI/lightning-thunder

sdpa_ex - Incorrect device in trace vs from actual computation

Open
#950 4 comments 0 reactions 1 assignee Claimed by @IvanYashchuk View on GitHub
sdpa triage review
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.