Fix gradient checkpointing (#2997)

Author: nathanielsimard

Cargo.lock

@@ -157,9 +157,9 @@ portable-atomic = { version = "1.11.0" }
 portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "eaf4b4b4814c6f5f8ea6d07184e3b4d4dba3b3ae" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "eaf4b4b4814c6f5f8ea6d07184e3b4d4dba3b3ae" }
-cubecl-std = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "eaf4b4b4814c6f5f8ea6d07184e3b4d4dba3b3ae" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c93be864495d1016b4223eb014c64b01d79aa56a" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c93be864495d1016b4223eb014c64b01d79aa56a" }
+cubecl-std = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c93be864495d1016b4223eb014c64b01d79aa56a" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

Cargo.toml

@@ -1,6 +1,6 @@
 use crate::{
     collections::HashMap,
-    graph::{ComputingProperty, NodeID, NodeSteps},
+    graph::{ComputingProperty, NodeID},
     tensor::AutodiffTensor,
 };
 use alloc::{boxed::Box, sync::Arc, vec::Vec};
@@ -108,8 +108,7 @@ impl CheckpointerBuilder {
         }
     }
 
-    pub(crate) fn build(self, graph: &NodeSteps) -> Checkpointer {
-        let node_tree = self.make_tree(graph);
+    pub(crate) fn build(self, node_tree: NodeTree) -> Checkpointer {
         let mut backward_states_map = HashMap::new();
         let mut retro_forwards_map = HashMap::new();
 
@@ -247,14 +246,6 @@ impl CheckpointerBuilder {
         }
     }
 
-    fn make_tree(&self, graph: &NodeSteps) -> NodeTree {
-        let mut tree = HashMap::default();
-        for (id, step) in graph {
-            tree.insert(*id, step.parents());
-        }
-        NodeTree::new(tree)
-    }
-
     fn update_n_required_of_parents(
         id: NodeID,
         n_required_map: &mut HashMap<NodeID, usize>,

crates/burn-autodiff/src/checkpoint/builder.rs

@@ -16,12 +16,23 @@ pub trait CheckpointStrategy: Clone + Copy + Debug + Default + Send + Sync + 'st
     fn compute_property<R: RetroForward>(retro_forward: R) -> ComputingProperty;
 
     /// Checkpoints parents if necessary in the strategy
-    fn checkpoint_parents<'a, B2, A>(parents: A, builder: &mut CheckpointerBuilder)
+    fn checkpoint_parents<'a, B2, A>(
+        parents: A,
+        builder: &mut CheckpointerBuilder,
+    ) -> Result<(), CheckpointingError>
     where
         B2: Backend,
         A: IntoIterator<Item = &'a AutodiffTensor<B2>>;
 }
 
+#[derive(Debug)]
+/// Error that can happen when trying to checkpoint a tensor.
+pub enum CheckpointingError {
+    /// When a parent is untracked, we can't easily checkpoint its state, since we don't know the
+    /// requirements in advanced.
+    UntrackedParent,
+}
+
 #[derive(Clone, Copy, Debug, Default)]
 /// All operations are considered compute bound, notwithstanding how they are marked
 pub struct NoCheckpointing {}
@@ -34,12 +45,16 @@ impl CheckpointStrategy for NoCheckpointing {
 
     /// An operation marked as memory bound is actually compute bound.
     /// It's therefore useless to checkpoint the parents
-    fn checkpoint_parents<'a, B2, A>(_parents: A, _builder: &mut CheckpointerBuilder)
+    fn checkpoint_parents<'a, B2, A>(
+        _parents: A,
+        _builder: &mut CheckpointerBuilder,
+    ) -> Result<(), CheckpointingError>
     where
         B2: Backend,
         A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
     {
         // Nothing to do here
+        Ok(())
     }
 }
 
@@ -59,13 +74,29 @@ impl CheckpointStrategy for BalancedCheckpointing {
     /// An operation marked as memory bound is really memory bound.
     /// Since the operation may not checkpoint its parents but may need them indirectly
     /// if asked to recompute itself, the method needs to know the parent tensors to maybe checkpoint them
-    fn checkpoint_parents<'a, B2, A>(parents: A, builder: &mut CheckpointerBuilder)
+    fn checkpoint_parents<'a, B2, A>(
+        parents: A,
+        builder: &mut CheckpointerBuilder,
+    ) -> Result<(), CheckpointingError>
     where
         B2: Backend,
         A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
     {
+        let mut can_checkpoint = true;
+
         for tensor in parents.into_iter() {
-            builder.checkpoint(tensor, ActionType::Backup);
+            if let crate::graph::Requirement::None = tensor.node.requirement {
+                can_checkpoint = false;
+            } else {
+                builder.checkpoint(tensor, ActionType::Backup);
+            }
         }
+
+        if !can_checkpoint {
+            *builder = CheckpointerBuilder::default();
+            return Err(CheckpointingError::UntrackedParent);
+        }
+
+        Ok(())
     }
 }

crates/burn-autodiff/src/checkpoint/strategy.rs

@@ -1,5 +1,5 @@
 use super::NodeID;
-use crate::{checkpoint::base::Checkpointer, collections::HashMap, grads::Gradients};
+use crate::{checkpoint::base::Checkpointer, grads::Gradients};
 use alloc::{boxed::Box, vec::Vec};
 
 /// Backward step for reverse mode autodiff.
@@ -15,4 +15,3 @@ pub trait Step: Send + core::fmt::Debug {
 }
 
 pub type StepBoxed = Box<dyn Step>;
-pub type NodeSteps = HashMap<NodeID, StepBoxed>;

crates/burn-autodiff/src/graph/base.rs

@@ -1,7 +1,7 @@
 use super::NodeRef;
 
 /// Requirement for each tensor in the graph.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Requirement {
     /// Operations that require gradients.
     Grad,

crates/burn-autodiff/src/graph/requirement.rs

@@ -106,13 +106,17 @@ where
         B2: Backend,
         A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
     {
-        C::checkpoint_parents(parents, &mut self.checkpointer_builder);
+        let compute_property = match C::checkpoint_parents(parents, &mut self.checkpointer_builder)
+        {
+            Ok(..) => self.compute_property,
+            Err(..) => ComputingProperty::ComputeBound,
+        };
 
         OpsPrep::new(
             self.nodes,
             self.requirement,
             self.backward,
-            self.compute_property,
+            compute_property,
             self.checkpointer_builder,
         )
     }

crates/burn-autodiff/src/ops/base.rs

@@ -1,7 +1,10 @@
 use super::memory_management::GraphMemoryManagement;
 use crate::{
     NodeID,
-    checkpoint::{base::Checkpointer, builder::CheckpointerBuilder},
+    checkpoint::{
+        base::{Checkpointer, NodeTree},
+        builder::CheckpointerBuilder,
+    },
     collections::HashMap,
     grads::Gradients,
     graph::{StepBoxed, traversal::BreadthFirstSearch},
@@ -34,8 +37,7 @@ impl AutodiffServer {
         );
         let builder = self.actions_builder.remove(&node_id).unwrap();
 
-        let (tape, builder) = self.build_tape(node_id, step, builder);
-        let checkpointer = builder.build(&self.steps);
+        let (tape, checkpointer) = self.build_tape(node_id, step, builder);
 
         let gradients = Self::execute_steps(tape, grads, checkpointer);
 
@@ -54,20 +56,25 @@ impl AutodiffServer {
         node: NodeID,
         node_step: StepBoxed,
         mut builder: CheckpointerBuilder,
-    ) -> (Vec<Vec<StepBoxed>>, CheckpointerBuilder) {
+    ) -> (Vec<Vec<StepBoxed>>, Checkpointer) {
         let mut tape = (0..node_step.depth())
             .map(|_| Vec::with_capacity(1))
             .collect::<Vec<_>>();
 
+        let mut tree = HashMap::default();
+
         BreadthFirstSearch.traverse(node, node_step, &mut self.steps, |id, step| {
             self.memory_management.consume_node(id);
 
             let depth = step.depth();
+
             if depth == 0 {
                 return;
             }
 
             if let Some(steps) = tape.get_mut(depth - 1) {
+                let parents = step.parents().into_iter().filter(|s| *s != id);
+                tree.insert(id, parents.collect());
                 steps.push(step);
             }
 
@@ -76,7 +83,9 @@ impl AutodiffServer {
             }
         });
 
-        (tape, builder)
+        let checkpointer = builder.build(NodeTree::new(tree));
+
+        (tape, checkpointer)
     }
 
     fn execute_steps(
@@ -93,6 +102,7 @@ impl AutodiffServer {
         #[cfg(feature = "export_tests")]
         // For checkpointing tests
         assert!(checkpointer.is_empty());
+
         grads
     }
 }