From 6646a66d416aac49dfd7259343b3135fa61c3fb7 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Mon, 22 Jun 2026 22:54:14 -0400 Subject: [PATCH] [BugFix] Use ModuleDiscovery for HSDP DiT sharding to support nested pipelines (#3456) Signed-off-by: Nick Cao Co-authored-by: Claude --- .buildkite/test-nightly.yml | 2 +- .../e2e/online_serving/test_ltx2_expansion.py | 95 +++++++++++++++++++ .../model_loader/diffusers_loader.py | 35 +++---- 3 files changed, 111 insertions(+), 21 deletions(-) create mode 100644 tests/e2e/online_serving/test_ltx2_expansion.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index dd6d5ec31..809e3feef 100644 --- a/.buildkite/test-nightly.yml +++ b/.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: diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py new file mode 100644 index 000000000..82855f62f --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 1aa40aae5..f070dbc3c 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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)