From a0918ce583985ce597d748fa223c0686204a1f5e Mon Sep 17 00:00:00 2001 From: boatman <109857087+sphinxkkkbc@users.noreply.github.com> Date: Tue, 5 May 2026 22:49:57 +0800 Subject: [PATCH] [Feat]add cpu-offload/layerwise-offload for stable-audio-open & fix output inconsistency with same seed (#2909) Signed-off-by: sphinxkkkbc Co-authored-by: sphinxkkkbc --- docs/user_guide/diffusion_features.md | 2 +- .../offline_inference/text_to_audio.md | 2 + .../text_to_audio/text_to_audio.py | 10 ++ .../test_diffusion_cpu_offload.py | 76 +++++++++++--- .../test_diffusion_layerwise_offload.py | 64 ++++++++---- .../stable_audio/pipeline_stable_audio.py | 17 ++-- .../stable_audio/stable_audio_transformer.py | 98 ++++++++++++++++++- 7 files changed, 229 insertions(+), 40 deletions(-) diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index e249d2e23..824dd9a0b 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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 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 a31a4d7a4..246314367 100644 --- a/docs/user_guide/examples/offline_inference/text_to_audio.md +++ b/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 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 2a1613e5e..3072cc06f 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/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 diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index d7fd6f72f..e7a8a44cb 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/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" diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index 4f19c1004..658fb09c3 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/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" ) diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index f582c4bc6..d6b4f9e93 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/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 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 7422f320f..83a669d92 100644 --- a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py +++ b/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")