[BugFix] Use ModuleDiscovery for HSDP DiT sharding to support nested pipelines (#3456)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -653,6 +653,13 @@ class DiffusersPipelineLoader:
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
+        # Discover pipeline components (DiT, encoders, VAEs) via
+        # ModuleDiscovery, which consults SupportsComponentDiscovery
+        # when available and falls back to well-known attribute names.
+        # This supports nested pipelines (e.g. LTX2TwoStagesPipeline
+        # where the transformer lives at "pipe.transformer").
+        discovered_modules = ModuleDiscovery.discover(model)
+
         # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
         # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
         # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
@@ -664,31 +671,19 @@ class DiffusersPipelineLoader:
                 prepare_fp8_layers_for_fsdp,
             )
 
-            for _attr in ("transformer", "transformer_2"):
-                _trans = getattr(model, _attr, None)
-                if _trans is not None:
-                    prepare_fp8_layers_for_fsdp(_trans)
+            for trans in discovered_modules.dits:
+                prepare_fp8_layers_for_fsdp(trans)
 
-        # Collect all transformers to shard (some models have transformer_2 for MoE)
-        transformers_to_shard = []
-        transformer = getattr(model, "transformer", None)
-        if transformer is None:
-            raise ValueError("Model has no transformer attribute for HSDP")
-        transformers_to_shard.append(("transformer", transformer))
+        if not discovered_modules.dits:
+            raise ValueError("No DiT modules discovered for HSDP sharding")
 
-        # Check for transformer_2 (MoE two-stage models like Wan2.2-I2V)
-        transformer_2 = getattr(model, "transformer_2", None)
-        if transformer_2 is not None:
-            transformers_to_shard.append(("transformer_2", transformer_2))
-
-        # Apply HSDP sharding to all transformers
-        for name, trans in transformers_to_shard:
+        # Apply HSDP sharding to all discovered DiT transformers
+        for name, trans in zip(discovered_modules.dit_names, discovered_modules.dits):
             logger.debug("Applying HSDP to %s", name)
             apply_hsdp_to_model(trans, hsdp_config, target_device=target_device)
 
-        # # HSDP only shards transformer modules. All other runtime modules must
-        # # be placed on the execution device explicitly after sharding.
-        discovered_modules = ModuleDiscovery.discover(model)
+        # HSDP only shards transformer modules. All other runtime modules must
+        # be placed on the execution device explicitly after sharding.
         modules_to_move: list[nn.Module] = []
         if discovered_modules.vaes is not None:
             modules_to_move.extend(discovered_modules.vaes)