[Breaking] Make Tensor, Module, Optimizer !Sync + Refactor Autodiff (#1575)

Author: nathanielsimard

crates/burn-autodiff/Cargo.toml

@@ -11,8 +11,9 @@ repository = "https://github.com/tracel-ai/burn/tree/main/burn-autodiff"
 version.workspace = true
 
 [features]
-default = []
+default = ["std"]
 export_tests = ["burn-tensor-testgen"]
+std = []
 
 [dependencies]
 burn-common = { path = "../burn-common", version = "0.13.0" }

crates/burn-autodiff/src/backend.rs

@@ -1,7 +1,7 @@
 use crate::{
     checkpoint::strategy::{CheckpointStrategy, NoCheckpointing},
     grads::Gradients,
-    graph::backward::backward,
+    runtime::AutodiffClient,
     tensor::AutodiffTensor,
     AutodiffBridge,
 };
@@ -53,7 +53,9 @@ impl<B: Backend, C: CheckpointStrategy> AutodiffBackend for Autodiff<B, C> {
     type Gradients = Gradients;
 
     fn backward<const D: usize>(tensor: AutodiffTensor<B, D>) -> Gradients {
-        backward(tensor)
+        let client = tensor.node.client.clone();
+
+        AutodiffClient::backward(&client, tensor)
     }
 
     fn grad<const D: usize>(
@@ -83,7 +85,7 @@ impl<B: Backend, C: CheckpointStrategy> AutodiffBackend for Autodiff<B, C> {
         grad: B::FloatTensorPrimitive<D>,
     ) {
         grads.remove(tensor);
-        grads.register::<B, D>(tensor.node.clone(), grad);
+        grads.register::<B, D>(tensor.node.id, grad);
     }
 
     fn int_inner<const D: usize>(

crates/burn-autodiff/src/bridge.rs

@@ -39,7 +39,7 @@ where
             _bridge: PhantomData<Bridge>,
         }
 
-        #[derive(new, Debug)]
+        #[derive(new, Debug, Clone)]
         struct RetroIntoTarget<B: Backend, Bridge: BackendBridge<B>, const D: usize> {
             tensor_id: NodeID,
             _backend: PhantomData<B>,
@@ -84,9 +84,9 @@ where
             _backend: PhantomData,
             _bridge: PhantomData,
         }
-        .prepare::<C>([tensor.node.clone()], [tensor.graph.clone()])
+        .prepare::<C>([tensor.node.clone()])
         .memory_bound()
-        .retro_forward(RetroIntoTarget::<B, Bridge, D>::new(tensor.node.id.clone()))
+        .retro_forward(RetroIntoTarget::<B, Bridge, D>::new(tensor.node.id))
         .parents([&tensor])
         .stateless(Bridge::into_target(tensor.primitive, None))
     }
@@ -101,7 +101,7 @@ where
             _bridge: PhantomData<Bridge>,
         }
 
-        #[derive(new, Debug)]
+        #[derive(new, Debug, Clone)]
         struct RetroFromTarget<B: Backend, Bridge: BackendBridge<B>, const D: usize> {
             tensor_id: NodeID,
             _backend: PhantomData<B>,
@@ -146,9 +146,9 @@ where
             _backend: PhantomData,
             _bridge: PhantomData,
         }
-        .prepare::<C>([tensor.node.clone()], [tensor.graph.clone()])
+        .prepare::<C>([tensor.node.clone()])
         .memory_bound()
-        .retro_forward(RetroFromTarget::<B, Bridge, D>::new(tensor.node.id.clone()))
+        .retro_forward(RetroFromTarget::<B, Bridge, D>::new(tensor.node.id))
         .parents([&tensor])
         .stateless(Bridge::from_target(tensor.primitive, None))
     }

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

@@ -1,22 +1,20 @@
-use std::collections::HashMap;
-
-use crate::graph::{NodeID, NodeRef};
-
 use super::{
     retro_forward::RetroForwards,
     state::{BackwardStates, State},
 };
+use crate::graph::NodeID;
+use std::collections::HashMap;
 
 #[derive(new, Debug)]
 /// Links a [NodeID] to its autodiff graph [NodeRef]
 pub(crate) struct NodeTree {
-    map: HashMap<NodeID, NodeRef>,
+    map: HashMap<NodeID, Vec<NodeID>>,
 }
 
 impl NodeTree {
     /// Gives the parents of the node in the autodiff graph
     pub(crate) fn parents(&self, node_id: &NodeID) -> Option<Vec<NodeID>> {
-        self.map.get(node_id).map(|node| node.parents.clone())
+        self.map.get(node_id).cloned()
     }
 }
 
@@ -33,14 +31,12 @@ impl Checkpointer {
     /// or give their pre-computed tensors.
     pub fn retrieve_node_output<T>(&mut self, node_id: NodeID) -> T
     where
-        T: Clone + Send + Sync + 'static,
+        T: Clone + Send + 'static,
     {
-        self.topological_sort(node_id.clone())
-            .into_iter()
-            .for_each(|node| {
-                self.retro_forwards
-                    .execute_retro_forward(node, &mut self.backward_states)
-            });
+        self.topological_sort(node_id).into_iter().for_each(|node| {
+            self.retro_forwards
+                .execute_retro_forward(node, &mut self.backward_states)
+        });
 
         self.backward_states.get_state::<T>(&node_id)
     }

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

@@ -1,11 +1,9 @@
-use std::{any::Any, collections::HashMap, sync::Arc};
-
-use burn_tensor::backend::Backend;
-
 use crate::{
-    graph::{ComputingProperty, NodeID, NodeRef, NodeSteps},
+    graph::{ComputingProperty, NodeID, NodeSteps},
     tensor::AutodiffTensor,
 };
+use burn_tensor::backend::Backend;
+use std::{any::Any, collections::HashMap, sync::Arc};
 
 use super::{
     base::{Checkpointer, NodeTree},
@@ -20,31 +18,34 @@ pub enum CheckpointingAction {
     /// The node's already computed output should be saved
     Computed {
         /// The node
-        node_ref: NodeRef,
+        node_id: NodeID,
         /// The node's output
-        state_content: Box<dyn Any + Send + Sync>,
+        state_content: Box<dyn Any + Send>,
     },
     /// The node should recompute itself when asked
     Recompute {
         /// The node
-        node_ref: NodeRef,
+        node_id: NodeID,
         /// How the node should recompute itself
         retro_forward: Arc<dyn RetroForward>,
     },
 }
 
+// TODO: Remove that when proper client server.
+unsafe impl Send for CheckpointingAction {}
+
 impl CheckpointingAction {
     /// Utilitary function to access the id of the node of the checkpointing action
     pub fn id(&self) -> NodeID {
         match self {
             CheckpointingAction::Computed {
-                node_ref,
+                node_id: node_ref,
                 state_content: _,
-            } => node_ref.id.clone(),
+            } => *node_ref,
             CheckpointingAction::Recompute {
-                node_ref,
+                node_id: node_ref,
                 retro_forward: _,
-            } => node_ref.id.clone(),
+            } => *node_ref,
         }
     }
 }
@@ -83,13 +84,13 @@ impl CheckpointerBuilder {
         match &tensor.node.properties {
             ComputingProperty::ComputeBound | ComputingProperty::Ambiguous => {
                 action_list.push(CheckpointingAction::Computed {
-                    node_ref: tensor.node.clone(),
+                    node_id: tensor.node.id,
                     state_content: Box::new(tensor.primitive.clone()),
                 })
             }
             ComputingProperty::MemoryBound { retro_forward } => {
                 action_list.push(CheckpointingAction::Recompute {
-                    node_ref: tensor.node.clone(),
+                    node_id: tensor.node.id,
                     retro_forward: retro_forward.clone(),
                 })
             }
@@ -105,10 +106,6 @@ impl CheckpointerBuilder {
         }
     }
 
-    pub(crate) fn len(&self) -> usize {
-        self.explicit_actions.len() + self.backup_actions.len()
-    }
-
     pub(crate) fn build(self, graph: &NodeSteps) -> Checkpointer {
         let node_tree = self.make_tree(graph);
         let mut backward_states_map = HashMap::new();
@@ -143,11 +140,11 @@ impl CheckpointerBuilder {
         {
             match action {
                 CheckpointingAction::Computed {
-                    node_ref,
+                    node_id: node_ref,
                     state_content: _,
-                } => stop_nodes.push(node_ref.id.clone()),
+                } => stop_nodes.push(*node_ref),
                 CheckpointingAction::Recompute {
-                    node_ref: _,
+                    node_id: _,
                     retro_forward: _,
                 } => {}
             }
@@ -165,10 +162,10 @@ impl CheckpointerBuilder {
         for action in self.explicit_actions.iter() {
             match action {
                 CheckpointingAction::Computed {
-                    node_ref,
+                    node_id: node_ref,
                     state_content: _,
                 } => {
-                    let id = node_ref.id.clone();
+                    let id = *node_ref;
                     match n_required_map.remove(&id) {
                         Some(n) => {
                             n_required_map.insert(id, n + 1);
@@ -179,10 +176,10 @@ impl CheckpointerBuilder {
                     };
                 }
                 CheckpointingAction::Recompute {
-                    node_ref,
+                    node_id: node_ref,
                     retro_forward: _,
                 } => {
-                    let id = node_ref.id.clone();
+                    let id = *node_ref;
                     Self::update_n_required_of_parents(
                         id,
                         &mut n_required_map,
@@ -229,13 +226,13 @@ impl CheckpointerBuilder {
 
             match action {
                 CheckpointingAction::Computed {
-                    node_ref: _,
+                    node_id: _,
                     state_content,
                 } => {
                     self.checkpoint_compute(backward_states_map, node_id, state_content, n_required)
                 }
                 CheckpointingAction::Recompute {
-                    node_ref: _,
+                    node_id: _,
                     retro_forward,
                 } => self.checkpoint_lazy(
                     backward_states_map,
@@ -251,7 +248,7 @@ impl CheckpointerBuilder {
     fn make_tree(&self, graph: &NodeSteps) -> NodeTree {
         let mut tree = HashMap::default();
         for (id, step) in graph {
-            tree.insert(id.clone(), step.node());
+            tree.insert(*id, step.parents());
         }
         NodeTree::new(tree)
     }
@@ -267,7 +264,7 @@ impl CheckpointerBuilder {
                 n_required_map.insert(id, n + 1);
             }
             None => {
-                n_required_map.insert(id.clone(), 1);
+                n_required_map.insert(id, 1);
                 if !stop_nodes.contains(&id) {
                     if let Some(parents) = node_tree.parents(&id) {
                         for p in parents {
@@ -288,7 +285,7 @@ impl CheckpointerBuilder {
         &self,
         backward_states_map: &mut HashMap<NodeID, State>,
         node_id: NodeID,
-        state_content: Box<dyn Any + Send + Sync>,
+        state_content: Box<dyn Any + Send>,
         n_required: usize,
     ) {
         backward_states_map.insert(
@@ -308,7 +305,7 @@ impl CheckpointerBuilder {
         retro_forward: Arc<dyn RetroForward>,
         n_required: usize,
     ) {
-        retro_forward_map.insert(node_id.clone(), retro_forward);
-        backward_states_map.insert(node_id.clone(), State::Recompute { n_required });
+        retro_forward_map.insert(node_id, retro_forward);
+        backward_states_map.insert(node_id, State::Recompute { n_required });
     }
 }

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

@@ -7,7 +7,7 @@ use super::state::{BackwardStates, State};
 /// Definition of the forward function of a node, called during retropropagation only.
 /// This is different from the normal forward function because it reads and writes from
 /// the [InnerStates] map instead of having a clear function signature.
-pub trait RetroForward: Debug + Send + Sync + 'static {
+pub trait RetroForward: Debug + Send + 'static {
     fn forward(&self, states: &mut BackwardStates, out_node: NodeID);
 }
 
@@ -31,7 +31,7 @@ impl RetroForwards {
         {
             // Retro forwards are always used only once because afterwards their state is computed
             let retro_forward = self.map.remove(&node_id).unwrap();
-            retro_forward.forward(backward_states, node_id.clone());
+            retro_forward.forward(backward_states, node_id);
         }
     }
 
@@ -48,7 +48,7 @@ macro_rules! retro_unary_scalar {
         $name:ident,
         $ops:expr
     ) => {
-        #[derive(new, Debug)]
+        #[derive(new, Debug, Clone)]
         struct $name<B: Backend, const D: usize> {
             lhs_id: NodeID,
             rhs: FloatElem<B>,
@@ -72,7 +72,7 @@ macro_rules! retro_unary {
         $name:ident,
         $ops:expr
     ) => {
-        #[derive(new, Debug)]
+        #[derive(new, Debug, Clone)]
         struct $name<B: Backend, const D: usize> {
             input_id: NodeID,
             _backend: PhantomData<B>,
@@ -95,7 +95,7 @@ macro_rules! retro_binary {
         $name:ident,
         $ops:expr
     ) => {
-        #[derive(new, Debug)]
+        #[derive(new, Debug, Clone)]
         struct $name<B: Backend, const D: usize> {
             lhs_id: NodeID,
             rhs_id: NodeID,

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

@@ -3,7 +3,7 @@ use std::{any::Any, collections::HashMap};
 use crate::graph::NodeID;
 
 /// In order to accept arbitrary node output in the same hashmap, we need to upcast them to any.
-pub(crate) type StateContent = Box<dyn Any + Send + Sync>;
+pub(crate) type StateContent = Box<dyn Any + Send>;
 
 #[derive(Debug)]
 /// The state contained at one node. Encapsulates the node output if precomputed,
@@ -71,7 +71,7 @@ impl BackwardStates {
     /// This function always gives ownership of the output, but will clone it if needed for further uses.
     pub(crate) fn get_state<T>(&mut self, node_id: &NodeID) -> T
     where
-        T: Clone + Send + Sync + 'static,
+        T: Clone + Send + 'static,
     {
         // Fetch the state and decrement its number of required
         let state = self.map.remove(node_id).unwrap();
@@ -97,7 +97,7 @@ impl BackwardStates {
                 .unwrap()
                 .clone();
 
-            self.insert_state(node_id.clone(), new_stored_state);
+            self.insert_state(*node_id, new_stored_state);
 
             downcasted
         } else {
@@ -119,7 +119,7 @@ impl BackwardStates {
 
     pub(crate) fn save<T>(&mut self, node_id: NodeID, saved_output: T)
     where
-        T: Clone + Send + Sync + 'static,
+        T: Clone + Send + 'static,
     {
         let n_required = self.get_state_ref(&node_id).unwrap().n_required();
         self.insert_state(