Lightning-AI / Lightning-AI/lightning-thunder

Errors on in-place ops on tensor aliases unresolved by proxy substitution

Open
#2,766 13 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

aliasing in-place
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

Repro:

```py
import torch, thunder

def f(a, b):
return a.exp_() * b.tanh_()

def g(a, _):
b = a.view(5,5)
return a.exp_() * b.tanh_()

def h(a, _):
b = a[0,0]
return a.exp_() * b.tanh_()

for fn in [f, g, h]:
jf = thunder.jit(fn)
x = torch.randn(5, 5, device='cuda')
x_ = x.detach().clone()
out = jf(x, x[0, 0])
out_ = fn(x_, x_[0, 0])

torch.testing.assert_close(out, out_)
# AssertionError on f, g and h

```

Found in https://github.com/Lightning-AI/lightning-thunder/pull/2760#issuecomment-3562836815 and https://github.com/Lightning-AI/lightning-thunder/pull/2760#issuecomment-3562688372 by @beverlylytle.

Trace of `f` after `update_aliases.py`:

```py
# Constructed by Update aliases for 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 f32[5, 5]"
# b: "cuda:0 f32[]"
(t6,) = prims.update_aliases((a,))

# /opt/pytorch/lightning-thunder/tmp/main.py:4: return a.exp_() * b.tanh_()
t1 = ltorch.exp_(t6) # t1: "cuda:0 f32[5, 5]"
# t0 = ltorch.exp(t6) # t0: "cuda:0 f32[5, 5]"
# t0 = prims.exp(t6) # t0: "cuda:0 f32[5, 5]"
# t1 = prims.copy_(t0, t6, grad_enabled=True) # t1: "cuda:0 f32[5, 5]"
(t7,) = prims.update_aliases((b,))

# /opt/pytorch/lightning-thunder/tmp/main.py:4: return a.exp_() * b.tanh_()
t3 = ltorch.tanh_(t7) # t3: "cuda:0 f32[]"
# t2 = ltorch.tanh(t7) # t2: "cuda:0 f32[]"
# t2 = prims.tanh(t7) # t2: "cuda:0 f32[]"
# t3 = prims.copy_(t2, t7, grad_enabled=True) # t3: "cuda:0 f32[]"
t5 = ltorch.mul(t1, t3) # t5: "cuda:0 f32[5, 5]"
# t4 = prims.broadcast_in_dim(t3, (5, 5), ()) # t4: "cuda:0 f32[5, 5]"
# t5 = prims.mul(t1, t4) # t5: "cuda:0 f32[5, 5]"
return {'output': (t5,), 'flat_args': [t1, t3]}
```

Trace after fusion:

```py
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(a, b):
# a: "cuda:0 f32[5, 5]"
# b: "cuda:0 f32[]"
(t10,) = update_aliases((a,))
del a
[t1] = nvFusion0(t10)
# t0 = prims.exp(t10) # t0: "cuda:0 f32[5, 5]"
# t1 = prims.copy_(t0, t10, grad_enabled=True) # t1: "cuda:0 f32[5, 5]"
del t10
(t11,) = update_aliases((b,))
del b
[t3, t5] = nvFusion1(t11, t1)
# t2 = prims.tanh(t11) # t2: "cuda:0 f32[]"
# t3 = prims.copy_(t2, t11, grad_enabled=True) # t3: "cuda:0 f32[]"
# t4 = prims.broadcast_in_dim(t3, (5, 5), ()) # t4: "cuda:0 f32[5, 5]"
# t5 = prims.mul(t1, t4) # t5: "cuda:0 f32[5, 5]"
del t11
return {'output': (t5,), 'flat_args': [t1, t3]}
```

The problem here is that `nvFusion1` does not know that `t11` and `t1` share memory.

In order to make sure that `t3 = prims.copy_(t2, t11, grad_enabled=True)` completes before `t5 = prims.mul(t1, t4)`, we could insert `prims.update_aliases` before `prims.mul`, which would fix the bugs because `prims.update_aliases` is unfusible.

Such solutions create more fusion breaks, so we want to minimize the use of `prims.update_aliases`. Ideally, we hope to make `prims.update_aliases` a fusible op and let nvFuser handle memory aliases in its combined region.

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 update_aliases.py and the implementation of prims.update_aliases, then trace how nvFusion regions are formed around the shown f, g, and h reproductions. Verify that in-place alias dependencies preserve the required ordering without inserting unnecessary fusion breaks; the reproductions should match their eager results.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, performance
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.