WIP Debug

tracel-ai · nathanielsimard · Apr 26, 2024 · Apr 22, 2024 · Apr 22, 2024 · Apr 25, 2024
commit 73a330a6d4070afbe16577830f906f2453bde1df
diff --git a/backend-comparison/Cargo.toml b/backend-comparison/Cargo.toml
@@ -88,6 88,10 @@ name = "custom-gelu"
 path = "benches/custom_gelu.rs"
 harness = false
 
 [[bench]]
 name = "autodiff"
 harness = false
 
 [[bin]]
 name = "burnbench"
 path = "src/bin/burnbench.rs"
diff --git a/backend-comparison/benches/autodiff.rs b/backend-comparison/benches/autodiff.rs
@@ -0,0 1,81 @@
 use backend_comparison::persistence::save;
 use burn::{
  module::Module,
  nn,
  tensor::{
  backend::{AutodiffBackend, Backend},
  Distribution, Tensor,
  },
 };
 use burn_common::benchmark::{run_benchmark, Benchmark};
 
 pub struct AutodiffOverheadBenchmark<B: AutodiffBackend> {
  config: nn::LstmConfig,
  lstm: nn::Lstm<B>,
  device: B::Device,
 }
 
 impl<B: AutodiffBackend> Benchmark for AutodiffOverheadBenchmark<B> {
  type Args = Tensor<B, 3>;
 
  fn name(&self) -> String {
  "autodiff_overhead".into()
  }
 
  fn shapes(&self) -> Vec<Vec<usize>> {
  vec![]
  }
 
  fn execute(&self, input: Self::Args) {
  for _ in 0..20 {
  let input = input.clone().detach();
  let mut cell = input.clone();
  let lstm = self.lstm.clone().fork(&input.device());
 
  for i in 0..10 {
  let (cells, _) = lstm.forward(input.clone(), None);
  cell = cell cells;
  }
 
  cell.backward();
  }
  }
 
  fn prepare(&self) -> Self::Args {
  let shape = [1, 3, self.config.d_hidden];
  Tensor::random(shape, Distribution::Default, &self.device)
  }
 
  fn sync(&self) {
  B::sync(&self.device)
  }
 }
 
 #[allow(dead_code)]
 fn bench<B: Backend>(
  device: &B::Device,
  feature_name: &str,
  url: Option<&str>,
  token: Option<&str>,
 ) {
  let config = nn::LstmConfig::new(3, 3, true);
  let lstm = config.init(device);
  let benchmark = AutodiffOverheadBenchmark::<burn::backend::Autodiff<B>> {
  lstm,
  config,
  device: device.clone(),
  };
 
  save::<B>(
  vec![run_benchmark(benchmark)],
  device,
  feature_name,
  url,
  token,
  )
  .unwrap();
 }
 
 fn main() {
  backend_comparison::bench_on_backend!();
 }
diff --git a/backend-comparison/src/burnbenchapp/base.rs b/backend-comparison/src/burnbenchapp/base.rs
@@ -98,6 98,8 @@ enum BenchmarkValues {
  MaxPool2d,
  #[strum(to_string = "load-record")]
  LoadRecord,
  #[strum(to_string = "autodiff")]
  Autodiff,
 }
 
 pub fn execute() {

diff --git a/crates/burn-autodiff/src/runtime/client.rs b/crates/burn-autodiff/src/runtime/client.rs
@@ -15,9 15,9 @@ pub trait AutodiffClient: Send Clone {
 }
 
 /// Client implementation in used.
-#[cfg(feature = "std")]
-pub type AutodiffClientImpl = super::mspc::ChannelClient;
-
-/// Client implementation in used.
-#[cfg(not(feature = "std"))]
 // #[cfg(feature = "std")]
 // pub type AutodiffClientImpl = super::mspc::ChannelClient;
 //
 // /// Client implementation in used.
 // #[cfg(not(feature = "std"))]
 pub type AutodiffClientImpl = super::mutex::MutexClient;
diff --git a/crates/burn-autodiff/src/runtime/memory_management.rs b/crates/burn-autodiff/src/runtime/memory_management.rs
@@ -11,6 11,8 @@ use std::{
 pub struct GraphMemoryManagement {
  graphs: HashMap<GraphId, GraphState>,
  owned: HashSet<GraphId>,
  nodes: HashMap<NodeRefCount, Vec<NodeID>>,
  leafs: HashSet<NodeID>,
 }
 
 #[derive(new, Hash, PartialEq, Eq, Clone, Copy, Debug)]
@@ -78,7 80,7 @@ impl GraphMemoryManagement {
  pub fn find_orphan_graphs(&self) -> Vec<GraphId> {
  self.owned
  .iter()
- .filter(|id| self.is_orphan(id))
  // .filter(|id| self.is_orphan(id))
  .copied()
  .collect()
  }
@@ -174,19 176,19 @@ impl core::fmt::Display for GraphMemoryManagement {
  self.owned.len(),
  self.graphs.len()
  ))?;
- for (id, state) in self.graphs.iter() {
- f.write_fmt(format_args!("Graph {} => ", id.node.value))?;
- match state {
- GraphState::Merged(id) => f.write_fmt(format_args!("Merged {}", id.node.value))?,
- GraphState::Owned(nodes) => {
- f.write_str("Owned")?;
- for node in nodes {
- f.write_fmt(format_args!(" {}", node.value))?;
- }
- }
- }
- f.write_str("\n")?;
- }
  // for (id, state) in self.graphs.iter() {
  //  f.write_fmt(format_args!("Graph {} => ", id.node.value))?;
  //  match state {
  //  GraphState::Merged(id) => f.write_fmt(format_args!("Merged {}", id.node.value))?,
  //  GraphState::Owned(nodes) => {
  //  f.write_str("Owned")?;
  //  for node in nodes {
  //  f.write_fmt(format_args!(" {}", node.value))?;
  //  }
  //  }
  //  }
  //  f.write_str("\n")?;
  // }
 
  Ok(())
  }

diff --git a/crates/burn-autodiff/src/runtime/mod.rs b/crates/burn-autodiff/src/runtime/mod.rs
@@ -2,10 2,10 @@ mod client;
 mod memory_management;
 mod server;
 
-#[cfg(not(feature = "std"))]
 // #[cfg(not(feature = "std"))]
 pub mod mutex;
 
-#[cfg(feature = "std")]
-pub mod mspc;
 // #[cfg(feature = "std")]
 // pub mod mspc;
 
 pub use client::*;
diff --git a/crates/burn-autodiff/src/runtime/server.rs b/crates/burn-autodiff/src/runtime/server.rs
@@ -8,6 8,8 @@ use crate::{
 };
 use std::collections::HashMap;
 
 static GRAPH_MM_ACTIVATED: bool = true;
 
 #[derive(Default)]
 pub struct AutodiffServer {
  steps: HashMap<NodeID, StepBoxed>,
@@ -20,7 22,9 @@ impl AutodiffServer {
  let parents = step.parents();
  let node_id = *rc.as_ref();
 
- self.memory_management.register(rc, parents);
  if GRAPH_MM_ACTIVATED {
  self.memory_management.register(rc, parents);
  }
 
  self.steps.insert(node_id, step);
  self.actions_builder.insert(node_id, actions);
@@ -33,20 37,30 @@ impl AutodiffServer {
  );
  let builder = self.actions_builder.remove(&node_id).unwrap();
 
 
  let orphans = self.memory_management.find_orphan_graphs();
 
  let (tape, builder) = self.build_tape(node_id, step, builder);
  let checkpointer = builder.build(&self.steps);
 
  let gradients = Self::execute_steps(tape, grads, checkpointer);
 
- // Cleanup
- let mut on_free_graph = |node_id: &NodeID| {
- self.steps.remove(node_id);
- self.actions_builder.remove(node_id);
- };
  if !GRAPH_MM_ACTIVATED {
  self.steps.clear();
  self.actions_builder.clear();
  } else {
  // Cleanup
  let mut on_free_graph = |node_id: &NodeID| {
  self.steps.remove(node_id);
  self.actions_builder.remove(node_id);
  };
 
  for graph_id in orphans {
  self.memory_management
  .free_graph(graph_id, &mut on_free_graph);
  }
 
- for graph_id in self.memory_management.find_orphan_graphs() {
- self.memory_management
- .free_graph(graph_id, &mut on_free_graph);
  println!("{}", self.memory_management);
  }
 
  gradients
@@ -63,7 77,9 @@ impl AutodiffServer {
  .collect::<Vec<_>>();
 
  BreadthFirstSearch.traverse(node, node_step, &mut self.steps, |id, step| {
- self.memory_management.free_node(id);
  if GRAPH_MM_ACTIVATED {
  self.memory_management.free_node(id);
  }
 
  let depth = step.depth();
  if depth == 0 {