SD3 Controlnet Train Example, run out of memory on validation step

[Bex0n]

Describe the bug

On default settings provided in SD3 controlnet example, with 2 validation images training will error out with out of memory during validation on single A100 80GB.

04/07/2025 21:15:15 - INFO - __main__ - ***** Running training *****    
04/07/2025 21:15:15 - INFO - __main__ -   Num examples = 10000                 
04/07/2025 21:15:15 - INFO - __main__ -   Num batches each epoch = 10000       
04/07/2025 21:15:15 - INFO - __main__ -   Num Epochs = 2                                                                                                      
04/07/2025 21:15:15 - INFO - __main__ -   Instantaneous batch size per device = 1                                                                             
04/07/2025 21:15:15 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 4               
04/07/2025 21:15:15 - INFO - __main__ -   Gradient Accumulation steps = 4      
04/07/2025 21:15:15 - INFO - __main__ -   Total optimization steps = 4000                                                                                     
Steps:   0%|          | 5/4000 [00:21<4:38:36,  4.18s/it, loss=0.00669, lr=1e-5]04/07/2025 21:15:36 - INFO - __main__ - Running validation... 
{'controlnet', 'image_encoder', 'feature_extractor'} was not found in config. Values will be initialized to default values.                                   Keyword arguments {'safety_checker': None} are not expected by StableDiffusion3ControlNetPipeline and will be ignored.                                                                                                             Loaded tokenizer as CLIPTokenizer from `tokenizer` subfolder of stabilityai/stable-diffusion-3.5-medium.e components...:   0%|          | 0/9 [00:00<?, ?it/s]                                                                                         {'invert_sigmas', 'base_shift', 'base_image_seq_len', 'use_dynamic_shifting', 'shift_terminal', 'time_shift_type', 'use_exponential_sigmas', 'max_shift', 'use
_karras_sigmas', 'max_image_seq_len', 'use_beta_sigmas'} was not found in config. Values will be initialized to default values.
Loaded scheduler as FlowMatchEulerDiscreteScheduler from `scheduler` subfolder of stabilityai/stable-diffusion-3.5-medium.                
Instantiating AutoencoderKL model under default dtype torch.float32.                                                                                          
All model checkpoint weights were used when initializing AutoencoderKL.                                                                                       
                                                                                                                                                              
All the weights of AutoencoderKL were initialized from the model checkpoint at /home/jakubdawidowicz/.cache/huggingface/hub/models--stabilityai--stable-diffus
ion-3.5-medium/snapshots/b940f670f0eda2d07fbb75229e779da1ad11eb80/vae.

If your task is similar to the task the model of the checkpoint was trained on, you can already use AutoencoderKL for predictions without further training.
Loaded vae as AutoencoderKL from `vae` subfolder of stabilityai/stable-diffusion-3.5-medium.
                                                                             Loaded text_encoder as CLIPTextModelWithProjection from `text_encoder` subfolder of stabilityai/stable-diffusion-3.5-medium.    | 3/9 [00:00<00:00, 13.70it/s]
Loaded tokenizer_2 as CLIPTokenizer from `tokenizer_2` subfolder of stabilityai/stable-diffusion-3.5-medium.
                                                                             Loaded text_encoder_2 as CLIPTextModelWithProjection from `text_encoder_2` subfolder of stabilityai/stable-diffusion-3.5-medium.| 5/9 [00:00<00:00,  8.40it/s]
Loading checkpoint shards: 100%|██████████| 2/2 [00:12<00:00,  6.02s/it]
Loaded text_encoder_3 as T5EncoderModel from `text_encoder_3` subfolder of stabilityai/stable-diffusion-3.5-medium.
Loading checkpoint shards: 100%|██████████| 2/2 [00:12<00:00,  5.98s/it]     Loaded tokenizer_3 as T5TokenizerFast from `tokenizer_3` subfolder of stabilityai/stable-diffusion-3.5-medium..:  78%|███████▊  | 7/9 [00:14<00:07,  3.66s/it]
                                                                             Instantiating SD3Transformer2DModel model under default dtype torch.float32.
All model checkpoint weights were used when initializing SD3Transformer2DModel.

All the weights of SD3Transformer2DModel were initialized from the model checkpoint at /home/jakubdawidowicz/.cache/huggingface/hub/models--stabilityai--stable-diffusion-3.5-medium/snapshots/b940f670f0eda2d07fbb75229e779da1ad11eb80/transformer.
If your task is similar to the task the model of the checkpoint was trained on, you can already use SD3Transformer2DModel for predictions without further training.
Loaded transformer as SD3Transformer2DModel from `transformer` subfolder of stabilityai/stable-diffusion-3.5-medium.
Loading pipeline components...: 100%|██████████| 9/9 [00:26<00:00,  3.00s/it]
Traceback (most recent call last):0%|██████████| 9/9 [00:26<00:00,  5.44s/it]
  File "/home/jakubdawidowicz/diffusers/examples/controlnet/train_controlnet_sd3.py", line 1429, in <module>
    main(args)
  File "/home/jakubdawidowicz/diffusers/examples/controlnet/train_controlnet_sd3.py", line 1377, in main
    image_logs = log_validation(
                 ^^^^^^^^^^^^^^^
  File "/home/jakubdawidowicz/diffusers/examples/controlnet/train_controlnet_sd3.py", line 83, in log_validation
    pipeline = pipeline.to(torch.device(accelerator.device))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/diffusers/src/diffusers/pipelines/pipeline_utils.py", line 482, in to
    module.to(device, dtype)
  File "/opt/diffusers/src/diffusers/models/modeling_utils.py", line 1351, in to
    return super().to(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1343, in to
    return self._apply(convert)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 903, in _apply
    module._apply(fn)
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 903, in _apply
    module._apply(fn)
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 903, in _apply
    module._apply(fn)
  [Previous line repeated 1 more time]
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 930, in _apply
    param_applied = fn(param)
                    ^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1329, in convert
    return t.to(
           ^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 54.00 MiB. GPU 0 has a total capacity of 79.25 GiB of which 4.75 MiB is free. Including non-PyTorch memory, this process has 79.24 GiB memory in use. Of the allocated memory 76.93 GiB is allocated by PyTorch, and 1.81 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Reproduction

Running SD3 controlnet example (adjusted dataset size and validation steps to decrease job time).

export MODEL_DIR="stabilityai/stable-diffusion-3.5-medium"
export OUTPUT_DIR="sd3-controlnet-out"

accelerate launch train_controlnet_sd3.py \
    --pretrained_model_name_or_path=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --train_data_dir="fill50k" \
    --resolution=1024 \
    --learning_rate=1e-5 \
    --max_train_samples=10000 \
    --max_train_steps=4000 \
    --checkpointing_steps=500 \
    --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
    --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
    --validation_steps=5 \
    --train_batch_size=1 \
    --gradient_accumulation_steps=4

Logs

System Info

System Info

• 🤗 Diffusers version: 0.33.0.dev0
• Platform: Linux-6.8.0-53-generic-x86_64-with-glibc2.39
• Running on Google Colab?: No
• Python version: 3.12.9
• PyTorch version (GPU?): 2.6.0+cu124
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.29.3
• Transformers version: 4.50.3
• Accelerate version: 1.5.2
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.5.3
• xFormers version: not installed
• Accelerator: NVIDIA A100 80GB PCIe, 81920 MiB

Who can help?

@sayakpaul

[sayakpaul]

Would it be possible to call enable_model_cpu_offload() on the pipeline object before calling the validation function?

[sayakpaul]

Interesting. Then it should be possible to replicate this issue without having to run the entire training script. Could you please provide a minimal reproducer with just the ControlNe pipeline and the input image?

[Bex0n]

Pipeline object is initialized inside validation function, changed pipeline = pipeline.to(torch.device(accelerator.device)) to pipeline.enable_model_cpu_offload(), but the error propagated until generating validation samples:

Traceback (most recent call last):
  File "/home/jakubdawidowicz/diffusers/examples/controlnet/train_controlnet_sd3.py", line 1430, in <module>
    main(args)
  File "/home/jakubdawidowicz/diffusers/examples/controlnet/train_controlnet_sd3.py", line 1378, in main
    image_logs = log_validation(
                 ^^^^^^^^^^^^^^^
  File "/home/jakubdawidowicz/diffusers/examples/controlnet/train_controlnet_sd3.py", line 116, in log_validation
    image = pipeline(
            ^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/diffusers/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py", line 1174, in __call__
    control_block_samples = self.controlnet(
                            ^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/accelerate/hooks.py", line 171, in new_forward
    args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/accelerate/hooks.py", line 722, in pre_forward
    module.to(self.execution_device)
  File "/opt/diffusers/src/diffusers/models/modeling_utils.py", line 1351, in to
    return super().to(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1343, in to
    return self._apply(convert)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 903, in _apply
    module._apply(fn)
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 903, in _apply
    module._apply(fn)
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 903, in _apply
    module._apply(fn)
  [Previous line repeated 3 more times]
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 930, in _apply
    param_applied = fn(param)
                    ^^^^^^^^^
  File "/opt/miniconda/envs/control/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1329, in convert
    return t.to(
           ^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 36.00 MiB. GPU 0 has a total capacity of 79.25 GiB of which 30.75 MiB is free. Including non-PyTorch memory, this process has 79.21 GiB memory in use. Of the allocated memory 77.74 GiB is allocated by PyTorch, and 988.33 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

[Johnson-yue]

@Bex0n I have same question, GPU is A100 ,,80G , but is still OOM in log_validation() functions

[Johnson-yue]

I think the problem is 80G GPU can not load two SD3 pipeline, so in log_validation() function , the SD3 pipeline should reuse the training weight , not load from_pretrained ,maybe work

[sayakpaul]

Checking it.

[Johnson-yue]

Another question , in the #Run a final round of validation. why set the controlnet is None in log_validation()??

diffusers/examples/controlnet/train_controlnet_sd3.py

Line 1401 in f685981

controlnet=None,

It is means the trained controlnet not test in log_validation() during a final round ??

[sayakpaul]

It's loaded from the output_dir:

diffusers/examples/controlnet/train_controlnet_sd3.py

Line 73 in f685981

controlnet = SD3ControlNetModel.from_pretrained(args.output_dir, torch_dtype=weight_dtype)

The output_dir is expected to contain the trained params of the ControlNet as we run:

diffusers/examples/controlnet/train_controlnet_sd3.py

Line 1395 in f685981

controlnet.save_pretrained(args.output_dir)

[Johnson-yue]

@sayakpaul oh， I see，you are right.

It's loaded from the output_dir:

diffusers/examples/controlnet/train_controlnet_sd3.py

Line 73 in f685981

controlnet = SD3ControlNetModel.from_pretrained(args.output_dir, torch_dtype=weight_dtype)
The output_dir is expected to contain the trained params of the ControlNet as we run:

diffusers/examples/controlnet/train_controlnet_sd3.py

Line 1395 in f685981

controlnet.save_pretrained(args.output_dir)

[sayakpaul]

Additionally, here's the fix: #11238

[Johnson-yue]

A question about process 'controlnet_image',
In train phase:
controlnet_image -> conditioning_image_transforms convert data type from PIL.Image to (0, 1) tensor -> then into vae_encoder

but in log_vailidation():
controlnet_image ->StableDiffusion3ControlNetPipeline.calll() function ,and StableDiffusion3ControlNetPipeline.prepare_image() function convert PIL.Image to (-1, 1) tensor, because vae.config.do_normalize is True. -> then into vae_encoder

the result is in log_validation() pipeline out image is nan, tensorboard log image is all black

and command line has a warning

/git_repos/diffusers/src/diffusers/image_processor.py:147: RuntimeWarning: invalid value encountered in cast
images = (images * 255).round().astype("uint8")

[sayakpaul]

Can you investigate this yourself by printing the value ranges? Also, can we please move to a separate discussion instead of an issue thread as we're digressing from the original reported issue?