Lightning-AI / Lightning-AI/lightning-thunder
A reference cycle is detected related to the ThunderModule
@t-vi is already working on this.
Since Aug 26, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.
## 🐛 Bug
### To Reproduce
```
import torch
import thunder
import weakref
import gc
mod = torch.nn.ReLU()
ref = weakref.ref(mod, lambda _: print("mod deleted!"))
opt_mod = thunder.jit(mod)
# opt_mod = torch.compile(mod)
ref_opt_mod = weakref.ref(opt_mod, lambda _: print("opt_mod deleted!"))
x = torch.randn(10, 10)
refx = weakref.ref(x, lambda _: print("x deleted!"))
opt_mod(x)
del x
del mod
del opt_mod
# gc.collect()
print("done!") # done!
```
with the line of torch.comile(), it outputs:
```
x deleted!
opt_mod deleted!
mod deleted!
done!
```
with thunder it outputs:
```
x deleted!
done!
mod deleted!
opt_mod deleted!
```
@kshitij12345 detected there's a reference cycle:
```
import torch
import thunder
import weakref
import gc
mod = torch.nn.ReLU()
ref = weakref.ref(mod, lambda _: print("mod deleted!"))
opt_mod = thunder.jit(mod)
# opt_mod = torch.compile(mod)
ref_opt_mod = weakref.ref(opt_mod, lambda _: print("opt_mod deleted!"))
x = torch.randn(10, 10)
refx = weakref.ref(x, lambda _: print("x deleted!"))
opt_mod(x)
del x
del mod
del opt_mod
# gc.collect()
print("done!") # done!
if ref_opt_mod() is not None:
import refcycle
graph = refcycle.snapshot()
try:
cycle = graph.shortest_cycle(ref_opt_mod())
print("CYCLE FOUND FROM MOD")
except ValueError:
print("NO CYCLE FROM MOD")
pass
# Save the latest cycle.
cycle.export_json("cycle.json")
cycle.export_image("cycle.png")
```

cc @apaz-cli
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.