enable FSDP for qwen vl

[akaitsuki-ii]

1. 更改FSDP的wrap策略，按照指定的模块类进行wrap；
2. 启用FSDP时也应用于Qwen VL (后面可以加上对其他模型的text encoder的支持)；
3. 修复启用FSDP前如果模型已经在GPU上初始化，会占用额外显存的问题

[gemini-code-assist]

Summary of Changes

Hello @akaitsuki-ii, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

此拉取请求主要重构了 FSDP（Fully Sharded Data Parallel）的模块包装机制，使其能够通过指定模块类进行更精确的控制。它为 Qwen VL 模型添加了 FSDP 支持，并引入了一项重要的显存优化，确保在启用 FSDP 时模型首先在 CPU 上初始化，避免了不必要的 GPU 显存占用。这些改进旨在提升分布式训练的效率和资源利用率。

Highlights

• FSDP 包装策略更新: 将 FSDP 的自动包装策略从基于 lambda 函数的名称匹配更改为基于指定模块类的 transformer_auto_wrap_policy，提高了策略的精确性和类型安全性。
• Qwen VL 模型支持 FSDP: 为 Qwen2_5_VLForConditionalGeneration 模型启用了 FSDP 支持，通过实现 get_fsdp_module_cls 方法，明确指定了需要进行 FSDP 包装的模块类（如 Qwen2_5_VisionBlock 和 Qwen2_5_VLDecoderLayer）。
• FSDP 显存优化: 修复了在启用 FSDP 前模型可能在 GPU 上重复初始化导致额外显存占用的问题。现在，如果启用了 FSDP，模型会在 CPU 上初始化，从而节省 GPU 显存。
• 统一 FSDP 模块类定义: 多个模型（如 FluxDiT、QwenImageDiT、WanS2VDiT）的 FSDP 模块获取方法已更新并重命名为 get_fsdp_module_cls，以返回需要包装的模块类集合，与新的 FSDP 策略保持一致。

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

本次 PR 的核心变更是重构了 FSDP 的封装策略，从原先基于模块名称的自定义封装方式，改为使用 transformer_auto_wrap_policy 的、基于模块类的标准封装方式。这是一个很好的改进，提升了代码质量和可维护性。此外，本次变更也为 Qwen VL 模型启用了 FSDP，并通过确保模型在 FSDP 启用时先在 CPU 上初始化，修复了额外占用 GPU 显存的问题。这些改动在多个模型和 pipeline 中都得到了正确的实现。

不过，我在 diffsynth_engine/pipelines/qwen_image.py 文件中发现一处潜在的逻辑错误，其中为 encoder 设置设备的代码在启用 FSDP 时似乎不正确。