Lightning-AI / Lightning-AI/lightning-thunder

fusion_type="dataflow" can lead to invalid trace

Open
#1,858 1 comment 1 reaction 1 assignee View on GitHub

@IvanYashchuk is already working on this.

Since Mar 26, 2025.

fusion logic in-place
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

  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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.