Skip to content

Manually call autodiff cleanup #4075

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.

Sign up for GitHub

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

Jump to bottom
Draft
swfsql wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Manually call autodiff cleanup #4075

swfsql wants to merge 2 commits into from
+24 −27

Conversation

@swfsql
Copy link
Contributor

@swfsql swfsql commented Nov 26, 2025 •
edited
Loading

Pull Request Template

Checklist

  • Confirmed that cargo run-checks command has been executed.
  • Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Changes

  • Adds the graph_cleanup method to the burn::tensor::backend::AutodiffBackend trait.
  • Also changed how the ordering of node removal is made, ensuring all leaf nodes are removed first.
    • This is not necessary since a given parent can only have a single child, but I think this could be a good little change.

Testing

Example of usage:

use burn::prelude::*;

type NdArray = burn::backend::NdArray<f32, i32>;
type NdArrayAuto = burn::backend::Autodiff<NdArray>;

fn main() {
    train::<NdArrayAuto>();
}

pub fn train<AutoB: burn::tensor::backend::AutodiffBackend>() {
    let train_device = <<AutoB as Backend>::Device>::default();
    for _ in 0..10_000_000 {
        let a: Tensor<AutoB, 2> = Tensor::zeros([2, 2], &train_device);
        let a = a.require_grad(); // an autodiff node is created
        drop(a); // the tensor is dropped but its unusable graph persists
        AutoB::graph_cleanup(); // either here, after the tensor is dropped, or after the loop
    }
    // in here, without graph_cleanup, AutoB/train_device cannot be used to cleanup the "hanging nodes"
}

Then by accompanying the system's ram usage, the memory usage doesn't grow (in case it's in the loop). Also, the are no vector/hashmaps/hashsets re-allocations, possibly improving runtime by a little.

I recently rebased the graph-sweep branch to main, as I'm using it locally.
In the current main, the "hot loop stress test" (increased to 10 million iterations) eventually still has the ram going up to around ~1GB, whereas graph-sweep leaves it limited at ~138kB. I haven't looked much into it why some nodes are missed on main, but I suspect it could be related to Arc counting (but I'm not sure).

Well, the reason why I created another PR is because all updates to the graph-sweep branch are ignored on the previous PR, because it is closed. And If the PR was already closed and if I rebased the branch (i.e. force-pushed the branch), I can no longer re-open it (-.-)', it appears to be a Github thing. So if it is ok, and contrary to what I did, I'd like to leave this PR open, even if it's open-ended.

@swfsql swfsql mentioned this pull request Nov 26, 2025
Closed
2 tasks
@swfsql swfsql marked this pull request as draft November 26, 2025 22:23
@swfsql
Copy link
Contributor Author

swfsql commented Nov 27, 2025 •
edited
Loading

To note, this PR probably has deadlocks on it, given the github action test runs.

@vadixidav
Copy link

vadixidav commented Dec 3, 2025

Let me know when it is stable again and I will battle test your branch. I am running training continuously and benefit a lot for my application from graph cleanup. I will say rc4 has been better with automated graph cleanup, but it isn't perfect. It helps a whole lot to call graph cleanup after I am done running all my backwards calls.

@swfsql
Copy link
Contributor Author

swfsql commented Dec 3, 2025 •
edited
Loading

@vadixidav I still haven't investigated the deadlocks, but I often rebase it to main (those are the force-push updates). But I'll ping you once I do -- but plz note that I'm not planning to investigate it too soon as I'm not getting the deadlocks for my own use.

And also to note, this PR as it is, you don't need to manually run it (you can if you want), but it automatically runs on every backward as well.

Edit: @laggui Also to note, if I moved the "extra cleanup procedure" from this PR to below the current cleanup procedure (not above as it currently is), then it fails (deadlocks aside) in cleaning up. So this is why I believe that this is related to Arc counting on the standard main cleanup. But then again, I haven't checked this too closely.

@github-actions
Copy link
Contributor

github-actions bot commented Jan 17, 2026

This PR has been marked as stale because it has not been updated for over a month

@github-actions github-actions bot added the stale The issue or pr has been open for too long label Jan 17, 2026
@swfsql
Copy link
Contributor Author

swfsql commented Jan 22, 2026

This linked PR (#4342) solves the need for the extra cleanup, and in current main I no longer see any leaks during the loop (considering the manual cleanup call).

So now this PR is mostly kept as a drafty way to directly call to the cleanup.

@swfsql swfsql changed the title Sweep autodiff graphs (2) Manually call autodiff cleanup Jan 22, 2026
@codecov
Copy link

codecov bot commented Jan 22, 2026

Codecov Report

❌ Patch coverage is 75.00000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.65%. Comparing base (5843c6e) to head (90dec2a).

Files with missing lines Patch % Lines
crates/burn-autodiff/src/backend.rs 0.00% 3 Missing ⚠️

❌ Your patch check has failed because the patch coverage (75.00%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage.
❌ Your project check has failed because the head coverage (68.65%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files 
@@            Coverage Diff             @@
##             main    #4075      +/-   ##
==========================================
- Coverage   68.65%   68.65%   -0.01%     
==========================================
  Files        1411     1411              
  Lines      168676   168671       -5     
==========================================
- Hits       115804   115796       -8     
- Misses      52872    52875       +3

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@github-actions github-actions bot removed the stale The issue or pr has been open for too long label Jan 22, 2026
@laggui
Copy link
Member

laggui commented Jan 23, 2026

This linked PR (#4342) solves the need for the extra cleanup, and in current main I no longer see any leaks during the loop (considering the manual cleanup call).

So now this PR is mostly kept as a drafty way to directly call to the cleanup.

Yeah that PR fixed a little oversight during the refactor that missed re-using the existing remove_entry cleanup method 😬 It was a good catch.

I appreciate you following up on the changes! Given this report, I think we'll eventually close this draft and you won't have to keep it up to date 😄

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@swfsql @vadixidav @laggui