diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index f827e79b3..765fb9535 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -82,6 +82,33 @@ mutual exclusion: when one group runs, the other moves to CPU. - Support single GPU only for now +### Component offloading for split models (e.g. Cosmos3) + +Some models split their transformer into mutually-exclusive *components* that run +in different phases of a single forward pass rather than as separate pipeline +components -- e.g. Cosmos3's understanding (reasoner) component runs once per +generation while the generation (generator) component runs every denoising step. +Such models have no separate text encoder to swap against, so the transformer +owns a small model-local offload path and wraps each phase with +`with self._offload_context(name):` + +```python +class Cosmos3VFMTransformer(nn.Module): + def forward(self, ...): + with self._offload_context("reasoner"): + ... # understanding pass, runs once + with self._offload_context("generator"): + ... # denoising pass, runs every step +``` + +Model-level offloading then keeps exactly one component GPU-resident at a time +(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The +pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the +transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading +works for these models too -- each component declares its own block container via +`_layerwise_offload_blocks_attrs`. + + ## Layerwise (Blockwise) Offloading ### How It Works @@ -145,9 +172,9 @@ class Flux2Transformer2DModel(nn.Module): ``` ### Limitations -- Cold start latency increases because of - 1) components are loaded to CPU first at the very first during initialization, - 2) weight consolidation and pinning +- Cold start latency increases because offloaded components must be moved to CPU + during setup; layerwise offload may add extra weight consolidation and pinning + work. - Performance depends on compute cost and H2D bandwidth as well - Support single GPU only for now @@ -178,11 +205,18 @@ Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelH ``` OffloadBackend (base class) -├── ModelLevelOffloadBackend → uses SequentialOffloadHook +├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap) +│ (delegates to a pipeline's enable_omni_model_cpu_offload +│ for split models like Cosmos3) └── LayerWiseOffloadBackend → uses LayerwiseOffloadHook ``` -Factory function `get_offload_backend()` selects the appropriate backend based on configuration. +Factory function `get_offload_backend()` selects the appropriate backend based on +configuration. + +For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's +`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its +reasoner/generator components inside the model forward pass. ## Supported Models diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 37bc99aa0..d3ed3bd67 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -105,7 +105,7 @@ The following tables show which models support each feature: > Notes: -> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. +> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md). > 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details. ### ImageGen @@ -144,7 +144,7 @@ The following tables show which models support each feature: > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary. -> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload. +> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload. > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. ### VideoGen diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 77215f79b..f0159f82e 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -20,6 +20,16 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_pipeline_declares_layerwise_offload_components() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + assert Cosmos3OmniDiffusersPipeline._dit_modules == ["transformer.language_model", "transformer"] + assert Cosmos3OmniDiffusersPipeline._encoder_modules == [] + assert Cosmos3OmniDiffusersPipeline._vae_modules == ["vae"] + assert Cosmos3OmniDiffusersPipeline._resident_modules == [] + assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload") + + class StubScheduler: def __init__( self, diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 70598c952..13d1c399b 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -10,9 +10,19 @@ import pytest import torch from torch import nn +from vllm_omni.platforms import current_omni_platform + pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.fixture +def accelerator_device() -> torch.device: + """Provide an accelerator device, skipping when none is available.""" + if current_omni_platform.get_device_count() == 0: + pytest.skip("Accelerator required for this test") + return current_omni_platform.get_torch_device(0) + + def _tiny_cosmos3_config(**overrides): config = { "hidden_size": 8, @@ -88,7 +98,7 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: - from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer model = object.__new__(Cosmos3VFMTransformer) nn.Module.__init__(model) @@ -104,6 +114,7 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: ] assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"] + assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"] assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"] model.latent_patch_size = 2 @@ -131,6 +142,97 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No assert tuple(output.shape) == (1, 2, 1, 2, 2) +def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + + class ToyReasonerLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + def forward(self, hidden: torch.Tensor, freqs): + del freqs + kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1) + return hidden + self.weight.to(hidden.dtype), kv, kv + + class ToyGeneratorLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + def forward(self, hidden: torch.Tensor, **kwargs): + del kwargs + return hidden + self.weight.to(hidden.dtype) + + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + model.language_model.layers = nn.ModuleList([ToyReasonerLayer()]) + model.gen_layers = nn.ModuleList([ToyGeneratorLayer()]) + model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False) + + assert model._model_cpu_offload_enabled + assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"} + assert model.device == torch.device("cpu") + + output = model( + hidden_states=torch.zeros(1, 2, 1, 2, 2), + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + ) + + assert tuple(output.shape) == (1, 2, 1, 2, 2) + # The reasoner runs once then the generator component stays resident for GEN. + assert model._active_model_cpu_offload_component == "generator" + + model.disable_model_cpu_offload() + assert not model._model_cpu_offload_enabled + + +def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device( + monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device +) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)]) + model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) + model.to(accelerator_device) + + reasoner_param = model.language_model.layers[0].weight + generator_param = model.gen_layers[0].weight + + model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False) + + assert model._model_cpu_offload_enabled + # On enable, every group is parked on CPU until a phase activates it. + assert reasoner_param.device.type == "cpu" + assert generator_param.device.type == "cpu" + + model._activate_model_cpu_offload_component("reasoner") + assert reasoner_param.device == accelerator_device + assert generator_param.device.type == "cpu" + + model._activate_model_cpu_offload_component("generator") + assert reasoner_param.device.type == "cpu" + assert generator_param.device == accelerator_device + + model.disable_model_cpu_offload() + assert not model._model_cpu_offload_enabled + # Disable restores both components onto the device. + assert reasoner_param.device == accelerator_device + assert generator_param.device == accelerator_device + + def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py index bc7fd97d1..e97c1953a 100644 --- a/tests/diffusion/offloader/test_module_collector.py +++ b/tests/diffusion/offloader/test_module_collector.py @@ -149,6 +149,40 @@ class MultiVaePipeline(nn.Module, SupportsComponentDiscovery): self.audio_vae = nn.Linear(10, 10) +class SiblingDiTsPipeline(nn.Module, SupportsComponentDiscovery): + """Two independent DiTs (e.g. a MoE transformer pair).""" + + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.transformer_2 = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class NestedDiTPipeline(nn.Module, SupportsComponentDiscovery): + """Cosmos3-like: a DiT plus one of its own submodules listed as a DiT. + + The nested entry exists so offloading can treat the inner stack as an + independent ring; the outer DiT's _hsdp_shard_conditions still own all + nested blocks for sharding. Order is nested-first, matching Cosmos3. + """ + + _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Module() + self.transformer.language_model = nn.Linear(10, 10) + self.transformer.gen_layers = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -238,3 +272,35 @@ class TestProtocolDiscovery: assert len(result.vaes) == 2 assert result.vaes[0] is pipeline.vae assert result.vaes[1] is pipeline.audio_vae + + +class TestOutermostDits: + """PipelineModules.outermost_dits() drops DiTs nested inside another DiT. + + HSDP shards each returned DiT via its own _hsdp_shard_conditions, so a DiT + nested inside another discovered DiT must be excluded: it is already covered + by its ancestor's conditions and may not declare conditions of its own. + """ + + def test_keeps_independent_sibling_dits(self): + result = ModuleDiscovery.discover(SiblingDiTsPipeline()) + + names, modules = result.outermost_dits() + + assert names == ["transformer", "transformer_2"] + assert [id(m) for m in modules] == [id(m) for m in result.dits] + + def test_drops_nested_dit_keeping_only_ancestor(self): + pipeline = NestedDiTPipeline() + result = ModuleDiscovery.discover(pipeline) + + # Discovery sees both, nested-first (matches Cosmos3's _dit_modules order). + assert result.dit_names == ["transformer.language_model", "transformer"] + + names, modules = result.outermost_dits() + + # Only the ancestor survives; sharding the nested condition-less stack + # again would double-wrap its blocks and raise in apply_hsdp_to_model. + assert names == ["transformer"] + assert len(modules) == 1 + assert modules[0] is pipeline.transformer diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index 2539cc068..dc74248e5 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -7,7 +7,8 @@ import pytest import torch from torch import nn -from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -41,6 +42,40 @@ def _track_pin_memory_calls(): return tracker, mock +def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None: + class CustomPipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.enable_args = None + self.disable_called = False + + def enable_omni_model_cpu_offload(self, **kwargs) -> None: + self.enable_args = kwargs + + def disable_omni_model_cpu_offload(self) -> None: + self.disable_called = True + + pipeline = CustomPipeline() + backend = ModelLevelOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert backend.enabled is True + assert pipeline.enable_args == { + "device": torch.device("cpu"), + "pin_memory": False, + "use_hsdp": False, + } + + backend.disable() + + assert backend.enabled is False + assert pipeline.disable_called is True + + class TestMoveParamsPinMemory: def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """DTensor should skip pin_memory to avoid RuntimeError.""" diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d8b9c90c5..e1aa0fd1c 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -636,6 +636,14 @@ class DiffusersPipelineLoader: # where the transformer lives at "pipe.transformer"). discovered_modules = ModuleDiscovery.discover(model) + # Shard only the outermost DiTs. A pipeline may list a DiT and one of its + # submodules as separate DiTs (e.g. Cosmos3's transformer and the nested + # transformer.language_model) for offload's independent rings; for HSDP an + # inner DiT is already covered by its ancestor's _hsdp_shard_conditions, so + # sharding it again would double-wrap blocks and require the inner stack to + # declare its own conditions. + outer_dit_names, outer_dits = discovered_modules.outermost_dits() + # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params. @@ -647,14 +655,14 @@ class DiffusersPipelineLoader: prepare_fp8_layers_for_fsdp, ) - for trans in discovered_modules.dits: + for trans in outer_dits: prepare_fp8_layers_for_fsdp(trans) - if not discovered_modules.dits: + if not outer_dits: raise ValueError("No DiT modules discovered for HSDP sharding") - # Apply HSDP sharding to all discovered DiT transformers - for name, trans in zip(discovered_modules.dit_names, discovered_modules.dits): + # Apply HSDP sharding to each outermost DiT transformer + for name, trans in zip(outer_dit_names, outer_dits): logger.debug("Applying HSDP to %s", name) apply_hsdp_to_model(trans, hsdp_config, target_device=target_device) diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 3dc0d4899..4153eaa67 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -57,6 +57,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.interface import ( ReferenceVideoDecodeSpec, SupportImageInput, + SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin @@ -628,7 +629,12 @@ def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig): # Pipeline # --------------------------------------------------------------------------- class Cosmos3OmniDiffusersPipeline( - nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin + nn.Module, + CFGParallelMixin, + SupportImageInput, + SupportsComponentDiscovery, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, ): """Cosmos3 text/image/video/sound/action pipeline. @@ -671,6 +677,10 @@ class Cosmos3OmniDiffusersPipeline( support_image_input: ClassVar[bool] = True color_format: ClassVar[str] = "RGB" + _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] @classmethod def reference_video_decode_spec( @@ -714,12 +724,6 @@ class Cosmos3OmniDiffusersPipeline( prefix: str = "", ) -> None: super().__init__() - if od_config.enable_cpu_offload: - raise ValueError( - "Cosmos3 has no separate text encoder, so CPU offloading " - "(transformer↔encoder swapping) is not supported. " - "Use --enable-layerwise-offload instead." - ) self.od_config = od_config self.device = get_local_device() self.dtype = od_config.dtype @@ -823,6 +827,32 @@ class Cosmos3OmniDiffusersPipeline( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + def enable_omni_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool = True, + use_hsdp: bool = False, + ) -> None: + """Enable Cosmos3 component-level model offload. + + Cosmos3 has a nested reasoner/generator transformer instead of separate + text-encoder and DiT pipeline components, so the transformer owns the + mutual-exclusion swaps. The VAE stays resident on GPU like the generic + model-level offloader. + """ + self.vae.to(device, non_blocking=True) + if isinstance(self._sound_tokenizer, nn.Module): + self._sound_tokenizer.to(device) + self.transformer.enable_model_cpu_offload( + device=device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + ) + + def disable_omni_model_cpu_offload(self) -> None: + self.transformer.disable_model_cpu_offload() + # -- Weight loading -------------------------------------------------------- @staticmethod diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 261573335..21e723ced 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -12,7 +12,9 @@ Ported from the TRT-LLM integration (tekit branch user/shreyasm/cosmos3). from __future__ import annotations import math -from typing import Any +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import TYPE_CHECKING, Any import torch import torch.distributed as dist @@ -38,6 +40,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_ from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm from vllm_omni.platforms import current_omni_platform +if TYPE_CHECKING: + from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook + logger = init_logger(__name__) @@ -869,6 +874,8 @@ class Cosmos3LanguageModel(nn.Module): computed once and reused across all sampling steps. """ + _layerwise_offload_blocks_attrs = ["layers"] + def __init__( self, *, @@ -959,7 +966,13 @@ class Cosmos3VFMTransformer(nn.Module): runs at each denoising step over the target video/image latent stream and optional transfer-control, action, and sound latent streams. - Layerwise offloading uses ``gen_layers`` as the block container. + Layerwise offloading uses ``gen_layers`` as the GEN block container. The + nested UND language model declares its own ``layers`` container so the two + pathways are offloaded as independent rings. + + Model-level CPU offload is Cosmos3-local: the UND ``reasoner`` and GEN + ``generator`` components are swapped at the phase boundaries marked by + ``with self._offload_context(...)``. Sequence parallelism uses ``_sp_plan`` to shard/gather the GEN pathway at module boundaries. ``Cosmos3CrossAttention`` checks @@ -1148,10 +1161,125 @@ class Cosmos3VFMTransformer(nn.Module): self.cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None self.cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None + self._model_cpu_offload_enabled = False + self._model_cpu_offload_device: torch.device | None = None + self._model_cpu_offload_mover: SequentialOffloadHook | None = None + self._active_model_cpu_offload_component: str | None = None + @property def device(self) -> torch.device: + offload_device = getattr(self, "_model_cpu_offload_device", None) + if getattr(self, "_model_cpu_offload_enabled", False) and offload_device is not None: + return offload_device return next(self.parameters()).device + def _model_cpu_offload_components(self) -> dict[str, list[nn.Module]]: + """Cosmos3's mutually-exclusive reasoner/generator component sets.""" + return { + "reasoner": [self.language_model.layers], + "generator": [self.gen_layers], + } + + def _model_cpu_offload_component_tensor_ids(self) -> set[int]: + component_tensors: set[int] = set() + for modules in self._model_cpu_offload_components().values(): + for module in modules: + component_tensors.update(id(param) for param in module.parameters()) + component_tensors.update(id(buffer) for buffer in module.buffers()) + return component_tensors + + def _move_model_cpu_offload_residents(self) -> None: + """Keep non-reasoner/non-generator weights resident on the target device.""" + component_tensors = self._model_cpu_offload_component_tensor_ids() + device = self._model_cpu_offload_device + if device is None: + return + with torch.no_grad(): + for param in self.parameters(): + if id(param) not in component_tensors and param.data.device != device: + param.data = param.data.to(device, non_blocking=False) + for buffer in self.buffers(): + if id(buffer) not in component_tensors and buffer.device != device: + buffer.data = buffer.data.to(device, non_blocking=False) + + def _offload_model_cpu_component(self, name: str) -> None: + mover = self._model_cpu_offload_mover + if mover is None: + raise RuntimeError("Cosmos3 model CPU offload is not enabled") + for module in self._model_cpu_offload_components()[name]: + mover._to_cpu(module) + + def _load_model_cpu_component(self, name: str) -> None: + mover = self._model_cpu_offload_mover + if mover is None: + raise RuntimeError("Cosmos3 model CPU offload is not enabled") + for module in self._model_cpu_offload_components()[name]: + mover._to_gpu(module) + + def enable_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool = True, + use_hsdp: bool = False, + ) -> None: + """Enable Cosmos3 reasoner/generator CPU swapping inside ``forward``.""" + if getattr(self, "_model_cpu_offload_enabled", False): + return + + from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook + + self._model_cpu_offload_device = torch.device(device) + self._model_cpu_offload_mover = SequentialOffloadHook( + offload_targets=[], + device=self._model_cpu_offload_device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + ) + self._move_model_cpu_offload_residents() + for name in self._model_cpu_offload_components(): + self._offload_model_cpu_component(name) + self._model_cpu_offload_enabled = True + self._active_model_cpu_offload_component = None + logger.info("Cosmos3 component-level CPU offload enabled on %s", self._model_cpu_offload_device) + + def disable_model_cpu_offload(self) -> None: + if not getattr(self, "_model_cpu_offload_enabled", False): + return + for name in self._model_cpu_offload_components(): + self._load_model_cpu_component(name) + if self._model_cpu_offload_device is not None and self._model_cpu_offload_device.type != "cpu": + current_omni_platform.synchronize() + self._model_cpu_offload_enabled = False + self._model_cpu_offload_device = None + self._model_cpu_offload_mover = None + self._active_model_cpu_offload_component = None + + def _activate_model_cpu_offload_component(self, name: str) -> None: + if not getattr(self, "_model_cpu_offload_enabled", False): + return + components = self._model_cpu_offload_components() + if name not in components: + raise ValueError(f"Unknown Cosmos3 offload component: {name!r} (known: {list(components)})") + if self._active_model_cpu_offload_component == name: + return + self._active_model_cpu_offload_component = None + for other in components: + if other != name: + self._offload_model_cpu_component(other) + self._load_model_cpu_component(name) + self._active_model_cpu_offload_component = name + + @contextmanager + def _model_cpu_offload_context(self, name: str) -> Iterator[None]: + self._activate_model_cpu_offload_component(name) + yield + + def _offload_context(self, name: str) -> AbstractContextManager[None]: + if not getattr(self, "_model_cpu_offload_enabled", False): + return nullcontext() + return self._model_cpu_offload_context(name) + # -- Patchify / Unpatchify ----------------------------------------------- def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]: @@ -1483,26 +1611,14 @@ class Cosmos3VFMTransformer(nn.Module): # Query Ulysses state at runtime ulysses_size, _, _ = _get_ulysses_state() - # Patchify latents and project to hidden space - hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w)) - s_video = hidden_video.shape[1] - s_control = 0 - hidden_controls: list[torch.Tensor] = [] - for idx, control in enumerate(control_latent_list): - if control.shape != hidden_states.shape: - raise ValueError( - "Cosmos3 transfer control latent shape must match target latent shape: " - f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}." - ) - hidden_control = self.proj_in( - self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w) - ) - hidden_controls.append(hidden_control) - s_control += hidden_control.shape[1] + # Pack action/sound tokens (no learned weights) up front so the UND + # cache sizing knows their token lengths. The modality projections are + # deferred into the generator offload context below, so model-level + # offload does not stage GEN weights before swapping to the reasoner. + action_tokens = None + sound_tokens = None s_action = 0 - hidden_action = None s_sound = 0 - hidden_sound = None if action_latents is not None: if action_latents.shape[0] != hidden_states.shape[0]: raise ValueError( @@ -1511,60 +1627,16 @@ class Cosmos3VFMTransformer(nn.Module): ) if action_domain_ids is None: action_domain_ids = torch.zeros(action_latents.shape[0], dtype=torch.long, device=action_latents.device) - hidden_action = self.action_proj_in(self.pack_action(action_latents), action_domain_ids) - hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) - s_action = hidden_action.shape[1] + action_tokens = self.pack_action(action_latents) + s_action = action_tokens.shape[1] if sound_latents is not None: if sound_latents.shape[0] != hidden_states.shape[0]: raise ValueError( "Cosmos3 sound and video batch sizes must match: " f"video={hidden_states.shape[0]}, sound={sound_latents.shape[0]}." ) - hidden_sound = self.audio_proj_in(self.pack_sound(sound_latents)) - hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype) - s_sound = hidden_sound.shape[1] - - # Timestep embedding (fp32 for precision). - # For I2V: only add to noisy tokens, not conditioned ones. - # Conditioned frames are clean context and should not receive - # the diffusion timestep signal. - with torch.autocast(current_omni_platform.device_type, enabled=False): - time_embed = self.time_embedder((timestep * self.timestep_scale).float()) - time_embed = time_embed.to(hidden_states.dtype) - - if noisy_frame_mask is not None: - # Build per-token mask from per-frame mask. - # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1] - token_noisy_mask = ( - noisy_frame_mask[:, 0, :, 0, 0] # [B, t] - .unsqueeze(-1) # [B, t, 1] - .expand(-1, -1, hp * wp) # [B, t, hp*wp] - .reshape(hidden_video.shape[0], -1, 1) # [B, t*hp*wp, 1] - ) - hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask - else: - hidden_video = hidden_video + time_embed.unsqueeze(1) - - if hidden_action is not None: - if action_noisy_mask is None: - hidden_action = hidden_action + time_embed.unsqueeze(1) - else: - if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1): - raise ValueError( - "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], " - f"got {tuple(action_noisy_mask.shape)}." - ) - action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device) - hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask - - if hidden_sound is not None: - hidden_sound = hidden_sound + time_embed.unsqueeze(1) - hidden_parts = [*hidden_controls, hidden_video] - if hidden_action is not None: - hidden_parts.append(hidden_action) - if hidden_sound is not None: - hidden_parts.append(hidden_sound) - hidden_gen = torch.cat(hidden_parts, dim=1) + sound_tokens = self.pack_sound(sound_latents) + s_sound = sound_tokens.shape[1] # Run UND pathway once and cache K/V (replicated across all ranks) if self.cached_kv is None: @@ -1583,90 +1655,162 @@ class Cosmos3VFMTransformer(nn.Module): num_vision_items=len(control_latent_list) + 1, share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) - cached_kv_full = self.language_model(text_ids, freqs_und) + with self._offload_context("reasoner"): + cached_kv_full = self.language_model(text_ids, freqs_und) self.cached_freqs_gen = freqs_gen # Trim to real text length (remove padding). K/V stay replicated; # the framework Attention layer head-slices them via joint_key/value. self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full] - # Run GEN layers. UND K/V (replicated) is passed to each layer; - # the Cosmos3CrossAttention forwards them as joint_key/value so the - # framework Attention handles the Ulysses head-slicing internally. - if self.cached_kv is None or self.cached_freqs_gen is None: - raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") - self._validate_gen_sequence_parallel( - s_gen=hidden_gen.shape[1], - s_video=s_video, - s_control=s_control, - s_action=s_action, - s_sound=s_sound, - has_action=has_action, - has_sound=has_sound, - has_control=has_control, - ulysses_size=ulysses_size, - ) - freqs_cos, freqs_sin = self.cached_freqs_gen - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) - freqs_gen = (freqs_cos, freqs_sin) - - if len(self.gen_layers) == len(self.cached_kv): - for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): - hidden_gen = layer( - hidden_gen, - k_und=k_und, - v_und=v_und, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, + with self._offload_context("generator"): + # Patchify latents and project to hidden space after UND cache + # construction, so model-level offload does not stage GEN weights + # before immediately swapping to the reasoner. + hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w)) + s_video = hidden_video.shape[1] + s_control = 0 + hidden_controls: list[torch.Tensor] = [] + for idx, control in enumerate(control_latent_list): + if control.shape != hidden_states.shape: + raise ValueError( + "Cosmos3 transfer control latent shape must match target latent shape: " + f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}." + ) + hidden_control = self.proj_in( + self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w) ) - # Cache-dit's block wrapper may return a tuple; unwrap it. - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] - else: - # Cache-dit patches gen_layers to a grouped wrapper. - for layer in self.gen_layers: - hidden_gen = layer( - hidden_gen, - cached_kv=self.cached_kv, - freqs_gen=freqs_gen, + hidden_controls.append(hidden_control) + s_control += hidden_control.shape[1] + hidden_action = None + hidden_sound = None + if action_tokens is not None: + assert action_domain_ids is not None + hidden_action = self.action_proj_in(action_tokens, action_domain_ids) + hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) + if sound_tokens is not None: + hidden_sound = self.audio_proj_in(sound_tokens) + hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype) + + # Timestep embedding (fp32 for precision). + # For I2V: only add to noisy tokens, not conditioned ones. + # Conditioned frames are clean context and should not receive + # the diffusion timestep signal. + with torch.autocast(current_omni_platform.device_type, enabled=False): + time_embed = self.time_embedder((timestep * self.timestep_scale).float()) + time_embed = time_embed.to(hidden_states.dtype) + + if noisy_frame_mask is not None: + # Build per-token mask from per-frame mask. + # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1] + token_noisy_mask = ( + noisy_frame_mask[:, 0, :, 0, 0] # [B, t] + .unsqueeze(-1) # [B, t, 1] + .expand(-1, -1, hp * wp) # [B, t, hp*wp] + .reshape(hidden_video.shape[0], -1, 1) # [B, t*hp*wp, 1] ) - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] + hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask + else: + hidden_video = hidden_video + time_embed.unsqueeze(1) - hidden_gen = self.gen_sp_gather(hidden_gen) + if hidden_action is not None: + if action_noisy_mask is None: + hidden_action = hidden_action + time_embed.unsqueeze(1) + else: + if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1): + raise ValueError( + "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], " + f"got {tuple(action_noisy_mask.shape)}." + ) + action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device) + hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask - # Final norm and project back to latent space - hidden_gen = self.norm_moe_gen(hidden_gen) - if not has_action and not has_sound and not has_control: - return self.unpatchify(self.proj_out(hidden_gen), t, h, w) + if hidden_sound is not None: + hidden_sound = hidden_sound + time_embed.unsqueeze(1) + hidden_parts = [*hidden_controls, hidden_video] + if hidden_action is not None: + hidden_parts.append(hidden_action) + if hidden_sound is not None: + hidden_parts.append(hidden_sound) + hidden_gen = torch.cat(hidden_parts, dim=1) - split_sizes = [] - if has_control: - split_sizes.append(s_control) - split_sizes.append(s_video) - if has_action: - split_sizes.append(s_action) - if has_sound: - split_sizes.append(s_sound) - split_hidden = hidden_gen.split(split_sizes, dim=1) - split_idx = 0 - if has_control: + # Run GEN layers. UND K/V (replicated) is passed to each layer; + # the Cosmos3CrossAttention forwards them as joint_key/value so the + # framework Attention handles the Ulysses head-slicing internally. + if self.cached_kv is None or self.cached_freqs_gen is None: + raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") + self._validate_gen_sequence_parallel( + s_gen=hidden_gen.shape[1], + s_video=s_video, + s_control=s_control, + s_action=s_action, + s_sound=s_sound, + has_action=has_action, + has_sound=has_sound, + has_control=has_control, + ulysses_size=ulysses_size, + ) + freqs_cos, freqs_sin = self.cached_freqs_gen + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + freqs_gen = (freqs_cos, freqs_sin) + + if len(self.gen_layers) == len(self.cached_kv): + for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): + hidden_gen = layer( + hidden_gen, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + ) + # Cache-dit's block wrapper may return a tuple; unwrap it. + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + else: + # Cache-dit patches gen_layers to a grouped wrapper. + for layer in self.gen_layers: + hidden_gen = layer( + hidden_gen, + cached_kv=self.cached_kv, + freqs_gen=freqs_gen, + ) + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + + hidden_gen = self.gen_sp_gather(hidden_gen) + + # Final norm and project back to latent space + hidden_gen = self.norm_moe_gen(hidden_gen) + if not has_action and not has_sound and not has_control: + return self.unpatchify(self.proj_out(hidden_gen), t, h, w) + + split_sizes = [] + if has_control: + split_sizes.append(s_control) + split_sizes.append(s_video) + if has_action: + split_sizes.append(s_action) + if has_sound: + split_sizes.append(s_sound) + split_hidden = hidden_gen.split(split_sizes, dim=1) + split_idx = 0 + if has_control: + split_idx += 1 + hidden_video = split_hidden[split_idx] split_idx += 1 - hidden_video = split_hidden[split_idx] - split_idx += 1 - video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) - if has_control: - return video_pred - outputs: list[torch.Tensor] = [video_pred] - if has_action: - hidden_action = split_hidden[split_idx] - split_idx += 1 - assert action_domain_ids is not None - outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids))) - if has_sound: - hidden_sound = split_hidden[split_idx] - outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) - return tuple(outputs) + video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) + if has_control: + return video_pred + outputs: list[torch.Tensor] = [video_pred] + if has_action: + hidden_action = split_hidden[split_idx] + split_idx += 1 + assert action_domain_ids is not None + outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids))) + if has_sound: + hidden_sound = split_hidden[split_idx] + outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) + return tuple(outputs) def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index 320f9bf38..17d4ca2fd 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -22,6 +22,38 @@ class PipelineModules: resident_modules: list[nn.Module] = field(default_factory=list) resident_names: list[str] = field(default_factory=list) + def outermost_dits(self) -> tuple[list[str], list[nn.Module]]: + """Discovered DiTs that are not nested inside another discovered DiT. + + Some pipelines declare a DiT and one of its own submodules as separate + DiTs (e.g. Cosmos3 lists both ``transformer`` and the nested + ``transformer.language_model``) so that offloading can treat them as + independent rings. HSDP instead shards each DiT through that DiT's own + ``_hsdp_shard_conditions``; an inner DiT's blocks are already covered by + its ancestor's conditions, so sharding it again would double-wrap them + (and the inner stack may not declare conditions at all). Keeping only the + outermost DiTs leaves the ancestor's conditions as the single source of + truth for every nested block. Order follows ``dit_names``. + """ + + def is_nested_in_another_dit(dit: nn.Module) -> bool: + for other in self.dits: + if other is dit: + continue + for submodule in other.modules(): + if submodule is dit: + return True + return False + + outer_names: list[str] = [] + outer_modules: list[nn.Module] = [] + for name, dit in zip(self.dit_names, self.dits): + if is_nested_in_another_dit(dit): + continue + outer_names.append(name) + outer_modules.append(dit) + return outer_names, outer_modules + class ModuleDiscovery: """Discovers pipeline components. diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 0edf77651..0b031be3a 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -197,12 +197,35 @@ 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 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( + device=self.device, + pin_memory=self.config.pin_cpu_memory, + use_hsdp=self.config.use_hsdp, + ) + self._custom_pipeline = pipeline + self.enabled = True + logger.info( + "Model-level offloading enabled through %s.enable_omni_model_cpu_offload", + pipeline.__class__.__name__, + ) + return + modules = ModuleDiscovery.discover(pipeline) # Move encoders to GPU @@ -259,6 +282,15 @@ class ModelLevelOffloadBackend(OffloadBackend): if not self.enabled: 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 = None + self.enabled = False + logger.info("Model-level offloading disabled") + return + remove_sequential_offload(self._offload_modules) self._offload_modules.clear()