Skip to content

Fixing leaking autograd graph memory in Trainer - #292

Merged
aarontuor merged 3 commits into
pnnl:developfrom
EladMichael:bugfix/devloss_autograd
Sep 14, 2026
Merged

aarontuor merged 3 commits into
pnnl:developfrom
EladMichael:bugfix/devloss_autograd

Conversation

@EladMichael

Copy link
Copy Markdown
Collaborator

Trainer has two memory leaks, both from the same source, we're not detaching the grad graphs from the metrics before we store them. This fixes both of those, as well as including a benchmark so you can see the difference. Check out the second commit, run the benchmark, then check out the third commit and run the benchmark.

No longer storing hundreds of MB or a few GB of autograd graph on accident per long training run!

… trainer process. Two losses are stored without .detach() being called, the self.best_devloss = output[self.eval_metric] is the big one, and self.loss_history[train].append(mean_loss) is the smaller
@aarontuor
aarontuor merged commit 7e2a233 into pnnl:develop Sep 14, 2026
1 check passed
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 this pull request may close these issues.

2 participants