Lightning-AI / Lightning-AI/lightning-thunder
fusion_type="dataflow" can lead to invalid trace
@IvanYashchuk is already working on this.
Since Mar 26, 2025.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
NOTE - WAR is to use `fusion_type="consecutive"`.
Found while working on #1850
```python
import torch
import thunder
def fwd_body_0(ctx : torch.autograd.function.Function, values: "bf16[86, 32, 96]", indices: "i64[86]"):
# File: /home/kkalambarkar/git/pytorch/torch/utils/_device.py:106 in __torch_function__, code: return func(*args, **kwargs)
output: "bf16[512, 32, 96]" = torch.zeros(512, 32, 96, device = torch.device(type='cuda', index=0), dtype = torch.bfloat16)
# File: /home/kkalambarkar/git/pytorch/torch/utils/_device.py:106 in __torch_function__, code: return func(*args, **kwargs)
output[indices] = values; setitem = output; values = setitem = None
return (output, [indices])
def bwd_body_0(ctx : torch.autograd.function.Function, grad_output: "bf16[512, 32, 96]", indices: "i64[86]"):
# No stacktrace found for following nodes
_set_grad_enabled = torch._C._set_grad_enabled(False); _set_grad_enabled = None
# File: /home/kkalambarkar/git/pytorch/torch/utils/_device.py:106 in __torch_function__, code: return func(*args, **kwargs)
grad_values: "bf16[86, 32, 96]" = grad_output[indices]; grad_output = indices = None
# No stacktrace found for following nodes
_set_grad_enabled_1 = torch._C._set_grad_enabled(True); _set_grad_enabled_1 = None
return (grad_values, None)
def forward(l_stack0_: "bf16[86, 32, 96]", l_indices_q_: "i64[86]"):
# File: /home/kkalambarkar/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/flash_attn-2.4.2-py3.10-linux-x86_64.egg/flash_attn/bert_padding.py:208 in pad_input, code: output = index_put_first_axis(hidden_states, indices, batch * seqlen)
output: "bf16[512, 32, 96]" = torch.ops.higher_order.autograd_function_apply(fwd_body_0, bwd_body_0, l_stack0_, l_indices_q_, args_tensor_mask = [True, True, False], non_differentiable_idx = [])
return output
with torch.device("cuda"):
l_stack0_ = torch.randn(86, 32, 96, dtype=torch.bfloat16, requires_grad=True)
l_indices_q_ = torch.randint(0, 10, (86,))
thunder.jit(forward, fusion_type="consecutive")(l_stack0_, l_indices_q_) # Works
thunder.jit(forward, fusion_type="dataflow")(l_stack0_, l_indices_q_)
```
Error
```python
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 631, in get_computation_and_inputs
computation_trc, backward_trc = split_forward_backward(computation_trc, cd, cs, *inps)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/torch_autograd.py", line 283, in split_forward_backward
fw_extrace = transform_for_execution(
^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/passes.py", line 127, in transform_for_execution
extrace = ex.fusion_pass(extrace)
^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/nvfuserex_impl.py", line 936, in fusion_pass
fusedtrace = self.cse(fusedtrace)
^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/nvfuserex_impl.py", line 809, in cse
assert return_bsym.sym.id == prims.PrimIDs.RETURN
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
```
The trace before the failing CSE pass
```python
# Constructed by Remove redundant casts (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(l_stack0_, l_indices_q_):
# l_stack0_: "cuda:0 bf16[86, 32, 96]"
# l_indices_q_: "cuda:0 i64[86]"
[t10] = nvFusion0()
# t10 = prims.full((512, 32, 96), 0, device=devices.Device("cuda:0"), dtype=dtypes.bfloat16) # t10: "cuda:0 bf16[512, 32, 96]"
t19 = copy_with_setitem_impl(t10, l_indices_q_, l_stack0_) # t19: "cuda:0 bf16[512, 32, 96]"
return {'output': (t10,), 'flat_args': [l_stack0_, l_indices_q_], 'flat_output': (t10,)}, ((l_indices_q_,), ())
nvFusion1(t19, t10)
# prims.copy_(t19, t10, grad_enabled=True)
```
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.
Assessment
This issue has not been assessed yet.