⚡ Bolt: Optimize autograd reverse pass to avoid redundant Tensor clones#127
⚡ Bolt: Optimize autograd reverse pass to avoid redundant Tensor clones#127teerthsharma wants to merge 1 commit into
Conversation
Changed the `backward` pass to use `Option::take()` instead of `.clone()` when retrieving gradients. Updated `accumulate_grad` to take its `Tensor` argument by value to eliminate cloning during the `.add()` and insertion operations. Added a critical performance learning entry to `.jules/bolt.md`. Co-authored-by: teerthsharma <78080953+teerthsharma@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
1 participant
💡 What: The
backwardpass of the autograd engine now usesOption::take()when fetching the gradient for the output node of an operation. Furthermore, theaccumulate_gradhelper was modified to take itsTensorargument by value.🎯 Why: In the previous implementation, fetching gradients from the
gradsvector via.clone()caused an unnecessary heap allocation for theTensormetadata (shape and strides), inside a very hot loop. By moving values with.take(), taking ownership inaccumulate_grad, and conditionally replacing them back, we eliminate redundant allocations during backpropagation.📊 Impact: Reduces memory allocations during backpropagation. Specifically, eliminates one
Tensormetadata heap allocation per operation block read, and one when inserting into a new gradient slot.🔬 Measurement: Run
cargo test -p aether-core --offlineto verify correctness. Theml::autograd::testssuite confirms no mathematical regressions.PR created automatically by Jules for task 10035790909556048901 started by @teerthsharma