feat(burn-cubecl): add deterministic feature for reproducible float reductions

[guoyucode]

Summary

Multi-step GPU training can produce different weights across identical runs even when application-level seeds, data order, and initialization are fixed. The root cause is in CubeCL float reduction: Tensor::mean() / float_sum use SumStrategy::default(), which with the default autotune feature resolves to SumStrategy::Autotune. Autotune frequently selects OneShot, whose shared_sum kernel finishes with a cross-cube float atomicAdd (fetch_add on the output). Floating-point addition is not associative, and atomic completion order is undefined, so the summed value (and therefore gradients / optimizer updates) can drift from run to run. After many Adam steps the drift is large enough to change early-stopping and final checkpoints.

This PR adds an opt-in deterministic feature on burn-cubecl that does not change default behavior:

• With deterministic: SumStrategy::default() → Chained(Unspecified) and KernelReduceStrategy::default() → Unspecified (tree/chained reduce, no cross-cube float atomics).
• Without deterministic: keep the existing Autotune / OneShot defaults.

Users who need bit-reproducible GPU training can enable the feature through the usual Cargo feature graph (e.g. via burn-wgpu / burn once forwarded, or by depending on burn-cubecl with features = [deterministic]).

Tests

Unit tests live next to the change in crates/burn-cubecl/src/kernel/reduce/base.rs:

Build	What it checks
--features deterministic	Default strategies are Chained/Unspecified; multi-step float-sum feedback with SumStrategy::default() is bit-stable across two runs
default features (no deterministic)	Defaults remain Autotune; multi-step feedback with OneShot(8) (Autotune’s common pick) diverges across runs

cargo test -p burn-cubecl --features deterministic multi_step_float_sum
cargo test -p burn-cubecl multi_step_float_sum
cargo test -p burn-cubecl --features deterministic default_sum_strategy
cargo test -p burn-cubecl default_sum_strategy

(GPU / Wgpu required for the multi-step probes; contract tests for Default are CPU-only.)

Test plan

• CI / local: contract tests pass with and without deterministic
• Local GPU: multi-step probe passes with deterministic (Δ ≈ 0) and fails/asserts divergence without it on OneShot
• Confirm default feature set (no deterministic) still builds and behaves as before
• Optional follow-up: forward deterministic from burn-wgpu / burn for a single user-facing switch

[Copilot]

Pull request overview

Adds an opt-in deterministic feature to burn-cubecl to make default float reduction behavior reproducible by avoiding reduction paths that rely on cross-cube floating-point atomics (which can introduce run-to-run drift).

Changes:

• Introduce a deterministic Cargo feature for burn-cubecl.
• Change SumStrategy::default() / KernelReduceStrategy::default() to select deterministic-safe defaults when the feature is enabled.
• Add unit tests intended to enforce the feature contract and probe (GPU) divergence vs reproducibility.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
crates/burn-cubecl/src/kernel/reduce/base.rs	Switch default reduction strategies under deterministic and add contract/probe tests.
crates/burn-cubecl/Cargo.toml	Add deterministic feature and a new dev-dependency for test support.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Charles23R]

That's valid

[Charles23R]

These tests should probably be in burn-wgpu directly (mod multi_step_float_sum)

[Charles23R]

That's valid

[Charles23R]

These tests should probably be in burn-wgpu directly (mod multi_step_float_sum)