Transformer2DModel does not support gradient checkpointing.

[lawrence-cj]

Describe the bug

Transformer2DModel does not support gradient checkpointing.

Reproduction

transformer = Transformer2DModel.from_pretrained(args.pretrained_teacher_model, subfolder="transformer") transformer.enable_gradient_checkpointing()

Logs

Traceback (most recent call last):
  File "/home/xieenze/chenjunsong/code/diffusion/Pixart/train_scripts/train_lcm_lora.py", line 1179, in <module>
    main(args)
  File "/home/xieenze/chenjunsong/code/diffusion/Pixart/train_scripts/train_lcm_lora.py", line 803, in main
    student_transformer.enable_gradient_checkpointing()
  File "/home/xieenze/chenjunsong/code/diffusion/diffusers/src/diffusers/models/modeling_utils.py", line 229, in enable_gradient_checkpointing
    raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
ValueError: Transformer2DModel does not support gradient checkpointing.

System Info

diffusers >= 0.23.0

Who can help?

@DN6 @yiyixuxu @sayakpaul @DN6 @patrickvonplaten

[patrickvonplaten]

Yes it looks like we will need to add it

[Ghy1209]

Hi. After solving this problem. did you meet the error:RuntimeError: output with shape [8, 16, 4096, 1] doesn't match the broadcast shape [8, 16, 4096, 121]. How did you solve it? Thank you!

[Ghy1209]

Describe the bug

Transformer2DModel does not support gradient checkpointing.

Reproduction

transformer = Transformer2DModel.from_pretrained(args.pretrained_teacher_model, subfolder="transformer") transformer.enable_gradient_checkpointing()

Logs

Traceback (most recent call last):
  File "/home/xieenze/chenjunsong/code/diffusion/Pixart/train_scripts/train_lcm_lora.py", line 1179, in <module>
    main(args)
  File "/home/xieenze/chenjunsong/code/diffusion/Pixart/train_scripts/train_lcm_lora.py", line 803, in main
    student_transformer.enable_gradient_checkpointing()
  File "/home/xieenze/chenjunsong/code/diffusion/diffusers/src/diffusers/models/modeling_utils.py", line 229, in enable_gradient_checkpointing
    raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
ValueError: Transformer2DModel does not support gradient checkpointing.

System Info

diffusers >= 0.23.0

Who can help?

@DN6 @yiyixuxu @sayakpaul @DN6 @patrickvonplaten

Hi. After solving this problem. did you meet the error:RuntimeError: output with shape [8, 16, 4096, 1] doesn't match the broadcast shape [8, 16, 4096, 121]. How did you solve it? Thank you!