[Feat]add cpu-offload/layerwise-offload for stable-audio-open & fix output inconsistency with same seed (#2909)

Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com>
Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>

docs/user_guide/diffusion_features.md

@@ -156,7 +156,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

@@ -57,6 +57,8 @@ Key arguments:
 - `--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.
+- `--enable-cpu-offload`: enabling model-wise offloading to save gpu memory
+- `--enable-layerwise-offload`: enabling layerwise offloading to save gpu memory
 
 ## Example materials
 

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -114,6 +114,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--enable-cpu-offload", action="store_true", help="Enable ModelWise cpu offloading to save gpu memory"
+    )
+    parser.add_argument(
+        "--enable-layerwise-offload", action="store_true", help="Enable Layerwise cpu offloading to save gpu memory"
+    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -214,6 +220,8 @@ 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)'}")
+    print(f"  ModelWise Offload: {'Enabled' if args.enable_cpu_offload else 'None'}")
+    print(f"  LayerWise Offload: {'Enabled' if args.enable_layerwise_offload else 'None'}")
     if args.use_hsdp:
         print(f"  HSDP: enabled (shard_size={args.hsdp_shard_size}, replicate_size={args.hsdp_replicate_size})")
     else:
@@ -234,6 +242,8 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        enable_cpu_offload=args.enable_cpu_offload,
+        enable_layerwise_offload=args.enable_layerwise_offload,
     )
 
     # Calculate audio end time

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -1,5 +1,6 @@
 import gc
 
+import numpy as np
 import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
@@ -10,7 +11,28 @@ from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
-models = ["riverclouds/qwen_image_random"]
+AUDIO_MODEL = {
+    "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
+}
+
+IMAGE_VIDEO_MODELS = {
+    "riverclouds/qwen_image_random": {"cuda": 2500, "rocm": 2100},
+}
+
+MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+
+AUDIO_MODEL_PARAMS = {
+    "runner_params": {},
+    "sampler_params": {},
+}
+
+IMAGE_VIDEO_MODELS_PARAMS = {
+    "runner_params": {},
+    "sampler_params": {
+        "height": 256,
+        "width": 256,
+    },
+}
 
 
 def inference(model_name: str, offload: bool = True):
@@ -19,25 +41,28 @@ def inference(model_name: str, offload: bool = True):
     device_index = current_omni_platform.current_device()
     current_omni_platform.reset_peak_memory_stats()
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
-    monitor.start()
+
+    if model_name in AUDIO_MODEL:
+        params = AUDIO_MODEL_PARAMS
+    else:
+        params = IMAGE_VIDEO_MODELS_PARAMS
+
     with OmniRunner(
         model_name,
         # TODO: we might want to add overlapped feature e2e tests
         # cache_backend="cache_dit",
         enable_cpu_offload=offload,
+        **params["runner_params"],
     ) as runner:
         current_omni_platform.reset_peak_memory_stats()
-        height = 256
-        width = 256
-
-        runner.omni.generate(
+        monitor.start()
+        output = runner.omni.generate(
             "a photo of a cat sitting on a laptop keyboard",
             OmniDiffusionSamplingParams(
-                height=height,
-                width=width,
                 num_inference_steps=9,
                 guidance_scale=0.0,
                 generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+                **params["sampler_params"],
             ),
         )
     peak = monitor.peak_used_mb
@@ -46,26 +71,51 @@ def inference(model_name: str, offload: bool = True):
     gc.collect()
     current_omni_platform.empty_cache()
 
-    return peak
+    return peak, output
+
+
+def check_audio_determinism(audio1, audio2, atol=1e-2):
+    device = current_omni_platform.device_type
+    if isinstance(audio1, np.ndarray):
+        audio1 = torch.from_numpy(audio1).to(device)
+    if isinstance(audio2, np.ndarray):
+        audio2 = torch.from_numpy(audio2).to(device)
+
+    if not torch.allclose(audio1, audio2, atol=atol):
+        diff = torch.abs(audio1 - audio2)
+        print(f"Max difference: {diff.max().item()}")
+        print(f"Mean difference: {diff.mean().item()}")
+        raise AssertionError(f"Audio outputs differ beyond tolerance atol={atol}")
+    return True
 
 
 @pytest.mark.core_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"})
-@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("model_name", list(MODELS.keys()))
 def test_cpu_offload_diffusion_model(model_name: str):
     try:
-        offload_peak_memory = inference(model_name, offload=True)
+        offload_peak_memory, output_offload = inference(model_name, offload=True)
         cleanup_dist_env_and_memory()
-        no_offload_peak_memory = inference(model_name, offload=False)
+        no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
     except Exception:
         pytest.fail("Inference failed")
     print(f"Offload peak memory: {offload_peak_memory} MB")
     print(f"No offload peak memory: {no_offload_peak_memory} MB")
+
+    if model_name == "stabilityai/stable-audio-open-1.0":
+        audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
+        audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
+        check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
+
     # Set platform-specific VRAM saving thresholds to account
     # for varying runtime memory overhead and fragmentation between CUDA and ROCm.
     is_rocm = torch.version.hip is not None
-    threshold = 2500 if not is_rocm else 2100
+    platform = "rocm" if is_rocm else "cuda"
+    threshold = MODELS[model_name][platform]
+    if threshold is None:
+        pytest.skip(f"Threshold not defined for {platform} on {model_name}")
+
     assert offload_peak_memory + threshold < no_offload_peak_memory, (
         f"Offload peak memory {offload_peak_memory} MB should be less than "
         f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB"

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -2,15 +2,31 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.e2e.offline_inference.test_diffusion_cpu_offload import check_audio_determinism
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
-# Models to test and expected saved memory in MB, correspondingly
-MODELS_SAVED_MEMORY_MB = {
-    "riverclouds/qwen_image_random": 4500,
-    # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": 45000,
+AUDIO_MODEL = {
+    "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500},
+}
+
+IMAGE_VIDEO_MODELS = {
+    "riverclouds/qwen_image_random": {"cuda": 4500, "rocm": None},
+    # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": {"cuda": 45000, "rocm": None},
+}
+
+MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+
+AUDIO_MODEL_PARAMS = {
+    "runner_params": {},
+    "sampler_params": {},
+}
+
+IMAGE_VIDEO_MODELS_PARAMS = {
+    "runner_params": {"boundary_ratio": 0.875, "flow_shift": 5.0},
+    "sampler_params": {"height": 480, "width": 640, "num_frames": 5},
 }
 
 
@@ -24,41 +40,39 @@ def run_inference(
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
 
+    if model_name in AUDIO_MODEL:
+        params = AUDIO_MODEL_PARAMS
+    else:
+        params = IMAGE_VIDEO_MODELS_PARAMS
+
     with OmniRunner(
         model_name,
         enable_layerwise_offload=layerwise_offload,
         # TODO: we might want to add overlapped feature e2e tests
         # cache_backend="cache_dit",
-        boundary_ratio=0.875,
-        flow_shift=5.0,
+        **params["runner_params"],
     ) as runner:
         current_omni_platform.reset_peak_memory_stats()
 
         # Refer to tests/e2e/offline_inference/test_t2v_model.py
         # Use minimal settings for testing
-        height = 480
-        width = 640
-        num_frames = 5
-
-        runner.omni.generate(
+        output = runner.omni.generate(
             "A cat sitting on a table",
             OmniDiffusionSamplingParams(
-                height=height,
-                width=width,
                 generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
                 guidance_scale=1.0,
                 num_inference_steps=num_inference_steps,
-                num_frames=num_frames,
+                **params["sampler_params"],
             ),
         )
 
     peak = monitor.peak_used_mb
     monitor.stop()
 
-    return peak
+    return peak, output
 
 
-@pytest.mark.parametrize("model_name", MODELS_SAVED_MEMORY_MB.keys())
+@pytest.mark.parametrize("model_name", list(MODELS.keys()))
 def test_layerwise_offload_diffusion_model(model_name: str):
     """Test that layerwise offloading reduces GPU memory usage.
 
@@ -69,11 +83,11 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     """
     try:
         # Run without layerwise offloading (baseline)
-        no_offload_peak_memory = run_inference(model_name, layerwise_offload=False)
+        no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
         cleanup_dist_env_and_memory()
 
         # Run with layerwise offloading (1 layer on device)
-        layerwise_offload_peak_memory = run_inference(model_name, layerwise_offload=True)
+        layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
         cleanup_dist_env_and_memory()
     except Exception:
         pytest.fail("Inference failed")
@@ -81,9 +95,21 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB")
     print(f"No offload peak memory: {no_offload_peak_memory} MB")
 
+    if model_name == "stabilityai/stable-audio-open-1.0":
+        audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
+        audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
+        check_audio_determinism(audio_offload, audio_no_offload, atol=1e-3)
+
+    is_rocm = torch.version.hip is not None
+    platform = "rocm" if is_rocm else "cuda"
+    expected_saved_memory = MODELS[model_name][platform]
+
+    if expected_saved_memory is None:
+        pytest.skip(f"Threshold not defined for {platform} on {model_name}")
+
     # Verify that layerwise offloading significantly reduces memory usage
     # Passes only if the actual savings exceeds the expected savings
-    assert layerwise_offload_peak_memory + MODELS_SAVED_MEMORY_MB[model_name] < no_offload_peak_memory, (
+    assert layerwise_offload_peak_memory + expected_saved_memory < no_offload_peak_memory, (
         f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
         f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB"
     )

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -28,7 +28,10 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportAudioOutput
-from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import StableAudioDiTModel
+from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import (
+    StableAudioDiTModel,
+    StableAudioSchedulerWrapper,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
@@ -134,10 +137,12 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
         self.transformer = StableAudioDiTModel(od_config=od_config, **transformer_kwargs)
 
         # Load scheduler
-        self.scheduler = CosineDPMSolverMultistepScheduler.from_pretrained(
-            model,
-            subfolder="scheduler",
-            local_files_only=local_files_only,
+        self.scheduler = StableAudioSchedulerWrapper(
+            CosineDPMSolverMultistepScheduler.from_pretrained(
+                model,
+                subfolder="scheduler",
+                local_files_only=local_files_only,
+            )
         )
 
         # Compute rotary embedding dimension
@@ -560,7 +565,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
                 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
             # Scheduler step
-            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+            latents = self.scheduler.step(noise_pred, t, latents, generator).prev_sample
 
         self._current_timestep = None
 

vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py

@@ -22,6 +22,19 @@ from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection
 logger = init_logger(__name__)
 
 
+def _preprocess_stable_audio_weight(
+    param: torch.Tensor,
+    loaded_weight: torch.Tensor,
+) -> torch.Tensor:
+    if param.shape == loaded_weight.shape:
+        return loaded_weight
+
+    if loaded_weight.ndim + 1 == param.ndim and param.shape[-1] == 1 and loaded_weight.shape == param.shape[:-1]:
+        return loaded_weight.unsqueeze(-1)
+
+    return loaded_weight
+
+
 def apply_rotary_emb_stable_audio(
     hidden_states: torch.Tensor,
     freqs_cis: tuple[torch.Tensor, torch.Tensor],
@@ -56,6 +69,84 @@ def apply_rotary_emb_stable_audio(
     return torch.cat([x_rot, x_pass], dim=-1)
 
 
+class StableAudioSchedulerWrapper:
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+
+    def __getattr__(self, name):
+        return getattr(self.scheduler, name)
+
+    def _get_step_index(self, timestep):
+        step_index = getattr(self.scheduler, "step_index", None)
+        if step_index is not None:
+            return step_index
+
+        timesteps = getattr(self.scheduler, "timesteps", None)
+        if timesteps is None:
+            return None
+
+        timestep = torch.as_tensor(
+            timestep,
+            device=timesteps.device,
+            dtype=timesteps.dtype,
+        )
+        indices = (timesteps == timestep).nonzero()
+        if len(indices) == 0:
+            return None
+
+        return indices[0].item()
+
+    def _is_final_zero_sigma_step(self, timestep):
+        step_index = self._get_step_index(timestep)
+        if step_index is None:
+            return False
+
+        sigmas = self.scheduler.sigmas
+        if step_index + 1 >= len(sigmas):
+            return False
+
+        sigma = sigmas[step_index]
+        next_sigma = sigmas[step_index + 1]
+        sigma_min = torch.as_tensor(
+            self.scheduler.config.sigma_min,
+            device=sigma.device,
+            dtype=sigma.dtype,
+        )
+
+        return torch.isclose(sigma, sigma_min) and torch.isclose(
+            next_sigma,
+            torch.zeros_like(next_sigma),
+        )
+
+    def step(self, model_output, timestep, sample, generator=None, return_dict=True):
+        use_zero_noise = self._is_final_zero_sigma_step(timestep)
+        old_noise_sampler = None
+
+        if use_zero_noise:
+            old_noise_sampler = getattr(self.scheduler, "noise_sampler", None)
+            self.scheduler.noise_sampler = _StableAudioZeroNoiseSampler(sample)
+
+        try:
+            return self.scheduler.step(
+                model_output,
+                timestep,
+                sample,
+                generator=generator,
+                return_dict=return_dict,
+            )
+        finally:
+            if use_zero_noise:
+                self.scheduler.noise_sampler = old_noise_sampler
+
+
+class _StableAudioZeroNoiseSampler:
+    def __init__(self, sample: torch.Tensor):
+        self.sample = sample
+
+    def __call__(self, sigma, next_sigma):
+        return torch.zeros_like(self.sample)
+
+
 class StableAudioGaussianFourierProjection(GaussianFourierProjection):
     """Gaussian Fourier embeddings for noise levels.
 
@@ -374,6 +465,7 @@ class StableAudioDiTModel(nn.Module):
     """
 
     _repeated_blocks = ["StableAudioDiTBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]
 
     def __init__(
@@ -594,8 +686,12 @@ class StableAudioDiTModel(nn.Module):
 
             if mapped_name in params_dict:
                 param = params_dict[mapped_name]
+                loaded_weight = _preprocess_stable_audio_weight(param, loaded_weight)
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
+                try:
+                    weight_loader(param, loaded_weight)
+                except AssertionError as err:
+                    raise AssertionError(f"Failed to load Stable Audio weight {name!r} as {mapped_name!r}") from err
                 loaded_params.add(mapped_name)
             else:
                 logger.debug(f"Skipping weight {name} - not found in model")