`nvfuserex.fusion_pass` seems to fail when input tra... | Lightning-AI/lightning-thunder

Nobody has claimed this yet.

## 🐛 Bug

### To Reproduce

`python thunder/benchmarks/benchmark_peft.py --model deepseek-ai/DeepSeek-V2.5 --trust-remote-code --attn-implementation eager --max-steps 100 --mbs 1 --seq-length 128 --compile thunder --fixed-num-hidden-layers 2` would result in the following error

```
Traceback (most recent call last):
File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_peft.py", line 710, in
main(args)
File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_peft.py", line 555, in main
outputs = model(**batch)
^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 372, in __call__
return super().__call__(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 699, in compile_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/peft/peft_model.py", line 1845, in forward
return self.base_model(
^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/peft/tuners/tuners_utils.py", line 216, in forward
return self.model.forward(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.cache/huggingface/modules/transformers_modules/deepseek-ai/DeepSeek-V2.5/c85b5ede86f2a598af339624cac5723861e557ed/modeling_deepseek.py", line 1675, in forward
outputs = self.model(
^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.cache/huggingface/modules/transformers_modules/deepseek-ai/DeepSeek-V2.5/c85b5ede86f2a598af339624cac5723861e557ed/modeling_deepseek.py", line 1544, in forward
layer_outputs = decoder_layer(
^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.cache/huggingface/modules/transformers_modules/deepseek-ai/DeepSeek-V2.5/c85b5ede86f2a598af339624cac5723861e557ed/modeling_deepseek.py", line 1272, in forward
hidden_states = self.mlp(hidden_states)
^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.cache/huggingface/modules/transformers_modules/deepseek-ai/DeepSeek-V2.5/c85b5ede86f2a598af339624cac5723861e557ed/modeling_deepseek.py", line 571, in forward
topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/.cache/huggingface/modules/transformers_modules/deepseek-ai/DeepSeek-V2.5/c85b5ede86f2a598af339624cac5723861e557ed/modeling_deepseek.py", line 420, in forward
def forward(self, hidden_states):
File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 372, in __call__
return super().__call__(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 893, in _fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 840, in call_wrapped
return self._wrapped_call(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 416, in __call__
raise e
File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 403, in __call__
return super(self.cls, obj).__call__(*args, **kwargs) # type: ignore[misc]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".54", line 5, in forward
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/module.py", line 80, in forward
res = self._forward_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 839, in wrapped
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 879, in fn_
cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 818, in wrapped
cache_entry, inps, pro_to_epi = get_computation_and_inputs_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/langctxs.py", line 135, in _fn
result = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 238, in cache_info_wrapper
res = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 784, in get_computation_and_inputs
cache_entry = apply_transforms_and_build_cache_entry(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 589, in apply_transforms_and_build_cache_entry
extraces = transform_for_execution(
^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/common.py", line 654, in transform_for_execution
extrace = executors.passes.transform_for_execution(trace, executors_list)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/passes.py", line 132, in transform_for_execution
extrace = ex.fusion_pass(extrace)
^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/nvfuserex_impl.py", line 1045, in fusion_pass
fusedtrace = rematerialize(fusedtrace)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/rematerialization.py", line 594, in rematerialize
cut = find_cut(external_producer_outputs, current_producer, current_consumer)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/rematerialization.py", line 370, in find_cut
add_edge("source", var.name + "_in", capacity=float("inf"))
^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'name'
```

#### Relevant traces

This is the input trace to `ex.fusion_pass`:
```python
### Trace before thunder.executors.nvfuserex_impl.nvFuserExecutor('nvfuser').fusion_pass
# Constructed by Dead Code Elimination (took 1 milliseconds)
import thunder
import thunder.core.devices as devices
import thunder.core.dtypes as dtypes
import thunder.core.prims as prims
import thunder.torch as ltorch
import torch
import torch.nn.functional
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(l_hidden_states_, l_self_parameters_weight_):
# l_hidden_states_: "cuda:0 bf16[1, 128, 5120]"
# l_self_parameters_weight_: "cuda:0 bf16[160, 5120]"

# .50:5: hidden_states = l_hidden_states_.view(-1, 5120); l_hidden_states_ = None
t32 = prims.reshape(l_hidden_states_, (128, 5120)) # t32: "cuda:0 bf16[128, 5120]"

# .50:6: type_1 = hidden_states.type(torch.float32); hidden_states = None
t35 = prims.convert_element_type(t32, dtypes.float32) # t35: "cuda:0 f32[128, 5120]"

# .50:7: type_2 = l_self_parameters_weight_.type(torch.float32); l_self_parameters_weight_ = None
t38 = prims.convert_element_type(l_self_parameters_weight_, dtypes.float32) # t38: "cuda:0 f32[160, 5120]"

# .50:8: logits = torch._C._nn.linear(type_1, type_2, None); type_1 = type_2 = None
t133 = torch.nn.functional.linear(t35, t38, None) # t133: "cuda:0 f32[128, 160]"
# t133 = ltorch.linear(t35, t38, None) # t133: "cuda:0 f32[128, 160]"
# t133 = prims.linear(t35, t38, None) # t133: "cuda:0 f32[128, 160]"

# .50:9: scores = logits.softmax(dim = -1, dtype = torch.float32); logits = None
t57 = ltorch._softmax(t133, -1, dtype=torch.float32) # t57: "cuda:0 f32[128, 160]"
# t50 = ltorch.amax(t133, -1, True) # t50: "cuda:0 f32[128, 1]"
# t49 = prims.amax(t133, (1,)) # t49: "cuda:0 f32[128]"
# t50 = prims.broadcast_in_dim(t49, [128, 1], [0]) # t50: "cuda:0 f32[128, 1]"
# t52 = ltorch.sub(t133, t50, alpha=1) # t52: "cuda:0 f32[128, 160]"
# t51 = prims.broadcast_in_dim(t50, (128, 160), (0, 1)) # t51: "cuda:0 f32[128, 160]"
# t52 = prims.sub(t133, t51) # t52: "cuda:0 f32[128, 160]"
# t53 = ltorch.exp(t52) # t53: "cuda:0 f32[128, 160]"
# t53 = prims.exp(t52) # t53: "cuda:0 f32[128, 160]"
# t55 = ltorch.sum(t53, -1, True, dtype=None) # t55: "cuda:0 f32[128, 1]"
# t54 = prims.sum(t53, (1,)) # t54: "cuda:0 f32[128]"
# t55 = prims.broadcast_in_dim(t54, [128, 1], [0]) # t55: "cuda:0 f32[128, 1]"
# t57 = ltorch.true_divide(t53, t55) # t57: "cuda:0 f32[128, 160]"
# t56 = prims.broadcast_in_dim(t55, (128, 160), (0, 1)) # t56: "cuda:0 f32[128, 160]"
# t57 = prims.div(t53, t56) # t57: "cuda:0 f32[128, 160]"

# .50:10: view_1 = scores.view(128, 8, -1)
t65 = prims.reshape(t57, (128, 8, 20)) # t65: "cuda:0 f32[128, 8, 20]"

# .50:11: max_1 = view_1.max(dim = -1); view_1 = None
t68 = prims.amax(t65, (2,)) # t68: "cuda:0 f32[128, 8]"

# .50:13: topk = torch.topk(group_scores, k = 3, dim = -1, sorted = False)
(_, t85) = prims.topk(t68, 3, 1, True, False)

# .50:15: group_mask = torch.zeros_like(group_scores); group_scores = None
group_mask = prims.full((128, 8), 0, device=devices.Device("cuda:0"), dtype=dtypes.float32) # group_mask: "cuda:0 f32[128, 8]"
(t134,) = update_aliases((group_mask,))

# .50:16: scatter_ = group_mask.scatter_(1, group_idx, 1); group_idx = scatter_ = None
t136 = torch.scatter(t134, 1, t85, 1) # t136: "cuda:0 f32[128, 8]"
# t136 = ltorch.scatter(t134, 1, t85, 1, value=None, reduce=None) # t136: "cuda:0 f32[128, 8]"
# t136 = prims.scatter(t134, t85, 1, 1) # t136: "cuda:0 f32[128, 8]"
t96 = prims.copy_(t136, t134, grad_enabled=True) # t96: "cuda:0 f32[128, 8]"
(t137,) = update_aliases((t96,))

# .50:17: unsqueeze = group_mask.unsqueeze(-1); group_mask = None
t100 = prims.broadcast_in_dim(t137, [128, 8, 1], [0, 1]) # t100: "cuda:0 f32[128, 8, 1]"
(t138, t139) = update_aliases((t100, t137))

# .50:18: expand = unsqueeze.expand(128, 8, 20); unsqueeze = None
t107 = prims.broadcast_in_dim(t138, (128, 8, 20), (0, 1, 2)) # t107: "cuda:0 f32[128, 8, 20]"
(t140, _, _) = update_aliases((t107, t138, t139))

# .50:19: score_mask = expand.reshape(128, -1); expand = None
t117 = prims.reshape(t140, (128, 160)) # t117: "cuda:0 f32[128, 160]"
t58 = prims.get_grad(t57) # t58: "cuda:0 f32[128, 160]"
t59 = ltorch.mul(t57, t58) # t59: "cuda:0 f32[128, 160]"
# t59 = prims.mul(t57, t58) # t59: "cuda:0 f32[128, 160]"
t61 = ltorch.sum(t59, -1, True, dtype=None) # t61: "cuda:0 f32[128, 1]"
# t60 = prims.sum(t59, (1,)) # t60: "cuda:0 f32[128]"
# t61 = prims.broadcast_in_dim(t60, [128, 1], [0]) # t61: "cuda:0 f32[128, 1]"
t63 = ltorch.sub(t58, t61, alpha=1) # t63: "cuda:0 f32[128, 160]"
# t62 = prims.broadcast_in_dim(t61, (128, 160), (0, 1)) # t62: "cuda:0 f32[128, 160]"
# t63 = prims.sub(t58, t62) # t63: "cuda:0 f32[128, 160]"
t64 = ltorch.mul(t57, t63) # t64: "cuda:0 f32[128, 160]"
# t64 = prims.mul(t57, t63) # t64: "cuda:0 f32[128, 160]"
t43 = ltorch.reshape(t64, -1, 160) # t43: "cuda:0 f32[128, 160]"
# t43 = prims.reshape(t64, (128, 160)) # t43: "cuda:0 f32[128, 160]"
t143 = torch.matmul(t43, t38) # t143: "cuda:0 f32[128, 5120]"
# t143 = ltorch.matmul(t43, t38) # t143: "cuda:0 f32[128, 5120]"
# t143 = prims.matmul(t43, t38) # t143: "cuda:0 f32[128, 5120]"
t37 = prims.convert_element_type(t143, dtypes.bfloat16) # t37: "cuda:0 bf16[128, 5120]"
t34 = prims.reshape(t37, (1, 128, 5120)) # t34: "cuda:0 bf16[1, 128, 5120]"
return {'output': (t34, None), 'flat_args': [l_hidden_states_, l_self_parameters_weight_], 'fw_flat_out': (t117, t57)}
```

and this is the input trace to `find_cut`
```python
def computation(l_hidden_states_, l_self_parameters_weight_):
# l_hidden_states_: "cuda:0 bf16[1, 128, 5120]"
# l_self_parameters_weight_: "cuda:0 bf16[160, 5120]"