Lightning-AI / Lightning-AI/lightning-thunder
[type promotion] Incorrect dtype when using `div` (or other arithmetic operation) with scalar tensor.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Found while running torchbench's `hf_Reformer` with thunderFX path (see below for the larger graph that fails).
```python
import torch
import thunder
def forward(x, double_scalar_tensor):
return x / double_scalar_tensor
x = torch.randn(3, requires_grad=True)
# x = torch.randn(3,) # This works fine if `requires_grad=False`
y = torch.randn((), dtype=torch.float64)
e = forward(x, y) # e.dtype=torch.float
a = thunder.jit(forward)(x, y) # a.dtype=torch.float64
torch.testing.assert_close(a, e)
```
Error
```python
Traceback (most recent call last):
File "/opt/pytorch/lightning-thunder/test.py", line 63, in
torch.testing.assert_close(a, e)
File "/usr/local/lib/python3.10/dist-packages/torch/testing/_comparison.py", line 1530, in assert_close
raise error_metas[0].to_error(msg)
AssertionError: The values for attribute 'dtype' do not match: torch.float64 != torch.float32.
```
Computation Trace
```python
@torch.no_grad()
@no_autocast
def computation(x, double_scalar_tensor):
# x: "cpu f32[3]"
# double_scalar_tensor: "cpu f64[]"
t2 = Tensor.to(x, torch.float64, copy=True) # t2: "cpu f64[3]"
# t2 = ltorch.to(x, torch.float64, None, device=None, dtype=None, copy=True, memory_format=None) # t2: "cpu f64[3]"
# t2 = prims.convert_element_type(x, dtypes.float64) # t2: "cpu f64[3]"
t1 = torch.true_divide(t2, double_scalar_tensor) # t1: "cpu f64[3]"
# t1 = ltorch.true_divide(t2, double_scalar_tensor) # t1: "cpu f64[3]"
# t1 = prims.div(t2, double_scalar_tensor) # t1: "cpu f64[3]"
del t2
return {'output': t1, 'flat_args': [x, double_scalar_tensor], 'flat_output': (t1,)}, ((double_scalar_tensor,), ())
```
Larger Repro
```python
import torch
from torch import device
import thunder
def forward(self, key_vectors_1: "f32[4, 12, 4096, 64]", wrapped_sqrt: "f64[]", query_vectors_1: "f32[4, 12, 4096, 64]", value_vectors_1: "f32[4, 12, 4096, 64]", l_self_modules_attention_modules_output_modules_dense_parameters_weight_: "f32[256, 768]", l_prev_attn_output_: "f32[4, 4096, 256]"):
# File: /usr/local/lib/python3.10/dist-packages/transformers/models/reformer/modeling_reformer.py:1115 in forward, code: key_vectors = key_vectors / np.sqrt(self.attention_head_size)
key_vectors_2: "f32[4, 12, 4096, 64]" = key_vectors_1 / wrapped_sqrt; key_vectors_1 = wrapped_sqrt = None
# File: /usr/local/lib/python3.10/dist-packages/transformers/models/reformer/modeling_reformer.py:322 in _split_seq_length_dim_to, code: return torch.reshape(vectors, split_dim_shape + (attn_head_size,))
query_vectors_2: "f32[4, 12, 64, 64, 64]" = torch.reshape(query_vectors_1, (4, 12, -1, 64, 64)); query_vectors_1 = None
# File: /usr/local/lib/python3.10/dist-packages/transformers/models/reformer/modeling_reformer.py:322 in _split_seq_length_dim_to, code: return torch.reshape(vectors, split_dim_shape + (attn_head_size,))
key_vectors_3: "f32[4, 12, 64, 64, 64]" = torch.reshape(key_vectors_2, (4, 12, -1, 64, 64)); key_vectors_2 = None
# File: /usr/local/lib/python3.10/dist-packages/transformers/models/reformer/modeling_reformer.py:296 in _look_adjacent, code: slices.append(torch.cat([vectors[:, :, i:, ...], vectors[:, :, :i, ...]], dim=2))
getitem: "f32[4, 12, 1, 64, 64]" = key_vectors_3[(slice(None, None, None), slice(None, None, None), slice(-1, None, None), Ellipsis)]
getitem_1: "f32[4, 12, 63, 64, 64]" = key_vectors_3[(slice(None, None, None), slice(None, None, None), slice(None, -1, None), Ellipsis)]
cat: "f32[4, 12, 64, 64, 64]" = torch.cat([getitem, getitem_1], dim = 2); getitem = getitem_1 = None
# File: /usr/local/lib/python3.10/dist-packages/transformers/models/reformer/modeling_reformer.py:297 in _look_adjacent, code: return torch.cat(slices, dim=3)
key_vectors_4: "f32[4, 12, 64, 128, 64]" = torch.cat([cat, key_vectors_3], dim = 3); cat = key_vectors_3 = None
# File: /usr/local/lib/python3.10/dist-packages/transformers/models/reformer/modeling_reformer.py:1163 in forward, code: query_key_dots = torch.matmul(query_vectors, key_vectors.transpose(-1, -2))
transpose_3: "f32[4, 12, 64, 64, 128]" = key_vectors_4.transpose(-1, -2); key_vectors_4 = None
query_key_dots: "f32[4, 12, 64, 64, 128]" = torch.matmul(query_vectors_2, transpose_3); query_vectors_2 = transpose_3 = None
return query_key_dots
key_vectors_1 = torch.randn([4, 12, 4096, 64], requires_grad=True)
wrapped_sqrt = torch.randn([], dtype=torch.float64)
query_vectors_1 = torch.randn([4, 12, 4096, 64])
value_vectors_1 = torch.randn([4, 12, 4096, 64])
l_self_modules_attention_modules_output_modules_dense_parameters_weight_ = torch.randn([256, 768])
l_prev_attn_output_ = torch.randn([4, 4096, 256])
e = forward(None, key_vectors_1, wrapped_sqrt, query_vectors_1,
value_vectors_1, l_self_modules_attention_modules_output_modules_dense_parameters_weight_,
l_prev_attn_output_)
a = thunder.jit(forward)(None, key_vectors_1, wrapped_sqrt, query_vectors_1,
value_vectors_1, l_self_modules_attention_modules_output_modules_dense_parameters_weight_,
l_prev_attn_output_)
print(e.dtype, a.dtype)
```
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the minimal forward example using x / double_scalar_tensor, then compare eager PyTorch with thunder.jit and inspect the generated computation trace. Check related arithmetic and type-promotion entry points; done means scalar-tensor division and other arithmetic preserve the eager result dtype, including the shown float32 and float64 case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100