ValueError: Could not find the transformer layer class QwenImageTransformerBlock in the model.

[garychan22]

Hi, I am trying to fine-tuning qwen-image-edit using accelerate in FSDP mode. I want to warp the QwenImageTransformerBlock in transformer and Qwen2_5_VLVisionBlock,Qwen2_5_VLDecoderLayer in text_encoder. I set the environment param

def set_fsdp_env():
    os.environ["ACCELERATE_USE_FSDP"] = 'true'
    os.environ["FSDP_AUTO_WRAP_POLICY"] = 'TRANSFORMER_BASED_WRAP'
    os.environ["FSDP_BACKWARD_PREFETCH"] = 'BACKWARD_PRE'
    os.environ["FSDP_TRANSFORMER_CLS_TO_WRAP"] = 'QwenImageTransformerBlock,Qwen2_5_VLVisionBlock,Qwen2_5_VLDecoderLayer'
    os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = 'false'

and prepare the two models

transformer = accelerator.prepare(transformer)
text_encoder = accelerator.prepare(text_encoder)

Finally, I encountered the error raised from text_encoder = accelerator.prepare(text_encoder)

ValueError: Could not find the transformer layer class QwenImageTransformerBlock in the model.

How can I resolve this problem? Thanks!

[SunMarc]

Hey @garychan22, our fsdp integration doesn't handle well multi-models. I recommend training only model at a time for now. Or maybe you can try to fix how the auto wrap policy works and see if this fixes it. The issue right now is that the text encoder doesn't have QwenImageTransformerBlock, which triggers the error.

[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.