Closed
Description
Describe the bug
When training with --mixed_precision bf16 or fp16, the prompt_embeds and pooled_prompt_embeds tensors in the compute_text_embeddings function are not cast to the appropriate weight_dtype (matching the rest of the model inputs and parameters), causing a mismatch error during training.
Specifically, the tensors are generated as float32 (default) and not moved to bf16/fp16, which leads to issues when performing operations inside the transformer/controlnet forward pass.
Reproduction
accelerate launch train_controlnet_sd3.py \
--pretrained_model_name_or_path=$MODEL_DIR \
--output_dir=$OUTPUT_DIR \
--train_data_dir=$DATASET \
--resolution=512 \
--caption_column=$CAPTION_COLLUMN \
--output_dir=$OUTPUT_DIR \
--learning_rate=1e-5 \
--max_train_steps=15000 \
--validation_steps=100 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--mixed_precision=bf16
Logs
System Info
- 🤗 Diffusers version: 0.33.0.dev0
- Platform: Linux-6.5.0-45-generic-x86_64-with-glibc2.35
- Running on Google Colab?: No
- Python version: 3.11.10
- PyTorch version (GPU?): 2.4.1+cu124 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.29.3
- Transformers version: 4.49.0
- Accelerate version: 1.5.1
- PEFT version: not installed
- Bitsandbytes version: not installed
- Safetensors version: 0.5.3
- xFormers version: not installed
- Accelerator: NVIDIA A100 80GB PCIe, 81920 MiB
- Using GPU in script?:
- Using distributed or parallel set-up in script?: