⚡ Bolt: [performance improvement] Option::take() for Tensor ownership in backward pass#110
⚡ Bolt: [performance improvement] Option::take() for Tensor ownership in backward pass#110teerthsharma wants to merge 1 commit into
Conversation
… in backward pass What: Use `Option::take()` in `autograd::Context::backward` to extract `Tensor` ownership from `grads` and pass it to `accumulate_grad`. Why: Bypasses borrow checker conflicts that forced expensive `.clone()` operations on `Option<Tensor>` and its heap-allocated metadata (shape and strides) arrays inside the hot reverse-mode propagation loop. Impact: Eliminates O(N) heap allocations for every operation on the gradient tape (`Add`, `Mul`, `MatMul`, `ReLU`) when propagating derivatives backward. Measurement: Compare backpropagation time in deep network training loops before and after this optimization. Tests passed successfully with no functional changes. Co-authored-by: teerthsharma <[email protected]>
|
👋 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: Modified
accumulate_gradto take ownership (grad: Tensor) instead of a reference. In thebackwardmethod,Option::take()is now used to pull the gradient out of thegradsvector so that it can be passed by value, and then re-inserted afterwards. For operations likeOp::Mul,Op::MatMul, andOp::ReLUwhich dynamically create partial derivative tensors, this completely avoids cloning internalTensormetadata.🎯 Why: During reverse-mode automatic differentiation, the backward pass operates in a tight loop. Cloning a
Tensorto satisfy borrow checker constraints triggers new vector allocations for itsshapeandstridesfields. Removing this drastically reduces GC pressure and speeds up training.📊 Impact: Reduces memory overhead and allocation time inside the backward pass by avoiding metadata cloning for
Op::Add,Op::Mul,Op::MatMul, andOp::ReLU.🔬 Measurement: Verifiable by checking that
cargo test -p aether-corecontinues to pass correctly. Performance gains can be benchmarked on larger MLP models liketest_mlp_large_scale.PR created automatically by Jules for task 5931644591477573448 started by @teerthsharma