huggingface / huggingface/diffusers

Redundant reinitialization of text encoders in train_dreambooth_lora_flux

Abierto
#9,358 2 comentarios 0 reacciones 0 asignados Ver en GitHub
bug stale
Lenguaje dominante
Python
Estrellas
34.5k
Forks
7.3k
Merge medio
3 d 3 h
PR fusionados (30 d)
91

Descripción

### Describe the bug

In the `train_dreambooth_lora_flux.py` script, during each call to `log_validation`, the text encoders `text_encoder_one` and `text_encoder_two` are reinitialized. https://github.com/huggingface/diffusers/blob/8ba90aa706a733f45d83508a5b221da3c59fe4cd/examples/dreambooth/train_dreambooth_lora_flux.py#L1768
This occurs even when the text encoders do not need to be trained (`if not args.train_text_encoder`). This unnecessary reinitialization can lead to inefficient use of resources and may cause CUDA out-of-memory errors, especially in scenarios where VRAM is less than 48 GiB.

Since the validation prompt is fixed (only one prompt is used), we can optimize the process by precomputing the text embeddings during the instance prompt preprocessing. This would allow the model to fit within 40 GiB of VRAM, preventing CUDA OOM issues.

### Proposed Fix

To address this issue, add the following code snippet to precompute the validation prompt embeddings only once when the text encoders do not need to be trained and custom instance prompts are not used:

```python
if not args.train_text_encoder and not train_dataset.custom_instance_prompts:
instance_prompt_hidden_states, instance_pooled_prompt_embeds, instance_text_ids = compute_text_embeddings(
args.instance_prompt, text_encoders, tokenizers
)

if args.validation_prompt is not None:
validation_prompt_hidden_states, validation_pooled_prompt_embeds, _ = compute_text_embeddings(
args.validation_prompt, text_encoders, tokenizers
)
```

This change will prevent the unnecessary reinitialization of text encoders and reduce the VRAM usage during training.

### Reproduction

``` python
from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
"diffusers/dog-example",
local_dir=local_dir, repo_type="dataset",
ignore_patterns=".gitattributes",
)

export MODEL_NAME="black-forest-labs/FLUX.1-dev"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="trained-flux-lora"

accelerate launch train_dreambooth_lora_flux.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--mixed_precision="bf16" \
--instance_prompt="a photo of sks dog" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--learning_rate=1e-5 \
--report_to="wandb" \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=500 \
--validation_prompt="A photo of sks dog in a bucket" \
--validation_epochs=25 \
--seed="0" \
--push_to_hub
```

### Logs

_No response_

### System Info

- 🤗 Diffusers version: 0.31.0.dev0
- Platform: Linux-5.15.0-47-generic-x86_64-with-glibc2.35
- Running on Google Colab?: No
- Python version: 3.10.8
- PyTorch version (GPU?): 2.4.0+cu121 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.24.5
- Transformers version: 4.43.4
- Accelerate version: 0.30.1
- PEFT version: 0.12.0
- Bitsandbytes version: not installed
- Safetensors version: 0.4.3
- xFormers version: not installed
- Accelerator: NVIDIA A40, 49140 MiB
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help?

@sayakpaul

Guía de contribución

Abrir la guía de contribución

Línea de trabajo

Comienza en examples/dreambooth/train_dreambooth_lora_flux.py, en el código de log_validation alrededor de la línea 1768, y luego sigue el preprocesamiento del prompt de instancia y compute_text_embeddings. Precalcula los embeddings especificados solo para las condiciones indicadas y asegúrate de que la validación ya no reinicialice los text encoders no utilizados. Vuelve a ejecutar el comando de entrenamiento de accelerate proporcionado y verifica la reducción del uso de VRAM sin CUDA OOM.

Escrito por el modelo de indexación a partir del texto del issue.

Evaluación

Stack tecnológico
python, pytorch
Área
machine-learning, performance
Tipo de issue
Error
Dificultad
3/5
Tiempo estimado
1-2 días
Estado de actividad
Estancado
Claridad
Bastante claro
Aptitud para principiantes
50/100

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.