enable FSDP for qwen vl (#184)

* enable FSDP for qwen vl

* fix

Author: akaitsuki-ii

diffsynth_engine/models/base.py

@@ -57,7 +57,7 @@ def unload_loras(self):
     def get_tp_plan(self):
         raise NotImplementedError(f"{self.__class__.__name__} does not support TP")
 
-    def get_fsdp_modules(self):
+    def get_fsdp_module_cls(self):
         raise NotImplementedError(f"{self.__class__.__name__} does not support FSDP")
 
 

diffsynth_engine/models/flux/flux_dit.py

@@ -515,5 +515,5 @@ def compile_repeated_blocks(self, *args, **kwargs):
         for block in self.blocks:
             block.compile(*args, **kwargs)
 
-    def get_fsdp_modules(self):
-        return ["blocks", "single_blocks"]
+    def get_fsdp_module_cls(self):
+        return {FluxDoubleTransformerBlock, FluxSingleTransformerBlock}

diffsynth_engine/models/qwen_image/qwen2_5_vl.py

@@ -942,6 +942,8 @@ def _unmask_unattended(
 
 
 class Qwen2_5_VLForConditionalGeneration(PreTrainedModel):
+    _supports_parallelization = True
+
     def __init__(
         self,
         vision_config: Qwen2_5_VLVisionConfig,
@@ -1173,6 +1175,9 @@ def get_rope_index(
 
             return position_ids, mrope_position_deltas
 
+    def get_fsdp_module_cls(self):
+        return {Qwen2_5_VisionBlock, Qwen2_5_VLDecoderLayer}
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,

diffsynth_engine/models/qwen_image/qwen_image_dit.py

@@ -544,5 +544,5 @@ def compile_repeated_blocks(self, *args, **kwargs):
         for block in self.transformer_blocks:
             block.compile(*args, **kwargs)
 
-    def get_fsdp_modules(self):
-        return ["transformer_blocks"]
+    def get_fsdp_module_cls(self):
+        return {QwenImageTransformerBlock}

diffsynth_engine/models/wan/wan_audio_encoder.py

@@ -223,7 +223,6 @@ def convert(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor
 
 class Wav2Vec2Model(PreTrainedModel):
     converter = Wav2Vec2StateDictConverter()
-    _supports_parallelization = False
 
     def __init__(self, config: Wav2Vec2Config, device: str = "cuda:0", dtype: torch.dtype = torch.bfloat16):
         super().__init__()

diffsynth_engine/models/wan/wan_dit.py

@@ -502,5 +502,5 @@ def compile_repeated_blocks(self, *args, **kwargs):
         for block in self.single_blocks:
             block.compile(*args, **kwargs)
 
-    def get_fsdp_modules(self):
-        return ["blocks"]
+    def get_fsdp_module_cls(self):
+        return {DiTBlock}

diffsynth_engine/pipelines/flux_image.py

@@ -143,7 +143,7 @@ def _from_kohya(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dic
                 layer_id, layer_type = name.split("_", 1)
                 layer_type = layer_type.replace("self_attn_", "self_attn.").replace("mlp_", "mlp.")
                 rename = ".".join(["encoders", layer_id, clip_attn_rename_dict[layer_type]])
-                
+
                 lora_args = {}
                 lora_args["alpha"] = param
                 lora_args["up"] = lora_state_dict[origin_key.replace(".alpha", ".lora_up.weight")]
@@ -517,7 +517,7 @@ def _from_state_dict(cls, state_dicts: FluxStateDicts, config: FluxPipelineConfi
             if config.use_fbcache:
                 dit = FluxDiTFBCache.from_state_dict(
                     state_dicts.model,
-                    device=init_device,
+                    device=("cpu" if config.use_fsdp else init_device),
                     dtype=config.model_dtype,
                     in_channel=config.control_type.get_in_channel(),
                     attn_kwargs=attn_kwargs,
@@ -526,7 +526,7 @@ def _from_state_dict(cls, state_dicts: FluxStateDicts, config: FluxPipelineConfi
             else:
                 dit = FluxDiT.from_state_dict(
                     state_dicts.model,
-                    device=init_device,
+                    device=("cpu" if config.use_fsdp else init_device),
                     dtype=config.model_dtype,
                     in_channel=config.control_type.get_in_channel(),
                     attn_kwargs=attn_kwargs,

diffsynth_engine/pipelines/qwen_image.py

@@ -237,7 +237,7 @@ def _from_state_dict(cls, state_dicts: QwenImageStateDicts, config: QwenImagePip
             state_dicts.encoder,
             vision_config=vision_config,
             config=text_config,
-            device=init_device,
+            device=("cpu" if config.use_fsdp else init_device),
             dtype=config.encoder_dtype,
         )
         with open(QWEN_IMAGE_VAE_CONFIG_FILE, "r", encoding="utf-8") as f:
@@ -257,15 +257,15 @@ def _from_state_dict(cls, state_dicts: QwenImageStateDicts, config: QwenImagePip
             if config.use_fbcache:
                 dit = QwenImageDiTFBCache.from_state_dict(
                     state_dicts.model,
-                    device=init_device,
+                    device=("cpu" if config.use_fsdp else init_device),
                     dtype=config.model_dtype,
                     attn_kwargs=attn_kwargs,
                     relative_l1_threshold=config.fbcache_relative_l1_threshold,
                 )
             else:
                 dit = QwenImageDiT.from_state_dict(
                     state_dicts.model,
-                    device=init_device,
+                    device=("cpu" if config.use_fsdp else init_device),
                     dtype=config.model_dtype,
                     attn_kwargs=attn_kwargs,
                 )

diffsynth_engine/pipelines/sdxl_image.py

@@ -181,7 +181,7 @@ def from_pretrained(cls, model_path_or_config: SDXLPipelineConfig) -> "SDXLImage
 
     @classmethod
     def from_state_dict(cls, state_dicts: SDXLStateDicts, config: SDXLPipelineConfig) -> "SDXLImagePipeline":
-        init_device = "cpu" if config.offload_mode else config.device
+        init_device = "cpu" if config.offload_mode is not None else config.device
         tokenizer = CLIPTokenizer.from_pretrained(SDXL_TOKENIZER_CONF_PATH)
         tokenizer_2 = CLIPTokenizer.from_pretrained(SDXL_TOKENIZER_2_CONF_PATH)
         with LoRAContext():

diffsynth_engine/pipelines/wan_s2v.py

@@ -664,7 +664,7 @@ def _from_state_dict(
             dit = WanS2VDiT.from_state_dict(
                 state_dicts.model,
                 config=model_config,
-                device=init_device,
+                device=("cpu" if config.use_fsdp else init_device),
                 dtype=config.model_dtype,
                 attn_kwargs=attn_kwargs,
             )