RuntimeError: torch._dynamo.optimize is called on a non function object.

[xiaosaxuexige]

Describe the bug

When I am trying to run the example given by the authors, which is on the pokemon dataset. I just follows the instruction on huggingface and never revise the code. but I got this error:

RuntimeError:

torch._dynamo.optimize is called on a non function object.
If this is a callable class, please wrap the relevant code into a function and optimize the
wrapper function.

Reproduction

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=15000
--learning_rate=1e-05
--max_grad_norm=1
--lr_scheduler="constant" --lr_warmup_steps=0
--output_dir="sd-pokemon-model"

Logs

03/22/2023 14:32:42 - INFO - __main__ - ***** Running training *****
03/22/2023 14:32:42 - INFO - __main__ -   Num examples = 833
03/22/2023 14:32:42 - INFO - __main__ -   Num Epochs = 72
03/22/2023 14:32:42 - INFO - __main__ -   Instantaneous batch size per device = 1
03/22/2023 14:32:42 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 4
03/22/2023 14:32:42 - INFO - __main__ -   Gradient Accumulation steps = 4
03/22/2023 14:32:42 - INFO - __main__ -   Total optimization steps = 15000
Steps:   0%|                                                                                                                                  | 0/15000 [00:00<?, ?it/s]Traceback (most recent call last):
  File "train_text_to_image.py", line 789, in <module>
    main()
  File "train_text_to_image.py", line 730, in main
    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
  File "/home/wangfuling/wangfuling/ENTER/envs/dfs/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wangfuling/wangfuling/ENTER/envs/dfs/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/home/wangfuling/wangfuling/ENTER/envs/dfs/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 229, in __call__
    raise RuntimeError(
RuntimeError: 

torch._dynamo.optimize is called on a non function object.
If this is a callable class, please wrap the relevant code into a function and optimize the
wrapper function.

>> class CallableClass:
>>     def __init__(self):
>>         super().__init__()
>>         self.relu = torch.nn.ReLU()
>>
>>     def __call__(self, x):
>>         return self.relu(torch.sin(x))
>>
>>     def print_hello(self):
>>         print("Hello world")
>>
>> mod = CallableClass()

If you want to optimize the __call__ function and other code, wrap that up in a function

>> def wrapper_fn(x):
>>     y = mod(x)
>>     return y.sum()

and then optimize the wrapper_fn

>> opt_wrapper_fn = torch._dynamo.optimize(wrapper_fn)

System Info

• diffusers version: 0.15.0.dev0
• Platform: Linux-3.10.0-1160.66.1.el7.x86_64-x86_64-with-glibc2.10
• Python version: 3.8.0
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.13.2
• Transformers version: 4.27.1
• Accelerate version: 0.17.1
• xFormers version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[patrickvonplaten]

@sayakpaul do you maybe have some time to look into this?

[sayakpaul]

@muellerzr @pacman100 does this error sound familiar? If so, do you have any recommendations?

[muellerzr]

CC @sgugger

[sgugger]

I haven't tried torch.compile on unets, so nothing to contribute here (except that this issue has been opened in Transformers once too)

[sayakpaul]

Our training scripts do not use this torch.compile(). On the other hand, accelerate doesn't set it by default either. So, on the surface, it does seem a bit weird to me.

[patrickvonplaten]

Can we try to reproduce the error somehow? @xiaosaxuexige do you think you could try to reproduce the error in a google colab maybe?

[zhongshsh]

I had the same problem.

[rongxiaoqu]

same problem...

[sayakpaul]

Could we get a minimally reproducible code snippet (preferably a Google Colab)?

[Happenmass]

same problem too......

[patrickvonplaten]

@sayakpaul when running this:

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=15000
--learning_rate=1e-05
--max_grad_norm=1
--lr_scheduler="constant" --lr_warmup_steps=0
--output_dir="sd-pokemon-model"

with Torch 2.0 on GPU can you reproduce the error?

[fecet]

I have same problem with:

import torch
from accelerate import Accelerator

def prepare_model(model):
    accelerator = Accelerator()
    model = accelerator.prepare(model)
    return model

model = torch.nn.Linear(10, 2)

model = prepare_model(model)

inputs = torch.randn(4, 10)
outputs = model(inputs)
print(outputs.shape)

maybe it's a problem caused by accelerate/torch?

Edit: remove my accelerate config resolved this. I guess there is something wrong with accelerate's compile behaviour

[sayakpaul]

Pinging @muellerzr one more time (sorry) to check if this is a known issue.

[sgugger]

What is your accelerate configuration? And what is your setup? This script doesn't work on my side since the inputs are on the CPU but the model is moved on the GPU by Accelerate. Putting the inputs on the GPU makes it work without any issue.