diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index be5b2d10c..6e4a54a65 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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 diff --git a/docs/user_guide/examples/offline_inference/text_to_audio.md b/docs/user_guide/examples/offline_inference/text_to_audio.md index 62a70e525..a31a4d7a4 100644 --- a/docs/user_guide/examples/offline_inference/text_to_audio.md +++ b/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 diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md index 50bab3e2f..10a8ae37e 100644 --- a/examples/offline_inference/text_to_audio/README.md +++ b/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. diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 3adb3ad53..53f3daf8d 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/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, diff --git a/tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py b/tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py new file mode 100644 index 000000000..923b9a863 --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py index 4a4892673..a3d4dc517 100644 --- a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py +++ b/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,