Lightning-AI / Lightning-AI/lightning-thunder
Inplace functionalization misplaced cast
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Inplace where the inplace-modified operand is lower precision than other operand needs to put the cast outside the out of place operator, not in its subsymbols.
```python
def test_fn(a, b):
a += b
c = a.clone()
return c
a = torch.randn(5, dtype=torch.bfloat16, device="cuda")
b = torch.randn(5, dtype=torch.float32, device="cuda", requires_grad=True)
jfn = thunder.jit(test_fn)
jfn(a, b)
thunder.last_traces(jfn)[4]
```
```python
# Constructed by Functionalize in-place ops
import thunder
import thunder.core.prims as prims
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(a, b):
# a: "cuda:0 bf16[5]"
# b: "cuda:0 f32[5]"
# Functionalized from `t3 = add_(a,b,1)`
t2 = ltorch.add(a, b, alpha=1) # t2: "cuda:0 bf16[5]"
# t1 = ltorch.add(a, b, alpha=1) # t1: "cuda:0 f32[5]"
# t0 = prims.convert_element_type(a, dtypes.float32) # t0: "cuda:0 f32[5]"
# t1 = prims.add(t0, b) # t1: "cuda:0 f32[5]"
# t2 = prims.convert_element_type(t1, dtypes.bfloat16) # t2: "cuda:0 bf16[5]" <--------- This cast is not part of .add so it should not be in the subsymbols.
# /tmp/ipykernel_97531/446867340.py:3: c = a.clone()
# `t3` is replaced by `t2`
c = ltorch.clone(t2, memory_format=_torch_memory_format_0) # c: "cuda:0 bf16[5]"
# c = prims.clone(t2) # c: "cuda:0 bf16[5]"
# /tmp/ipykernel_97531/446867340.py:2: a += b
t3 = prims.copy_(t2, a, grad_enabled=True) # t3: "cuda:0 bf16[5]"
return {'output': (c,), 'flat_args': [t3, b]}
```
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
Reproduce the issue with the provided Python test_fn, torch tensors, thunder.jit, and thunder.last_traces(jfn)[4]. Start at the Functionalize handling of ltorch.add, prims.convert_element_type, and prims.copy_; done means the lower-precision cast is outside the out-of-place add rather than inside its subsymbols.
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
- 42/100