Hitting addition with overflow in FlushingPolicyState

[jeandudey]

System information:

• GPU: 7900 XTX 24 GB VRAM (not an OOM error since it seems to only use half of it).
• OS: Ubuntu 24.04 on Windows 11 using WSL2.
• Backend: Using ROCm with librocdxg on commit 9fa6ae22c722201a43901788f371a6904f81f600.
• ROCm version: 6.4.2
• Burn version: 0.21 from crates.io
• rustc 1.95.0 (59807616e 2026-04-14)

Disclaimer: this could very well me being the problem since I'm new to machine learning and I'm currently learning by trying to train a model.

I'm currently hitting an attempt add with overflow in the register method of FlushingPolicyState when trying to train a model using resnet-burn and burn.

The experiment.log, so the training goes well up until the point that it crashes with a panic (training tries to keep working but keeps failing):

thread 'DSD-0-0' (93753) panicked at /home/jeandudey/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/cubecl-runtime-0.10.0/src/memory_management/drop_queue/policy.rs:36:9:
attempt to add with overflow
=== PANIC ===
A fatal error happened, you can check the experiment logs here => './artifacts/experiment.log'
=============

I've tried searching if anyone had hit this but I didn't find anything, the code that reproduces this error can be found here:

https://github.com/jeandudey/burn-experiment/blob/main/src/main.rs

On commit: 77f18b171cb1c0aac34feed8b09bc033f6d1633d

It requires libmatio-dev and pkg-config and the 300W-LP dataset extracted in repo-directory/datasets/300W_LP being 300W_LP the root of the dataset.

I'm not sure what could be the cause of this, and I found that increasing the number of workers for the data loaders makes it happen quicker.

[djarb]

I ran into the same problem, using CUDA rather than ROCm. Given what I found in the code, it's not surprising that this is backend-agnostic. The problem appears to be the use of a u32 in the flushing policy, which can only describe byte sizes up to 4 GiB.

My specific case was loading the SDXL UNet (~10 GiB of staged weights on first load) through burn 0.21 + cubecl-runtime 0.10.0. Patched locally by widening bytes_size from u32 to u64. Diff in crates/cubecl-runtime/src/memory_management/drop_queue/policy.rs (no other call sites need changes):

   pub struct FlushingPolicy {
       /// Flush when this many allocations have been staged.
       pub max_bytes_count: u32,
       /// Flush when the total staged size reaches this many bytes.
  -    pub max_bytes_size: u32,
  +    pub max_bytes_size: u64,
   }
  
   impl Default for FlushingPolicy {
       fn default() -> Self {
           Self {
               max_bytes_count: 64,
               max_bytes_size: 64 * 1024 * 1024, // 64 MiB
           }
       }
   }

   #[derive(Default, Debug)]
   pub(crate) struct FlushingPolicyState {
       bytes_count: u32,
  -    bytes_size: u32,
  +    bytes_size: u64,
   }

   impl FlushingPolicyState {
       pub(crate) fn register(&mut self, bytes: &Bytes) {
           self.bytes_count += 1;
  -        self.bytes_size += bytes.len() as u32;
  +        self.bytes_size += bytes.len() as u64;
       }
       ...
   }

[for-lack-of-a-better-name-j]

I'm looking into this, I'll let you know if I find anything.

[for-lack-of-a-better-name-j]

Okay folks, quick update.

Platform details

OS: Docker container rocm/dev-ubuntu-24.04:6.4.2-complete on Arch Linux
Backend: ROCm 6.4.2
Burn version: 0.21.0
rustc 1.96.0
GPU: Radeon 7900XTX

I believe I may be reproducing the bug here, when I bring TrainingConfig.num_workers up quite a bit (38) my GPU's VRAM utilization starts going up like crazy and I get repeated panics failing during memory initialization:

0x5574dd457f38 - <cubecl_hip::compute::server::HipServer as cubecl_runtime::server::base::ComputeServer>::initialize_memory::h85e007c926d16095 at /root/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/cubecl-hip-0.10.0/src/compute/server.rs:84:46

I'm still getting familiar with this project, but my working theory is that we might be looking at a producer-consumer backpressure issue. I suspect that the reason why jeandudey didn't have this specific error comes down to differences between how WSL and native Linux allocate (and spill) GPU memory. I'll be digging deeper to see if I can find the root cause. Curious if anyone else has experienced this behavior on HIP.

[for-lack-of-a-better-name-j]

Hey folks, quick update. Still digging. I have found that my memory issue (hip initialize memory error) is likely unrelated to the u32 overflow previously mentioned. @djarb are you also on WSL? If not what platform?

[djarb]

Hey folks, quick update. Still digging. I have found that my memory issue (hip initialize memory error) is likely unrelated to the u32 overflow previously mentioned. @djarb are you also on WSL? If not what platform?

Arch Linux on amd64, up to date with the rolling release.

[for-lack-of-a-better-name-j]

So it turns out that the staging queues in the flushing policy can receive really big chunks of memory quite fast. I am not entirely certain yet but this may have caused the u32 overflow. I'll edit this comment with more detail in a day or two.

[for-lack-of-a-better-name-j]

I've been digging into the backend stream pipeline and have a working theory. On my machine, allocated chunks along this path almost never exceed 100MiB, so hitting a u32 tracking overflow is definitely an anomaly.

I noticed that FlushingPolicyState's should_flush is only evaluated when a registered kernel is executed. Because of this cadence, it seems entirely possible for the byte counter to get overfull or overflow between kernel runs if a rapid burst of allocations piles up asynchronously (like from aggressive dataloader prefetching) before hitting the next evaluation gate.

I'm also looking into a couple of alternative ideas regarding the buffer swapping logic to rule out any compounding ghost allocations. I'll go see if I can verify any of these ideas and I'll give another update.

[for-lack-of-a-better-name-j]

So far, I can see that the path all the way from loaded data to the DSD thread where flushing is called has bounds on the channel lengths that are used to send data, but I haven't seen anything that limits the size of the actual bytes that are being sent and registered into that FlushingPolicyState. I'm currently digging to see if I can find anything that limits that allocation size. My working theory is that in some cases the memory that FlushingPolicyState represents can get more than 4GiB if for some reason the GPU doesn't process things fast enough. I'll let you guys know what I find out.

[for-lack-of-a-better-name-j]

Folks, if I'm not seeing this wrong, today I actually found a codepath that allows multiple calls to push method of FlushingPolicy (and therefore calls FlushingPolicyState's register) without doing a should_flush check. I actually found this by accident while examining this in a debugger.

I'll be running the debugger again to see if I can repeat, and if I can, then I'll be able to check if those multiple calls to write_to_gpu are bounded or not bounded enough to cause our bug.

[for-lack-of-a-better-name-j]

Bug Reproduced!

Folks, I have successfully reproduced the bug:

Why it Happens

This bug occurs when multiple tensor allocations, in total containing greater than 4,294,967,295 bytes, are allocated between kernel launches.

How it Happens

The reason why this happens is because on tensor allocation, current.drop_queue.push(data) is run without doing a current.drop_queue.should_flush() when a CubeTensor is written to the gpu through do_create and write_to_gpu. While current.drop_queue.should_flush() is called on kernel launch, it is not called on the tensor creation codepath.

The Code to Reproduce the Bug

I originally forgot to add in the code that I wrote to trigger the bug:

fn trigger_overflow_burn<B: Backend>(device: &B::Device) {
    let mut tensors = Vec::new();
    let shape = [625_000_000]; // 625,000,000 f32s * 4 bytes = 2.5gb,
    println!("Big data dump into allocation queue");
    for i in 0..3 {
        println!("creating tensor {i}");
        let data = TensorData::ones::<f32, _>(shape);
        let tensor = Tensor::<B, 1>::from_data(data, device);
        tensors.push(tensor);
    }
    println!("chain some ops, force load to GPU and make it crash");
    let mut computation = tensors[0].clone();

    for tensor in tensors.iter().skip(1) {
        computation = computation * tensor.clone();
    }
    let raw_data = computation.into_data();
    println!("raw_data: {:?}", raw_data);
}

Two questions for maintainers

1. Is it intentional that FlushingPolicyState's bytes_size is a u32 and not a u64? I was thinking it was acting as a canary for overly large memory allocations.
2. Is it intentional that the new tensor creation path does not call PendingDropQueue::should_flush?

@nathanielsimard