Lightning-AI / Lightning-AI/lightning-thunder

A reference cycle is detected related to the ThunderModule

Open
#1,034 5 comments 0 reactions 1 assignee View on GitHub

@t-vi is already working on this.

Since Aug 26, 2024.

bug memory use thunderfx
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")
```
![image](https://github.com/user-attachments/assets/33c2550d-d84c-45c6-8989-d9f77e3ae7b2)

cc @apaz-cli

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.