3% Memory Regression when Compiling the Optimizer

[eellison]

🐛 Describe the bug

Compiling the optimizer introduced a ~3% memory regression as observed in the performance dashboard. @mlazos, @ngimel and others have tried debugging a bit but we haven't yet found the culprit.

3% is not huge but it's still worth getting to the bottom of this, and memory overheads add up.

Error logs

No response

Minified repro

No response

[mlazos]

The CheckFunctionManager contains a circular reference to itself, which keeps the output_graph alive. Within the output_graph the side_effects tracker would keep the grads from all of the parameters of the model in its keepalive list (to track mutation of the list containing all of the grads) which is necessary during analysis, but should be released once we've generated code. This was only noticeable after compiling the optimizer because the grads for the entire model were kept in a list, but this may help memory across a lot of models if tensors were stored in containers which could be mutated. I'm removing the full output graph from the CheckFunctionManager, and this should allow side_effects to get cleaned up after compilation.

[mlazos]

benchmark to repro https://gist.github.com/mlazos/7969081049d4e0681774defc0d0f50c7