Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Value propagation produces a large memory footprint #187

Open
neNasko1 opened this issue Oct 23, 2024 · 2 comments · May be fixed by #189
Open

Value propagation produces a large memory footprint #187

neNasko1 opened this issue Oct 23, 2024 · 2 comments · May be fixed by #189
Assignees

Comments

@neNasko1
Copy link
Contributor

To demonstrate the issue one may consult the following example:

import tracemalloc
import spox
from spox.opset.ai.onnx import v20 as op
import numpy as np
import gc

def app():
    a = op.const(np.zeros(10000))
    b = op.const(np.zeros(10000))
    for i in range(N):
        a = op.mul(a, a)
        gc.collect()
    return a

tracemalloc.start()

var = app()

print(tracemalloc.get_traced_memory())

tracemalloc.stop()

Upon execution, the outputs are as follows:

(spox) ➜  spox git:(main) ✗ python3 memtest10.py
(4494527, 4900843)
(spox) ➜  spox git:(main) ✗ python3 memtest100.py
(11821814, 12228285)

This example highlights that the Var-s persist with their values even when they are considered "unreachable" through the stable APIs. This issue stems from the current practice of consolidating all Var information in a singular linked location. Since most information is crucial during the subsequent building process, Node-s maintain references to their Input and Output entities, thus creating the present challenge.

Me and @cbourjau have been talking about ways to circumvent the issue at hand. In order of how hacky they are:

  • Have the ownership mechanism be represented by some weakref-s. This would be extremely hacky to pull off, since we only need to garbage collect the computed values when all the paths are from the form User->Var->Node->...->Var->_value
  • Have the user visible Var-s be a different object than the Var inside the creation graph. I am still a bit unsure whether this is something feasible as spox seems to rely on the fact that copying does not actually copy, however this invariant may be only needed for the building stage.
  • Create a InnerVar which holds all information except the _value which is held onto by every object in the graph and the current Var becomes a wrapper around that attaching a _value attribute. I am not extremely sure if this would yield a breaking change somewhere downstream.

This is the current stage of affairs, can you weight in? @jbachurski @adityagoel4512

@cbourjau
Copy link
Collaborator

Thanks for the summary, @neNasko1! Could you please elaborate a bit more on the difference between bullet points two and three?

@neNasko1
Copy link
Contributor Author

Thanks for the summary, @neNasko1! Could you please elaborate a bit more on the difference between bullet points two and three?

Sorry for not being clear enough, they should be mostly the same in functionality, but in 2) we have 2 Var-s instead of a Var and a InnerVar. I gave 2) to have a complete list of directions, but IMO is strictly worse than 3).

As this is pretty abstract at this point, I will try to have a draft PR implementing 3) this week to enable us to talk about things more concretely.

@neNasko1 neNasko1 self-assigned this Oct 23, 2024
@neNasko1 neNasko1 linked a pull request Nov 6, 2024 that will close this issue
8 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants