Summary
On the cuda backend, if a training loop reallocates a network + its Adam optimizer on every outer iteration, Optimizer::step gets progressively slower across iterations, and the slowdown spills over to .step() calls on other, persistent networks in the same process. The same code on wgpu (DX12 / Vulkan) is flat.
Environment
- cubecl 0.10, Burn 0.21
- GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition (sm_120)
- CUDA 12.8 toolkit + matching driver (RunPod
pytorch:2.4.0-py3.11-cuda12.8 image)
- rustc 1.92,
--release
- Backend:
burn::backend::Autodiff<burn_cuda::Cuda<f32>>
Reproducer
Minimal example code for reproduction: https://github.com/hypercat/cubecl-cuda-reinit-repro
Two modes over the same persistent network + Adam optimizer + inner loop:
--mode flat: only the persistent network is trained across iters.--mode climb: same persistent network + inner loop, PLUS a second network + fresh Adam optimizer freshly allocated inside each outer iter, trained for the same inner loop, then dropped.
Each stdout line: iter=<N> st=<X>ms, where st is the total wall-clock of the persistent network's Optimizer::step calls in iter N. The persistent network's state and op sequence are byte-identical between the two modes. Only the ephemeral net's per-iter alloc/dealloc differs.
cargo build --release
./target/release/cubecl-cuda-reinit-repro --mode flat --iters 500 # baseline
./target/release/cubecl-cuda-reinit-repro --mode climb --iters 500 # trigger
Symptom
On flat, st stays roughly constant across iters.
On climb, st on the persistent network climbs monotonically across iters. The network's own state doesn't change (it just keeps training on the same shapes); only the process is being asked to alloc+drop a second network + Adam optimizer per outer iter.
Numbers from three Adam-trained networks per outer iter, 50 Optimizer::step calls per network per iter, over 1500 iters, one of the three reallocated per iter (st per network):
Network A st: 1173 -> 1411 ms (+20.2%)
Network B st: 722 -> 1067 ms (+47.8%)
Network C st: 1380 -> 1842 ms (+33.5%)
GradientsParams::from_grads stayed at 1-3 ms flat throughout, so the growth is entirely inside Optimizer::step. fwd, bwd, prep also flat. Same workload on wgpu (DX12): flat end to end.
Isolation
Bisected by dropping each contributor one at a time (500-iter runs each):
| Config |
Persistent net st climb |
| Baseline: 3 networks trained per iter, one reallocated per iter |
+37.6% |
| Only the persistent net trained (drop the other two loops) |
-7.5% flat |
| Persistent + second identical net, neither reallocated |
-0.0% flat |
| Persistent + reallocated-per-iter net |
+41.3% |
| Same, but the reallocated net does half the inner steps |
+87.1% (larger %) |
| Persistent + second net, keep reallocation flag OFF |
-3.1% flat |
Halving the reallocated net's inner-step count increases the climb percentage, which means the accumulator is fixed per outer iter, not per Adam step. It scales with iter count, not with total Adam-step count. The only thing that flips the whole system from flat to climbing (or back) is whether or not the second network + its Adam optimizer are freshly allocated inside the outer loop.
Ruled out
- Autotune cache growth:
CUBECL_AUTOTUNE_LEVEL=minimal shows autotune events collapse to near zero after the first 100 iters, but st keeps climbing past that.
- Autograd graph accumulation:
GradientsParams::from_grads flat at 1-3 ms throughout, so the graph is being detached correctly on .backward().
- Adam step-counter arithmetic: a network trained continuously for 50000-150000 Adam steps without any reallocation stays flat.
- Fixed-shape workloads with no reallocation: 1 to 3 persistent networks + Adam, exhaustively varied (data content per step, mask ops, bootstrap-target-style extra forwards, etc.) all stay flat until a reallocation is added.
Hypothesis
Something in the cuda backend accumulates state across allocations of network params + Adam moment tensors that never fully returns when those tensors drop. Since wgpu doesn't share the code path, it doesn't reproduce. Reasonable suspects: the cuda memory pool's free-list / per-size buckets, or per-tensor caches keyed by handle.
Workaround for users
Keep networks alive across iterations (avoid Network::with_layers(...) + AdamConfig::new().init() inside your training loop). For loops that need "reset the optimizer between iterations" semantically, zero the Adam moment tensors in place rather than dropping and reallocating.
Happy to run additional probes on the same pod if that would help narrow it further from your side.
Summary
On the
cudabackend, if a training loop reallocates a network + its Adam optimizer on every outer iteration,Optimizer::stepgets progressively slower across iterations, and the slowdown spills over to.step()calls on other, persistent networks in the same process. The same code onwgpu(DX12 / Vulkan) is flat.Environment
pytorch:2.4.0-py3.11-cuda12.8image)--releaseburn::backend::Autodiff<burn_cuda::Cuda<f32>>Reproducer
Minimal example code for reproduction: https://github.com/hypercat/cubecl-cuda-reinit-repro
Two modes over the same persistent network + Adam optimizer + inner loop:
--mode flat: only the persistent network is trained across iters.--mode climb: same persistent network + inner loop, PLUS a second network + fresh Adam optimizer freshly allocated inside each outer iter, trained for the same inner loop, then dropped.Each stdout line:
iter=<N> st=<X>ms, wherestis the total wall-clock of the persistent network'sOptimizer::stepcalls in iter N. The persistent network's state and op sequence are byte-identical between the two modes. Only the ephemeral net's per-iter alloc/dealloc differs.Symptom
On
flat,ststays roughly constant across iters.On
climb,ston the persistent network climbs monotonically across iters. The network's own state doesn't change (it just keeps training on the same shapes); only the process is being asked to alloc+drop a second network + Adam optimizer per outer iter.Numbers from three Adam-trained networks per outer iter, 50
Optimizer::stepcalls per network per iter, over 1500 iters, one of the three reallocated per iter (stper network):GradientsParams::from_gradsstayed at 1-3 ms flat throughout, so the growth is entirely insideOptimizer::step.fwd,bwd,prepalso flat. Same workload onwgpu(DX12): flat end to end.Isolation
Bisected by dropping each contributor one at a time (500-iter runs each):
stclimbHalving the reallocated net's inner-step count increases the climb percentage, which means the accumulator is fixed per outer iter, not per Adam step. It scales with iter count, not with total Adam-step count. The only thing that flips the whole system from flat to climbing (or back) is whether or not the second network + its Adam optimizer are freshly allocated inside the outer loop.
Ruled out
CUBECL_AUTOTUNE_LEVEL=minimalshows autotune events collapse to near zero after the first 100 iters, butstkeeps climbing past that.GradientsParams::from_gradsflat at 1-3 ms throughout, so the graph is being detached correctly on.backward().Hypothesis
Something in the
cudabackend accumulates state across allocations of network params + Adam moment tensors that never fully returns when those tensors drop. Sincewgpudoesn't share the code path, it doesn't reproduce. Reasonable suspects: the cuda memory pool's free-list / per-size buckets, or per-tensor caches keyed by handle.Workaround for users
Keep networks alive across iterations (avoid
Network::with_layers(...)+AdamConfig::new().init()inside your training loop). For loops that need "reset the optimizer between iterations" semantically, zero the Adam moment tensors in place rather than dropping and reallocating.Happy to run additional probes on the same pod if that would help narrow it further from your side.