Extend Cosmos3 CPU offloading to component and layerwise paths (#4695)

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

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

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,

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
 

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

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."""

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)
 

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

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."""

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.

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()