A concurrent `backward()` frees another thread's autodiff steps, silently losing a gradient

[swfsql]

(Analysis and reproduction prepared by Claude)

Describe the bug

When more than one thread calls backward() in the same process, a leaf tensor can come
back absent from the Gradients map — tensor.grad(&grads) returns None for a
parameter that plainly participated in the loss. No panic, no warning. When only part of a
branch is lost the gradient is present but wrong instead.

GraphCleaner::cleanup_orphaned_entries() (burn-autodiff/src/runtime/graph.rs) runs at
the end of every backward() and sweeps every graph in the process, including
graphs another thread is still building. For a graph it judges un-useful,
free_unused_roots drops the steps. The other thread's backward then cannot traverse past
the missing step.

The judgement is GraphMemoryManagement::maybe_useful(), i.e. "does any node in this graph
have Arc::strong_count > 1". That is not a sound liveness test: a node can be live yet
unreferenced, because the only thing naming it is a child that has been created — holding
its id in Node::parents — and not yet registered. Single-threaded that window is never
observed, since the sweep only runs at the end of a backward on the same thread. Across
threads it is as wide as a tensor copy.

Reproduces on 1893b60 (0.22.0-pre.2), 13f0a12 (0.22.0-pre.3) and 1bceb88 (main), on
both the flex and ndarray backends.

To Reproduce

Two files, no local paths; cargo test -- --test-threads=1 --nocapture.

Cargo.toml

[package]
name = "burn-autodiff-race"
version = "0.1.0"
edition = "2024"

[dependencies.burn]
git = "https://github.com/tracel-ai/burn.git"
rev = "1bceb88"
default-features = false
features = ["autodiff", "flex", "std"]

# The race is timing-sensitive: at the stock `opt-level = 0` it still fails, but
# only ~3 rounds in 1800 and the test takes ~5 minutes. With these it is ~190 in
# 1800 and takes ~30 seconds.
[profile.dev]
opt-level = 1

[profile.dev.package."*"]
opt-level = 3

src/main.rs

//! A `backward()` on one thread can delete autodiff graph steps that another
//! thread is still building, so a leaf silently comes back with **no gradient**.
//!
//! ```text
//! cargo test -- --test-threads=1 --nocapture
//! ```
//!
//! `reused_leaf_keeps_its_gradient` fails; `fresh_leaf_keeps_its_gradient`
//! (the control) passes.

fn main() {}

#[cfg(test)]
mod tests {
    use burn::prelude::*;
    use burn::tensor::Distribution;
    use std::sync::atomic::{AtomicBool, Ordering};

    /// Bigger tensors widen the window inside `float_cat`.
    const SIZE: usize = 256;
    const ROUNDS: usize = 300;
    const VICTIM_THREADS: usize = 6;
    const SWEEPER_THREADS: usize = 6;

    fn leaf(size: usize, device: &Device) -> Tensor<2> {
        Tensor::<2>::from_inner(Tensor::<2>::random(
            [size, size],
            Distribution::Normal(0.0, 0.5),
            device,
        ))
        .require_grad()
    }

    /// One round. With `reuse_leaf` the leaf is put through a throwaway
    /// `backward()` first, which is what drops its `GraphLocator` entry and
    /// makes the next op start a one-node graph. Without it — the control —
    /// the leaf keeps its graph and the window never opens.
    fn round(reuse_leaf: bool, device: &Device) -> Result<(), &'static str> {
        let a = leaf(SIZE, device);
        let b = leaf(SIZE, device);

        if reuse_leaf {
            let _ = a.clone().sum().backward();
        }

        // First op on `a` this pass: one node, in a graph of its own.
        let head = a.clone().tanh();
        // `cat` releases `head`'s NodeRefCount before it copies and registers.
        let out = Tensor::cat(vec![head, b.clone().tanh()], 0);
        let grads = out.sum().backward();

        if a.grad(&grads).is_none() {
            return Err("gradient absent for `a`");
        }
        if b.grad(&grads).is_none() {
            return Err("gradient absent for `b`");
        }
        Ok(())
    }

    /// A tight loop of trivial backwards; each one ends in a sweep over every
    /// graph in the process.
    fn sweeper(stop: &AtomicBool, device: &Device) {
        while !stop.load(Ordering::Relaxed) {
            let _ = leaf(8, device).tanh().sum().backward();
        }
    }

    fn stress(reuse_leaf: bool) -> (usize, usize) {
        let device: Device = Default::default();
        let stop = AtomicBool::new(false);

        let failures: usize = std::thread::scope(|s| {
            for _ in 0..SWEEPER_THREADS {
                let device = device.clone();
                let stop = &stop;
                s.spawn(move || sweeper(stop, &device));
            }
            let handles: Vec<_> = (0..VICTIM_THREADS)
                .map(|_| {
                    let device = device.clone();
                    s.spawn(move || {
                        (0..ROUNDS)
                            .filter(|_| round(reuse_leaf, &device).is_err())
                            .count()
                    })
                })
                .collect();
            let failures = handles.into_iter().map(|h| h.join().unwrap()).sum();
            stop.store(true, Ordering::Relaxed);
            failures
        });

        let total = VICTIM_THREADS * ROUNDS;
        eprintln!("reuse_leaf = {reuse_leaf}: {failures} / {total} rounds lost a gradient");
        (failures, total)
    }

    /// The bug: a leaf reused after an earlier `backward()`, concurrently with
    /// other threads' backwards, sporadically comes back with no gradient.
    #[test]
    fn reused_leaf_keeps_its_gradient() {
        let (failures, total) = stress(true);
        assert_eq!(failures, 0, "{failures} / {total} rounds lost a gradient");
    }

    /// Control: the same graph, but the leaf has never been backwarded, so its
    /// first op joins the leaf's existing graph instead of starting a one-node
    /// one. This must never fail.
    #[test]
    fn fresh_leaf_keeps_its_gradient() {
        let (failures, total) = stress(false);
        assert_eq!(failures, 0, "{failures} / {total} rounds lost a gradient");
    }
}

Six threads each build two leaves, run one throwaway backward() on the first, then tanh
it and cat the two branches together; six more threads loop trivial backwards. Only stock
ops, no custom Backward node.

Rounds losing a gradient, out of 1800, on a 12-core box:

burn rev	flex	ndarray
1893b60 (0.22.0-pre.2)	280	553
13f0a12 (0.22.0-pre.3)	283	581
1bceb88 (main)	251 / 268 / 281 / 312 (four runs)	580

fresh_leaf_keeps_its_gradient — the control — is 0 / 1800 in every one of those
configurations. A from-scratch project with exactly the two files above gives 188 / 1800
on flex + main.

The rate is timing-sensitive, so it moves with the build profile and the core count: at
the stock opt-level = 0 the same code loses only ~3 rounds in 1800 (and takes ~5 minutes
instead of ~30 seconds), which is why the profile above is part of the reproduction.

Expected behavior

a.grad(&grads) returns Some for every leaf that contributed to the loss, regardless of
what other threads in the process are doing. A backward() on one thread should not be
able to observe, let alone free, another thread's graph state.

Additional context

The window, step by step

1. AutodiffServer::backward finishes by removing every consumed node from the
   process-global GraphLocator. A leaf survives the backward as a tensor, but its
   locator entry does not.

2. So the next op on that same leaf finds no graph for its parent
   (GraphLocator::analyse → graphs.is_empty()) and calls new_graph: a brand-new
   graph holding exactly one node.

3. float_cat (burn-autodiff/src/ops/tensor.rs) moves tensor.node and
   tensor.primitive out of each input and drops the remainder — releasing the input's
   NodeRefCount — and only then concatenates and registers its own step:

   tensors.into_iter().for_each(|tensor| {
       dim_sizes.push(tensor.primitive.shape()[dim]);
       nodes.push(tensor.node);
       primitives.push(tensor.primitive);
       // `tensor.rc` is dropped here
   });
   ...
   let output = B::float_cat(primitives, dim);   // the whole copy runs with rc == 1
   ...
   output.register_step(ops, checkpointer_builder)

   For the duration of that copy the one-node graph from (2) has
   Arc::strong_count == 1 on its only node, so maybe_useful() answers false even
   though a child step is about to reference it.

4. Another thread finishing a backward() sweeps that graph, finds it un-useful, and
   free_unused_roots frees the node: it is a root of its graph (parents_absent — its
   parent leaf lives in the previous pass's graph) and carries no status (statuses is
   empty on a graph that has never been backwarded), so clear_unused_roots collects it
   on strong_count == 1 alone.

5. The cat then registers a child of a node whose step no longer exists. build_tape's
   traversal stops there and the leaf below never receives a gradient.

Evidence that the sweep is the cause

Measured on the real workload that surfaced this (a chunked SSM training kernel's test
suite), 16 test threads, same binary, same box, against an instrumented burn-autodiff:

build	failing runs
stock	16 / 150
cleanup_orphaned_entries() made a no-op	0 / 150
sweep kept, but the locator entries not removed (only the steps freed)	13 / 150

The third row locates the damage in the freeing of the steps, not in the GraphLocator
bookkeeping that follows.

Tracing what the sweep collects: every collection observed during a failing run was
1 of 1 node, 156 of 186 of them from a graph created by a different thread — exactly the
shape step (3) predicts. Tagging a graph when a foreign thread's sweep collects from it,
then reporting any later use of that graph, fired in 7 of 12 failing runs and 0 of 68
passing runs.

Thread-count dependence of the same workload: 0/200 at 1 thread, 1/100 at 2, 9/100 at 4,
12/100 at 16.

Why it matters beyond a stress test: the trigger is "parameters reused across successive
backward() calls" — i.e. any training loop — plus "more than one thread calling
backward() in the process". Two models training on two threads over a CPU backend can
silently lose a parameter's gradient.