train_dreambooth_lora_flux.py distributed bugs

[neuron-party]

Describe the bug

AttributeError when running model parallel distributed training with accelerate

Reproduction

accelerate launch --config_file train_dreambooth_lora_flux.py
--resolution=1024
--mixed_precision=bf16
--pretrained_model_name_or_path=black-forest-labels/FLUX.1-dev
--num_validation_images=8
--validation_epochs=100
--rank=16
--train_batch_size=1
--learning_rate=1e-4
--guidance_scale=3.5
--checkpointing_steps=200
--instance_prompt=xyz
--instance_data_dir=xyz
--output_dir=xyz
--logging_dir=xyz
--validation_prompt=xyz

accelerate config:

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MULTI_GPU
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
use_cpu: false
gpu_ids: '0, 1'
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false

Logs

if transformer.config.guidance_embeds:

AttributeError: DistributedDataParallel object has no attribute config

System Info

diffusers from source
accelerate==0.33.0
transformers==4.44.1

training on A100s

Who can help?

No response

[tolgacangoz]

What happens if you unwrap it:

if accelerator.unwrap_model(transformer).config.guidance_embeds:

[maziyarpanahi]

Any progress on this? The full fine-tune gets OOM on 4x A100/80G and the LoRA results in this error.

What happens if you unwrap it:

if accelerator.unwrap_model(transformer).config.guidance_embeds:

This exist in the full fine-tune, there is no config.guidance_embeds inside train_dreambooth_lora_flux.py file.

[Adenialzz]

What happens if you unwrap it:

if accelerator.unwrap_model(transformer).config.guidance_embeds:

this works in my case.

But I got an another pytorch oom error in log_validation, even using a800 80g. How can I fix this?

[tolgacangoz]

For OOM, see https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md. Did you add --gradient_checkpointing, --use_8bit_adam, and --gradient_accumulation_steps=4 or 8? Is it possible for you to try without --validation_prompt? Also, could you try to use one of the latest versions of PyTorch?

@maziyarpanahi Could you elaborate? Isn't config.guidance_embeds in train_dreambooth_lora_flux.py:

diffusers/examples/dreambooth/train_dreambooth_lora_flux.py

Line 1656 in 214990e

if transformer.config.guidance_embeds:

[maziyarpanahi]

@maziyarpanahi Could you elaborate? Isn't config.guidance_embeds in train_dreambooth_lora_flux.py:

Sorry about that, I was looking in the wrong file. It does indeed exist in the lora file as well.

[Adenialzz]

For OOM, see https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md. Did you add --gradient_checkpointing, --use_8bit_adam, and --gradient_accumulation_steps=4 or 8? Is it possible for you to try without --validation_prompt? Also, could you try to use one of the latest versions of PyTorch?

@maziyarpanahi Could you elaborate? Isn't config.guidance_embeds in train_dreambooth_lora_flux.py:

diffusers/examples/dreambooth/train_dreambooth_lora_flux.py

Line 1656 in 214990e

if transformer.config.guidance_embeds:

thanks. I missed this oom guidance before. It helps a lot.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.