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

.buildkite/test-nightly.yml

@@ -1279,7 +1279,7 @@ steps:
       - label: ":full_moon: Diffusion X2V · Function Test"
         timeout_in_minutes: 90
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
         agents:
           queue: "mithril-h100-pool"
         plugins:

tests/e2e/online_serving/test_ltx2_expansion.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Expansion tests for LTX-2 two-stage pipelines.
+
+Coverage:
+- HSDP with LTX2TwoStagesPipeline (text-to-video)
+- HSDP with LTX2ImageToVideoTwoStagesPipeline (image-to-video)
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
+PROMPT = "A serene lake at sunset with mountains in the background."
+PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
+
+
+def _cases():
+    cases = []
+
+    # T2V: LTX2TwoStagesPipeline
+    cases.append(
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    *HSDP_ARGS,
+                    "--model-class-name",
+                    "LTX2TwoStagesPipeline",
+                ],
+            ),
+            id="t2v_hsdp",
+            marks=PARALLEL_MARKS,
+        )
+    )
+
+    # I2V: LTX2ImageToVideoTwoStagesPipeline
+    cases.append(
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    *HSDP_ARGS,
+                    "--model-class-name",
+                    "LTX2ImageToVideoTwoStagesPipeline",
+                ],
+            ),
+            id="i2v_hsdp",
+            marks=PARALLEL_MARKS,
+        )
+    )
+
+    return cases
+
+
+@pytest.mark.parametrize("omni_server", _cases(), indirect=True)
+def test_ltx2_two_stage_hsdp(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args)
+
+    # The two-stage pipeline generates at the requested resolution then 2x
+    # upsamples via LTX2LatentUpsamplerModel, so output dimensions differ
+    # from the request.  Omit height/width from form_data so the assertion
+    # helper skips the dimension check (it only asserts when the key is
+    # present).  The pipeline falls back to its own defaults.
+    form_data = {
+        "prompt": PROMPT,
+        "model": omni_server.model,
+        "num_frames": 9,
+        "fps": 8,
+        "num_inference_steps": 2,
+        "guidance_scale": 1.0,
+        "seed": 42,
+    }
+
+    request_config = {
+        "model": omni_server.model,
+        "form_data": form_data,
+    }
+
+    if is_i2v:
+        request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+
+    openai_client.send_video_diffusion_request(request_config)

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)