[Model] Add HSDP support for Stable-Audio-Open (#2982)

Signed-off-by: hanzheli <hanzheli@kuaishou.com>

docs/user_guide/diffusion_features.md

@@ -154,7 +154,7 @@ The following tables show which models support each feature:
 
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
 
 
 ## Feature Compatibility

docs/user_guide/examples/offline_inference/text_to_audio.md

@@ -29,6 +29,22 @@ python text_to_audio.py \
   --output stable_audio_output.wav
 ```
 
+To reduce per-GPU memory for multi-GPU inference, launch with HSDP:
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative-prompt "Low quality" \
+  --seed 42 \
+  --guidance-scale 7.0 \
+  --audio-length 10.0 \
+  --num-inference-steps 100 \
+  --use-hsdp \
+  --hsdp-shard-size 2 \
+  --output stable_audio_output.wav
+```
+
 Key arguments:
 
 - `--prompt`: text description (string).
@@ -37,6 +53,9 @@ Key arguments:
 - `--guidance-scale`: classifier-free guidance scale.
 - `--audio-length`: audio duration in seconds.
 - `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT.
+- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
+- `--hsdp-replicate-size`: number of HSDP replica groups.
 - `--output`: path to save the generated WAV file.
 
 ## Example materials

examples/offline_inference/text_to_audio/README.md

@@ -27,6 +27,22 @@ python text_to_audio.py \
   --output stable_audio_output.wav
 ```
 
+To reduce per-GPU memory for multi-GPU inference, launch with HSDP:
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative-prompt "Low quality" \
+  --seed 42 \
+  --guidance-scale 7.0 \
+  --audio-length 10.0 \
+  --num-inference-steps 100 \
+  --use-hsdp \
+  --hsdp-shard-size 2 \
+  --output stable_audio_output.wav
+```
+
 Key arguments:
 
 - `--prompt`: text description (string).
@@ -35,5 +51,8 @@ Key arguments:
 - `--guidance-scale`: classifier-free guidance scale.
 - `--audio-length`: audio duration in seconds.
 - `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT.
+- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
+- `--hsdp-replicate-size`: number of HSDP replica groups.
 - `--cache-backend`: cache acceleration backend. Stable Audio currently supports `tea_cache`.
 - `--output`: path to save the generated WAV file.

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -21,6 +21,7 @@ from pathlib import Path
 import numpy as np
 import torch
 
+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -113,6 +114,23 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--use-hsdp",
+        action="store_true",
+        help="Enable HSDP for Stable Audio DiT weight sharding.",
+    )
+    parser.add_argument(
+        "--hsdp-shard-size",
+        type=int,
+        default=1,
+        help="Number of GPUs to shard Stable Audio DiT weights across when HSDP is enabled.",
+    )
+    parser.add_argument(
+        "--hsdp-replicate-size",
+        type=int,
+        default=1,
+        help="Number of HSDP replica groups. Default 1 means pure sharding.",
+    )
     return parser.parse_args()
 
 
@@ -158,12 +176,23 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Guidance scale: {args.guidance_scale}")
     print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
+    if args.use_hsdp:
+        print(f"  HSDP: enabled (shard_size={args.hsdp_shard_size}, replicate_size={args.hsdp_replicate_size})")
+    else:
+        print("  HSDP: disabled")
     print(f"  Seed: {args.seed}")
     print(f"{'=' * 60}\n")
 
+    parallel_config = DiffusionParallelConfig(
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
+    )
+
     # Initialize Omni with Stable Audio model
     omni = Omni(
         model=args.model,
+        parallel_config=parallel_config,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,

tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py

@@ -0,0 +1,25 @@
+import pytest
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import StableAudioDiTModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_stable_audio_exposes_hsdp_shard_conditions_for_transformer_blocks():
+    model = object.__new__(StableAudioDiTModel)
+    nn.Module.__init__(model)
+    model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.proj_out = nn.Linear(4, 4)
+
+    conditions = getattr(model, "_hsdp_shard_conditions", None)
+
+    assert conditions is not None
+    assert len(conditions) == 1
+
+    matched = []
+    for name, module in model.named_modules():
+        if any(cond(name, module) for cond in conditions):
+            matched.append(name)
+
+    assert matched == ["transformer_blocks.0", "transformer_blocks.1"]

vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py

@@ -17,6 +17,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 
 logger = init_logger(__name__)
 
@@ -376,6 +377,7 @@ class StableAudioDiTModel(nn.Module):
     """
 
     _repeated_blocks = ["StableAudioDiTBlock"]
+    _hsdp_shard_conditions = [is_transformer_block_module]
 
     def __init__(
         self,