FLUX GPU training OOM

[hx89]

When testing FLUX GPU training with FSDP=8 using Madiffusion main branch and the same run script that worked before, we got OOM error.

The log had multiple error messages indicating XLA is allocating several large memories for some collectives such as:

allocator.cc:62 NCCL WARN Cuda failure 2 'out of memory'
external/xla/xla/stream_executor/integrations/stream_executor_allocator.cc:66] could not allocate collective of size: 20289079296

Our run command is:

python3 src/maxdiffusion/train_flux.py src/maxdiffusion/configs/base_flux_schnell.yml hardware=gpu run_name=flux attention=cudnn_flash_te max_train_steps=20 enable_profiler=True profiler_steps=5 skip_first_n_steps_for_profiler=10 profiler=xplane train_new_flux=True

I have tried various XLA_PYTHON_CLIENT_MEM_FRACTION values but didn't help.

[entrpn]

@hx89 can you try without profiling and without flash attention. What's the memory of your GPU?

[hx89]

Thanks @entrpn ! I was using 8 b200 GPUs with 192GB memory each.

I have disabled profiling and tried dot_product instead of cudnn_flash_te, it also OOM, and the size of the collective that could not be allocated is larger:

2025-08-01 15:07:01.437102: W external/xla/xla/stream_executor/integrations/stream_executor_allocator.cc:66] could not allocate collective of size: 47811428352

Could there be anything wrong in the collective allocation?

[entrpn]

@hx89 can you test with an older commit from when it was working and verify that works? I want to see if this issue could be related to jax or something that changed in maxdiffusion so that I can follow up. Thank you!