Remove reference cycle

[anijain2305]

For the following test, Dynamo was holding on to more memory than necessary. And this memory was released after calling gc.collect(). This indicated that their was a reference cycle somewhere.

def print_mem(name):
    print(name, torch.cuda.memory_allocated() / 10**9, "GB")


def fn1():
    x = torch.randn(1024, 1024, 1024, device="cuda")
    print_mem("Inside func")
    return x.sum()


fn1()
print_mem("Outside func")
print("------ Eager Done --------")
print("\n\n\n")

# torchdynamo.config.debug = True
# torchdynamo.config.trace = True

with torchdynamo.optimize("eager"):
    fn1()
    print_mem("Outside func")
print("------ TorchDynamo Done --------")
print("\n\n\n")


graph = refcycle.garbage()
sccs = graph.strongly_connected_components()

refcycles = graph.source_components()
for idx, r in enumerate(refcycles):
    r.export_image(f"cycle{idx}.svg")

print(refcycles)

I used refcycle package to find the reference cycles, and following is the culprit.

This PR does the following

• Breaks the cycle by setting output_graph.root_tx = None after we are done with tracing
• Also, cleanup the graphargs, the data structure that was holding on to the tensor.

[anijain2305]

Cross referencing #69 and pytorch/pytorch#93751

[jansel]

Awesome, glad we found this one!