[examples] ]Error due to mismatches in the dtype of UNet while running train_text_to_image.py

@kashif

Describe the bug

I am referring to this script: https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py.

To perform qualitative validation and following https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py, I have added some sample inference code:

        if args.validation_prompt is not None and epoch % args.validation_epochs == 0:
            logger.info(
                f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
                f" {args.validation_prompt}."
            )
            # create pipeline
            temp_unet = accelerator.unwrap_model(unet)
            pipeline = StableDiffusionPipeline.from_pretrained(
                args.pretrained_model_name_or_path,
                unet=temp_unet,
                revision=args.revision,
                torch_dtype=weight_dtype
            )
            pipeline = pipeline.to(accelerator.device)
            pipeline.set_progress_bar_config(disable=True)

            # run inference
            generator = torch.Generator(device=accelerator.device).manual_seed(
                args.seed
            )
            images = []
            for _ in range(args.num_validation_images):
                images.append(
                    pipeline(
                        args.validation_prompt,
                        num_inference_steps=30,
                        generator=generator,
                    ).images[0]
                )

Training runs smoothly but when it gets to the pipeline calling part, it runs into an error complaining about the mismatches in the dtypes of the model (UNet) and the inputs.

Cc: @kashif

Reproduction

Code: https://gist.github.com/sayakpaul/97ec0007ec423960912ca8822e4cb7be

Command:

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export dataset_name="lambdalabs/pokemon-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=250 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --validation_prompt="cute dragon creature" \
  --seed=666 \
  --output_dir="sd-pokemon-model"

Logs

No response

System Info

- `diffusers` version: 0.13.0.dev0
- Platform: Linux-4.19.0-23-cloud-amd64-x86_64-with-glibc2.10
- Python version: 3.8.15
- PyTorch version (GPU?): 1.13.1+cu116 (True)
- Huggingface_hub version: 0.12.0
- Transformers version: 4.26.0
- Accelerate version: 0.15.0
- xFormers version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

I just tried this with train_text_to_image_lora.py, script and couldn't reproduce, maybe check if the modifications you did to text_to_image script are similar.

Yes, I am aware that script won't lead to this issue. The setup here is a bit different, no since the actual UNet parameters are not being used in the fine-tuning process when using LoRA.

I will check the train_unconditional.py script (reference) to see if there's something I am missing.

Aah, you are right! This could be an issue if unet is being passed to accelerate.prepare. Before this, would you like to open a PR to add logging to text_to_image script, so we'll debug it there?

Let me do that :)

@patrickvonplaten

I looked into it a bit. This happens because when using mixed_precision, the vae and text_encoder weights are cast to fp16, while the unet is still in fp32 (but wrapped under a accelerate wrapper which uses autocast during forward and then casts the output back to fp32, classic mixed-precision training.) There are two cases here

When we unwrap the unet , it's in fp32, without the accelerate autocast wrapper, so it can't accept fp16 input. Hence the mismatch.
If we use the unet without unwrapping, the unet will run fine (because of the wrapper), but the output will be casted back to fp32 and then vae decode will fail.

I think the safest way here would be to use torch.autocast for inference when doing mixed-precision training.
cc @patrickvonplaten @pcuenca

@isamu-isozaki

Why don't we run inference in float32 in this case? I think it should be the baseline for diffusers, even when training in mixed precision.

In any case, we need to settle on a single method and apply it to all scripts. Some discussion was initiated by @isamu-isozaki in #2077 (comment)

Why don't we run inference in float32 in this case? I think it should be the baseline for diffusers, even when training in mixed precision.

I am opening a PR, let's continue the discussion there?

(Said PR: #2173)

@pcuenca

Why don't we run inference in float32 in this case? I think it should be the baseline for diffusers, even when training in mixed precision.

@pcuenca for mixed-precision, then we'll have to cast the vae and text_encoder to fp32 before inference and then back to fp16 for training.

@patil-suraj

@patil-suraj From @sayakpaul's snippet posted above, I would do this:

            pipeline = StableDiffusionPipeline.from_pretrained(
                args.pretrained_model_name_or_path,
                unet=accelerator.unwrap_model(unet),
                revision=args.revision,
            )

i.e., just remove torch_dtype. We are loading the components from the pretrained model anyway.

I see, that looks good. My concern is, wouldn't it require more memory, then? The vae and text_encoder will be newly loaded and moved to GPU.

If that's a concern then we need to reuse the components, yes (which was not happening above either). For me the question is more about what's diffusers supposed to default to. In our design philosophy we state that the default is float32 unless overridden by the user. Should we behave differently while evaluating training runs?

The default is indeed fp32 when not doing mixed precision. It'll only use half precision when--mixed_precision="fp16", else it'll be in fp32.

Up for reviews and discussions: #2173

@sayakpaul

@sayakpaul Thanks for the issue! And @pcuenca , thanks for the tag!
I think it'll be good as you mentioned in #2077 to have the option to load everything in fp16+reuse components but not by default. When I was testing around, reusing components+using torch_dtype=fp16 was the only way to do logging in 6gb of vram without oom. Also, the results are not visibly bad in my testing.
But I think this can be a separate pr as you mentioned too. Will make one today!
One idea I had was to save torch_dtype as an attribute so that we can convert all the intermediate tensors to that. Especially if we are reusing components since it won't be obvious which model will have the "correct" dtype.

Up.

Let's use torch.autocast to solve the issue no? As explained here: #2568

This is an end-to-end example that works with mixed-precision: https://gist.github.com/sayakpaul/075ee224a5332779710f6b649bfa304d

Closing the issue henceforth.

Navigation Menu

Repository navigation

[examples] ]Error due to mismatches in the dtype of UNet while running `train_text_to_image.py` #2163

Describe the bug

Reproduction

Logs

System Info

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[examples] ]Error due to mismatches in the dtype of UNet while running train_text_to_image.py #2163

Description

Describe the bug

Reproduction

Logs

System Info

Activity

patil-suraj commented on Jan 31, 2023

sayakpaul commented on Jan 31, 2023

patil-suraj commented on Jan 31, 2023

sayakpaul commented on Jan 31, 2023

patil-suraj commented on Jan 31, 2023

pcuenca commented on Jan 31, 2023

sayakpaul commented on Jan 31, 2023

pcuenca commented on Jan 31, 2023

patil-suraj commented on Jan 31, 2023

pcuenca commented on Jan 31, 2023

patil-suraj commented on Jan 31, 2023

pcuenca commented on Jan 31, 2023

patil-suraj commented on Jan 31, 2023

sayakpaul commented on Jan 31, 2023

isamu-isozaki commented on Jan 31, 2023

sayakpaul commented on Mar 1, 2023

patrickvonplaten commented on Mar 7, 2023

sayakpaul commented on Mar 13, 2023

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[examples] ]Error due to mismatches in the dtype of UNet while running `train_text_to_image.py` #2163