Lightning-AI / Lightning-AI/lightning-thunder

[type promotion] Incorrect dtype when using `div` (or other arithmetic operation) with scalar tensor.

Open
#1,239 5 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

huggingface type promotion
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.