From 260db8bab5443f52b38b8eb813899a076c18d54c Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Fri, 21 Aug 2026 18:30:53 +0800 Subject: [PATCH] [Bugfix][Diffusion] Fix MiniMax-H3 model-level CPU offload residency (#6072) Signed-off-by: Xiaodong Ye --- .../minimax_h3/test_minimax_h3_contract.py | 7 +- .../minimax_h3/test_minimax_h3_offload.py | 308 ++++++++++++++++++ .../offloader/test_sequential_backend.py | 61 +++- .../models/minimax_h3/pipeline_minimax_h3.py | 193 ++++++----- vllm_omni/diffusion/offloader/__init__.py | 5 +- vllm_omni/diffusion/offloader/base.py | 21 ++ .../diffusion/offloader/sequential_backend.py | 60 +++- 7 files changed, 553 insertions(+), 102 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 34250abf2..0f5a11279 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -1038,6 +1038,7 @@ def test_model_offload_uses_hooked_text_encoder_call(): ) expected = torch.ones(2, 3) pipeline.text_encoder = Mock(return_value=expected) + pipeline._model_cpu_offload_modules = [pipeline.text_encoder] input_ids = torch.tensor([1, 2]) vision_kwargs = {"pixel_values": torch.ones(1, 4)} @@ -1401,7 +1402,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch): lambda path: f"frames:{path}", ) - rows, shapes = pipeline._encode_video_conditions(None, count=2) + rows, shapes = pipeline._encode_video_conditions_resident(None, count=2) assert pipeline.video_vae.calls == [ "frames:video-1.mp4", @@ -1964,7 +1965,7 @@ def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch): waveform = torch.zeros(1, 10) with pytest.raises(ValueError, match="max_duration_seconds must be positive"): - pipeline._encode_audio_conditions( + pipeline._encode_audio_conditions_resident( [(waveform, 10)], max_duration_seconds=0, ) @@ -1985,7 +1986,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration(): pipeline.audio_vae = FakeAudioVAE() waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40) - encoded, lengths = pipeline._encode_audio_conditions( + encoded, lengths = pipeline._encode_audio_conditions_resident( [(waveform, 10)], max_duration_seconds=2.5, ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py new file mode 100644 index 000000000..c993de492 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from PIL import Image + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_h3_profiler_targets_cover_aggregate_reference_encoders(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + targets = MiniMaxH3Pipeline._PROFILER_TARGETS + assert "_encode_visual_conditions" in targets + assert "_encode_reference_audio_conditions" in targets + assert "_encode_video_conditions" not in targets + assert "_encode_video_audio_conditions" not in targets + assert "_encode_audio_conditions" not in targets + + +@pytest.mark.parametrize("distributed", [False, True]) +def test_manual_component_offload_wins_over_requested_model_offload(distributed): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=not distributed, + enable_distributed_layerwise_offload=distributed, + ) + pipeline._model_cpu_offload_modules = [] + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + component = Mock() + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + with pipeline._component_on_device(component): + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_not_called() + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + component.offload_to_cpu.assert_called_once_with() + + +def test_requested_model_offload_without_backend_uses_resident_components(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._model_cpu_offload_modules = [] + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.encode_ids.assert_called_once() + + +def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = torch.nn.Linear(2, 2) + pipeline.transformers_ref = torch.nn.Linear(2, 2) + pipeline.text_encoder = torch.nn.Linear(2, 2) + pipeline.video_vae = torch.nn.Linear(2, 2) + pipeline.audio_vae = torch.nn.Linear(2, 2) + apply_offload = Mock() + remove_offload = Mock() + monkeypatch.setattr(module, "apply_sequential_offload", apply_offload) + monkeypatch.setattr(module, "remove_sequential_offload", remove_offload) + + pipeline.enable_omni_model_cpu_offload( + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + ) + + dits = [pipeline.transformer, pipeline.transformers_ref] + stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae] + apply_offload.assert_called_once_with( + dit_modules=dits, + encoder_modules=stages, + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + offload_initial_dits=True, + ) + + pipeline.disable_omni_model_cpu_offload() + + remove_offload.assert_called_once_with([*dits, *stages]) + + +@pytest.mark.parametrize("decode_fails", [False, True]) +def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + component = torch.nn.Linear(2, 2) + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + monkeypatch.setattr(module, "sequential_offload_component", record_component) + pipeline._model_cpu_offload_modules = [component] + + def decode(): + with pipeline._component_on_device(component): + events.append(("decode", component)) + if decode_fails: + raise RuntimeError("decode failed") + + if decode_fails: + with pytest.raises(RuntimeError, match="decode failed"): + decode() + else: + decode() + + assert events == [ + ("activate", component), + ("decode", component), + ("offload", component), + ] + + +@pytest.mark.parametrize("encode_fails", [False, True]) +def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.video_vae = Mock() + pipeline._model_cpu_offload_modules = [pipeline.video_vae] + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + def encode_image(image): + events.append(("encode", image.size)) + if encode_fails and image.width == 32: + raise RuntimeError("encode failed") + return torch.ones(image.width // 16, 4) + + pipeline.video_vae.encode_image.side_effect = encode_image + monkeypatch.setattr(module, "sequential_offload_component", record_component) + monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1)) + monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value) + images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))] + + if encode_fails: + with pytest.raises(RuntimeError, match="encode failed"): + pipeline._encode_visual_conditions(images, None, video_count=0) + else: + rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0) + assert rows.shape == (3, 4) + assert shapes == [(1, 1, 1), (1, 1, 2)] + + assert events == [ + ("activate", pipeline.video_vae), + ("encode", (16, 16)), + ("encode", (32, 16)), + ("offload", pipeline.video_vae), + ] + + +def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.video_vae = Mock() + pipeline.video_vae.is_distributed_enabled.return_value = False + pipeline._model_cpu_offload_modules = [pipeline.video_vae] + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + def encode_image(image): + events.append(("image", image.size)) + return torch.ones(1, 4) + + def encode_video(frames): + events.append(("video", frames)) + return torch.ones(3, 4), (3, 2, 2) + + pipeline.video_vae.encode_image.side_effect = encode_image + pipeline.video_vae.encode_video.side_effect = encode_video + monkeypatch.setattr(module, "sequential_offload_component", record_component) + monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1)) + monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value) + monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}") + images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))] + prepared_videos = [{"prepared_path": "reference.mp4"}] + + rows, shapes = pipeline._encode_visual_conditions( + images, + prepared_videos, + video_count=1, + ) + + assert rows.shape == (5, 4) + assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)] + assert events == [ + ("activate", pipeline.video_vae), + ("image", (16, 16)), + ("image", (32, 16)), + ("video", "frames:reference.mp4"), + ("offload", pipeline.video_vae), + ] + + +@pytest.mark.parametrize("encode_fails", [False, True]) +def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeypatch, encode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = Mock() + pipeline._model_cpu_offload_modules = [pipeline.audio_vae] + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + def encode_embedded(*args, **kwargs): + events.append(("embedded",)) + if encode_fails: + raise RuntimeError("audio encode failed") + return torch.ones(1, 2), [1] + + def encode_standalone(*args, **kwargs): + events.append(("standalone",)) + return torch.ones(1, 2), [1] + + pipeline._encode_video_audio_conditions_resident = encode_embedded + pipeline._encode_audio_conditions_resident = encode_standalone + monkeypatch.setattr(module, "sequential_offload_component", record_component) + + def call(): + return pipeline._encode_reference_audio_conditions( + [{"input_has_audio": True}], + has_audio=[True], + standalone_audios=[(torch.ones(1), 16_000)], + max_duration_seconds=1.0, + ) + + if encode_fails: + with pytest.raises(RuntimeError, match="audio encode failed"): + call() + else: + result = call() + assert result[0].shape == (1, 2) + assert result[2].shape == (1, 2) + + assert events[0] == ("activate", pipeline.audio_vae) + assert events[-1] == ("offload", pipeline.audio_vae) + assert events.count(("activate", pipeline.audio_vae)) == 1 + assert events.count(("offload", pipeline.audio_vae)) == 1 diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index dc74248e5..ac3d3181d 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -8,7 +8,13 @@ import torch from torch import nn from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy -from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook +from vllm_omni.diffusion.offloader.sequential_backend import ( + ModelLevelOffloadBackend, + SequentialOffloadHook, + apply_sequential_offload, + remove_sequential_offload, + sequential_offload_component, +) from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -76,6 +82,52 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None: assert pipeline.disable_called is True +def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + dit = _create_simple_module() + encoder = _create_simple_module() + offloaded: list[nn.Module] = [] + monkeypatch.setattr( + SequentialOffloadHook, + "_to_cpu", + lambda self, module: offloaded.append(module), + ) + + apply_sequential_offload( + dit_modules=[dit], + encoder_modules=[encoder], + device=torch.device("cpu"), + offload_initial_dits=True, + ) + + assert offloaded == [dit] + remove_sequential_offload([dit, encoder]) + + +def test_direct_component_activation_failure_still_offloads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + component = _create_simple_module() + apply_sequential_offload( + dit_modules=[_create_simple_module()], + encoder_modules=[component], + device=torch.device("cpu"), + ) + hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + offloaded: list[nn.Module] = [] + + def fail_activation(module: nn.Module) -> None: + raise RuntimeError("activation failed") + + monkeypatch.setattr(hook, "pre_forward", fail_activation) + monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module)) + + with pytest.raises(RuntimeError, match="activation failed"): + with sequential_offload_component(component): + pass + + assert offloaded == [component] + + class TestMoveParamsPinMemory: def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """DTensor should skip pin_memory to avoid RuntimeError.""" @@ -85,7 +137,12 @@ class TestMoveParamsPinMemory: original_isinstance = isinstance def fake_isinstance(obj, cls): - if cls.__name__ == "DTensor": + # torchada's Tensor.to() checks tuples such as + # ``(str, torch.device)``. Keep the DTensor probe compatible with + # the normal isinstance contract instead of breaking those calls. + if type(cls) is tuple: + return any(fake_isinstance(obj, candidate) for candidate in cls) + if getattr(cls, "__name__", None) == "DTensor": return True return original_isinstance(obj, cls) diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 57fc26177..62d10488c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -8,7 +8,7 @@ import math import os import tempfile from collections.abc import Iterable, Mapping, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from itertools import groupby from pathlib import Path from typing import Any, ClassVar @@ -44,7 +44,13 @@ from vllm_omni.diffusion.models.interface import ( SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.offloader import OffloadPlan +from vllm_omni.diffusion.offloader import ( + OffloadPlan, + apply_sequential_offload, + remove_sequential_offload, + sequential_offload_component, +) +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) @@ -543,9 +549,8 @@ class MiniMaxH3Pipeline( _PROFILER_TARGETS: ClassVar[list[str]] = [ "_prepare_reference_videos", "encode_prompt", - "_encode_video_conditions", - "_encode_video_audio_conditions", - "_encode_audio_conditions", + "_encode_visual_conditions", + "_encode_reference_audio_conditions", "diffuse", "decode", ] @@ -1129,9 +1134,7 @@ class MiniMaxH3Pipeline( input_ids: torch.Tensor, vision_kwargs: dict[str, torch.Tensor], ) -> torch.Tensor: - if self.od_config.enable_cpu_offload and not getattr( - self.od_config, "enable_distributed_layerwise_offload", False - ): + if getattr(self, "_model_cpu_offload_modules", None): # Invoke nn.Module.__call__ so the generic model-level offloader # swaps the resident DiT and encoder. return self.text_encoder(input_ids, **vision_kwargs) @@ -1160,8 +1163,48 @@ class MiniMaxH3Pipeline( or getattr(od_config, "enable_distributed_layerwise_offload", False) ) + def enable_omni_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool, + use_hsdp: bool, + ) -> None: + if getattr(self, "_model_cpu_offload_modules", None): + return + + components = ModuleDiscovery.discover(self) + dits = components.dits + stages = [*components.encoders, *components.vaes] + modules = [*dits, *stages] + apply_sequential_offload( + dit_modules=dits, + encoder_modules=stages, + device=device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + offload_initial_dits=True, + ) + + self._model_cpu_offload_modules = modules + logger.info( + "MiniMax-H3 model-level CPU offload enabled for %d DiT(s), text encoder, video VAE, and audio VAE", + len(dits), + ) + + def disable_omni_model_cpu_offload(self) -> None: + modules = getattr(self, "_model_cpu_offload_modules", None) + if not modules: + return + remove_sequential_offload(modules) + self._model_cpu_offload_modules = [] + @contextmanager def _component_on_device(self, component: nn.Module): + if getattr(self, "_model_cpu_offload_modules", None): + with sequential_offload_component(component): + yield + return staged = self._uses_manual_component_offload() if staged: component.load_to_device() @@ -1171,21 +1214,6 @@ class MiniMaxH3Pipeline( if staged: component.offload_to_cpu() - def _encode_visual_condition( - self, - image: Image.Image, - ) -> torch.Tensor: - _, rank, _ = _dit_rank_world() - rows = None - if rank == 0: - with self._component_on_device(self.video_vae): - rows = self.video_vae.encode_image(image) - return _broadcast_tensor( - rows, - dtype=torch.float32, - device=self.device, - ) - def _encode_visual_conditions( self, images: list[Image.Image], @@ -1195,44 +1223,34 @@ class MiniMaxH3Pipeline( ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]: rows: list[torch.Tensor] = [] shapes: list[tuple[int, int, int]] = [] - for image in images: - rows.append(self._encode_visual_condition(image)) - shapes.append((1, image.height // 16, image.width // 16)) - if video_count: - video_rows, video_shapes = self._encode_video_conditions( - prepared_videos, - count=video_count, - ) - rows.append(video_rows) - shapes.extend(video_shapes) + _, rank, _ = _dit_rank_world() + # Keep image and video references in one residency window when both + # appear in a request; otherwise the video branch would reload the VAE. + needs_video_vae = video_count > 0 or (rank == 0 and bool(images)) + video_vae_context = self._component_on_device(self.video_vae) if needs_video_vae else nullcontext() + with video_vae_context: + if images: + image_rows = None + if rank == 0: + image_rows = torch.cat([self.video_vae.encode_image(image) for image in images]) + rows.append( + _broadcast_tensor( + image_rows, + dtype=torch.float32, + device=self.device, + ) + ) + shapes.extend((1, image.height // 16, image.width // 16) for image in images) + if video_count: + video_rows, video_shapes = self._encode_video_conditions_resident( + prepared_videos, + count=video_count, + ) + rows.append(video_rows) + shapes.extend(video_shapes) return (torch.cat(rows) if rows else None), shapes - def _encode_audio_condition( - self, - audio: tuple[torch.Tensor, int], - ) -> tuple[torch.Tensor, int]: - _, rank, _ = _dit_rank_world() - rows = None - audio_t = 0 - if rank == 0: - with self._component_on_device(self.audio_vae): - rows, audio_t = self.audio_vae.encode_waveform(*audio) - audio_t_tensor = torch.tensor( - [audio_t], - dtype=torch.long, - device=self.device, - ) - group, _, world_size = _dit_rank_world() - if world_size > 1: - dist.broadcast(audio_t_tensor, src=0, group=group) - rows = _broadcast_tensor( - rows, - dtype=torch.float32, - device=self.device, - ) - return rows, int(audio_t_tensor.item()) - - def _encode_audio_conditions( + def _encode_audio_conditions_resident( self, audios: list[tuple[torch.Tensor, int]], *, @@ -1269,15 +1287,6 @@ class MiniMaxH3Pipeline( [int(value) for value in lengths.tolist()], ) - def _encode_video_conditions( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - count: int, - ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: - with self._component_on_device(self.video_vae): - return self._encode_video_conditions_resident(prepared_videos, count=count) - def _encode_video_conditions_resident( self, prepared_videos: list[dict[str, Any]] | None, @@ -1325,15 +1334,6 @@ class MiniMaxH3Pipeline( [tuple(int(value) for value in item) for item in shapes.tolist()], ) - def _encode_video_audio_conditions( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - has_audio: list[bool], - ) -> tuple[torch.Tensor | None, list[int]]: - with self._component_on_device(self.audio_vae): - return self._encode_video_audio_conditions_resident(prepared_videos, has_audio=has_audio) - def _encode_video_audio_conditions_resident( self, prepared_videos: list[dict[str, Any]] | None, @@ -1377,6 +1377,34 @@ class MiniMaxH3Pipeline( [int(value) for value in lengths.tolist()], ) + def _encode_reference_audio_conditions( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + has_audio: list[bool], + standalone_audios: list[tuple[torch.Tensor, int]], + max_duration_seconds: float, + ) -> tuple[torch.Tensor | None, list[int], torch.Tensor | None, list[int]]: + # Embedded and standalone audio are consecutive direct Audio-VAE + # calls. Keep the component resident across both paths. + needs_audio_vae = any(has_audio) or bool(standalone_audios) + audio_vae_context = self._component_on_device(self.audio_vae) if needs_audio_vae else nullcontext() + with audio_vae_context: + embedded_condition, embedded_lengths = self._encode_video_audio_conditions_resident( + prepared_videos, + has_audio=has_audio, + ) + external_condition, external_lengths = self._encode_audio_conditions_resident( + standalone_audios, + max_duration_seconds=max_duration_seconds, + ) + return ( + embedded_condition, + embedded_lengths, + external_condition, + external_lengths, + ) + def _initial_noise( self, *, @@ -1755,12 +1783,15 @@ class MiniMaxH3Pipeline( prepared_videos, video_count=video_count, ) - embedded_audio_condition, embedded_audio_lengths = self._encode_video_audio_conditions( + ( + embedded_audio_condition, + embedded_audio_lengths, + external_audio_condition, + external_audio_lengths, + ) = self._encode_reference_audio_conditions( prepared_videos, has_audio=has_audio, - ) - external_audio_condition, external_audio_lengths = self._encode_audio_conditions( - standalone_audios, + standalone_audios=standalone_audios, max_duration_seconds=float(num_frames) / float(sampling.fps or MINIMAX_H3_FPS), ) audio_parts = [ diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 2502fa76b..d29daaded 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -8,7 +8,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig, OffloadStrategy +from .base import OffloadBackend, OffloadConfig, OffloadStrategy, SupportsModelCpuOffload from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names from .distributed_layerwise_backend import ( DistributedLayerwiseOffloadBackend, @@ -23,6 +23,7 @@ from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, remove_sequential_offload, + sequential_offload_component, ) from .tensor_utils import ( dtype_size, @@ -39,6 +40,7 @@ __all__ = [ "OffloadConfig", "OffloadPlan", "OffloadStrategy", + "SupportsModelCpuOffload", "LayerWiseOffloadBackend", "DistributedLayerwiseOffloadBackend", "DistributedLayerwiseOffloadHook", @@ -46,6 +48,7 @@ __all__ = [ "PinnedModuleStager", "apply_sequential_offload", "remove_sequential_offload", + "sequential_offload_component", "apply_distributed_block_hook", "remove_distributed_block_hook", "get_offload_backend", diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index a750d79f5..4e0eb1b5b 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum +from typing import Protocol, runtime_checkable import torch from torch import nn @@ -14,6 +15,26 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig logger = init_logger(__name__) +@runtime_checkable +class SupportsModelCpuOffload(Protocol): + """Pipeline-owned lifecycle for model-level CPU offload. + + Pipelines with non-forward component entry points (for example VAE + ``decode_latent`` methods) need to activate those stages explicitly, so + generic forward-hook discovery cannot manage their full lifecycle. + """ + + def enable_omni_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool, + use_hsdp: bool, + ) -> None: ... + + def disable_omni_model_cpu_offload(self) -> None: ... + + class OffloadStrategy(Enum): NONE = "none" MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 0b031be3a..36e0f7b22 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator +from contextlib import contextmanager + import torch from torch import nn from torch.distributed._tensor import DTensor # type: ignore[attr-defined] @@ -9,7 +12,7 @@ from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig +from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload from .module_collector import ModuleDiscovery logger = init_logger(__name__) @@ -123,6 +126,7 @@ def apply_sequential_offload( device: torch.device, pin_memory: bool = True, use_hsdp: bool = False, + offload_initial_dits: bool = False, ) -> None: """Apply sequential offloading hooks to DiT and encoder modules. @@ -136,6 +140,7 @@ def apply_sequential_offload( device: Target GPU device for loading pin_memory: Whether to pin CPU memory for faster transfers use_hsdp: Whether HSDP is enabled (affects non_blocking behavior) + offload_initial_dits: Whether to begin with all DiT modules on CPU. Example: >>> apply_sequential_offload( @@ -170,6 +175,14 @@ def apply_sequential_offload( registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) logger.debug("Registered offload hook for %s", enc.__class__.__name__) + if offload_initial_dits: + try: + for dit_mod in dit_modules: + _get_sequential_offload_hook(dit_mod)._to_cpu(dit_mod) + except Exception: + remove_sequential_offload([*dit_modules, *encoder_modules]) + raise + def remove_sequential_offload(modules: list[nn.Module]) -> None: """Remove sequential offloading hooks from modules. @@ -188,6 +201,31 @@ def remove_sequential_offload(modules: list[nn.Module]) -> None: logger.debug("Removed offload hook from %s", module.__class__.__name__) +def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook: + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None + if not isinstance(hook, SequentialOffloadHook): + raise RuntimeError(f"{module.__class__.__name__} has no sequential offload hook") + return hook + + +@contextmanager +def sequential_offload_component(module: nn.Module) -> Iterator[None]: + """Activate and release a hooked component called outside ``forward``.""" + hook = _get_sequential_offload_hook(module) + try: + hook.pre_forward(module) + yield + except BaseException: + try: + hook._to_cpu(module) + except Exception: + logger.exception("Failed to release %s after component failure", module.__class__.__name__) + raise + else: + hook._to_cpu(module) + + class ModelLevelOffloadBackend(OffloadBackend): """Model-level (sequential) offloading backend. @@ -197,23 +235,17 @@ class ModelLevelOffloadBackend(OffloadBackend): def __init__(self, config: OffloadConfig, device: torch.device): super().__init__(config, device) self._offload_modules: list[nn.Module] = [] # Track modules with hooks - self._custom_pipeline: nn.Module | None = None + self._custom_pipeline: SupportsModelCpuOffload | None = None def enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("ModelLevelOffloadBackend already enabled") return - # Pipelines with a nested transformer (e.g. Cosmos3's reasoner/generator - # pathways) own their own mutual-exclusion swaps and cannot be offloaded - # with the generic encoder<->DiT hooks. Delegate to the custom hook. - # TODO: Cosmos3 is the only implementer today, so we duck-type the - # enable/disable_omni_model_cpu_offload pair via getattr. Once a second - # pipeline needs it, formalize this as a Protocol (like - # SupportsComponentDiscovery) instead of getattr detection. - custom_enable = getattr(pipeline, "enable_omni_model_cpu_offload", None) - if callable(custom_enable): - custom_enable( + # Pipelines with non-forward component entry points own their complete + # mutual-exclusion lifecycle. Delegate through the explicit protocol. + if isinstance(pipeline, SupportsModelCpuOffload): + pipeline.enable_omni_model_cpu_offload( device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, @@ -283,9 +315,7 @@ class ModelLevelOffloadBackend(OffloadBackend): return if self._custom_pipeline is not None: - custom_disable = getattr(self._custom_pipeline, "disable_omni_model_cpu_offload", None) - if callable(custom_disable): - custom_disable() + self._custom_pipeline.disable_omni_model_cpu_offload() self._custom_pipeline = None self.enabled = False logger.info("Model-level offloading disabled")