added initial offloading implementation

Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>

Author: rahul-steiger-nv

.agents/skills/cosmos3-inference/SKILL.md

@@ -30,6 +30,7 @@ All paths below are relative to the cosmos3 package root (`../../../` from this
 | Which model should I use? (Nano vs Super, memory, shift)                | `README.md` § Models                                                                    |
 | Which modality? (t2i, t2v, i2v, examples)                               | `README.md` § Modalities                                                                |
 | What parallelism preset? (latency vs throughput)                        | `README.md` § Inference                                                                 |
+| How do I lower GPU memory / offload to CPU? (`--offload-stages`)        | `docs/inference.md` § CPU Offloading                                                    |
 | What input fields are available? (prompt, vision_path, num_frames, ...) | `docs/inference.md` § Sample Arguments                                                  |
 | What are the default parameter values?                                  | `cosmos_framework/inference/defaults/<model_mode>/sample_args.json` (per-modality JSON) |
 | How do I use custom defaults?                                           | `docs/inference.md` § Custom Defaults                                                   |

.claude/skills/cosmos3-inference/SKILL.md

@@ -30,6 +30,7 @@ All paths below are relative to the cosmos3 package root (`../../../` from this
 | Which model should I use? (Nano vs Super, memory, shift)                | `README.md` § Models                                                                    |
 | Which modality? (t2i, t2v, i2v, examples)                               | `README.md` § Modalities                                                                |
 | What parallelism preset? (latency vs throughput)                        | `README.md` § Inference                                                                 |
+| How do I lower GPU memory / offload to CPU? (`--offload-stages`)        | `docs/inference.md` § CPU Offloading                                                    |
 | What input fields are available? (prompt, vision_path, num_frames, ...) | `docs/inference.md` § Sample Arguments                                                  |
 | What are the default parameter values?                                  | `cosmos_framework/inference/defaults/<model_mode>/sample_args.json` (per-modality JSON) |
 | How do I use custom defaults?                                           | `docs/inference.md` § Custom Defaults                                                   |

cosmos_framework/inference/args_test.py

@@ -129,6 +129,32 @@ def check_model_equal(actual: pydantic.BaseModel, expected: pydantic.BaseModel):
     check_model_equal(OmniSetupOverrides.model_validate(args.model_dump()).build_setup(), args)
 
 
+def test_offload_stages(tmp_path: Path):
+    def _build(**kwargs):
+        return OmniSetupOverrides(
+            checkpoint_path=DEFAULT_CHECKPOINT_NAME,
+            output_dir=tmp_path / "outputs",
+            **kwargs,
+        ).build_setup()
+
+    # Default: offloading disabled.
+    args = _build()
+    assert args.offload_stages == ()
+
+    # Arena stages round-trip through build_setup.
+    args = _build(offload_stages=("reasoner", "generator", "vae"))
+    assert args.offload_stages == ("reasoner", "generator", "vae")
+
+    # Guardrail offloading is a separate flag, not an --offload-stages value.
+    for bad in (("guardrails",), ("bogus",)):
+        with pytest.raises(pydantic.ValidationError):
+            OmniSetupOverrides(
+                checkpoint_path=DEFAULT_CHECKPOINT_NAME,
+                output_dir=tmp_path / "outputs",
+                offload_stages=bad,
+            )
+
+
 def test_sample_args(tmp_path: Path):
     setup_args = OmniSetupOverrides(
         checkpoint_path=DEFAULT_CHECKPOINT_NAME,

cosmos_framework/inference/common/args.py

@@ -47,6 +47,11 @@
     Training = Suppress
 
 
+# Single-GPU CPU-offload stages selectable via ``--offload-stages``. (Guardrail
+# offloading has its own dedicated flag, ``--offload-guardrail-models``.)
+OffloadStage = Literal["reasoner", "generator", "vae"]
+
+
 IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp"]
 VIDEO_EXTENSIONS = [".mp4"]
 MEDIA_EXTENSIONS = IMAGE_EXTENSIONS + VIDEO_EXTENSIONS
@@ -643,6 +648,7 @@ class SetupArgs(ABC, CheckpointArgs, ParallelismArgs, GuardrailArgs):
     warmup: pydantic.NonNegativeInt
     max_model_len: pydantic.PositiveInt | None
     max_num_seqs: pydantic.PositiveInt | None
+    offload_stages: tuple[OffloadStage, ...]
 
     # Subclass must implement these fields/methods
     # ------------------------------------------------------------
@@ -693,6 +699,14 @@ class SetupOverrides(ABC, CheckpointOverrides, ParallelismOverrides, GuardrailOv
     max_num_seqs: pydantic.PositiveInt | None = 1
     """Maximum number of sequences per batch.  When set, samples are packed into
     batches by number of sequences."""
+    offload_stages: tuple[OffloadStage, ...] = ()
+    """Single-GPU CPU-offload stages. Each named component is offloaded to pinned CPU
+    storage and staged into one reusable GPU arena only while in use, reducing peak
+    GPU memory. Choices: 'reasoner' / 'generator' (the MoT towers — enabling either
+    runs the understanding pathway once as a prefill that caches the per-layer K/V, then
+    runs the denoise loop generator-only) and 'vae' (the vision tokenizer, staged around
+    encode/decode). Empty = off (joint path, unchanged). Single-GPU only; incompatible
+    with CUDA graphs. Guardrail offloading has its own flag, --offload-guardrail-models."""
 
     def _build_setup(self):
         pass

cosmos_framework/inference/inference.py

@@ -37,6 +37,7 @@
 from cosmos_framework.inference.common.inference import Inference, sync_distributed_errors
 from cosmos_framework.inference.common.init import get_rank, get_world_size
 from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel
+from cosmos_framework.inference.offloading import OffloadPipeline
 from cosmos_framework.inference.vision import (
     build_conditioned_video_batch,
     build_image_edit_batch,
@@ -59,6 +60,91 @@
 
 UpsampleTask = Literal["t2i", "t2v", "i2v"]
 
+# Reasoner/generator CPU-offload group names (single-GPU). The two MoT pathways
+# are interleaved per layer, so the groups gather the per-layer understanding vs
+# generation submodules (plus the generation-side diffusion heads) by reference.
+REASONER_OFFLOAD_PART = "reasoner"
+GENERATOR_OFFLOAD_PART = "generator"
+VAE_OFFLOAD_PART = "vae"
+# Offload stages that ride the diffusion GPU arena (``guardrails`` is handled
+# separately by the guardrail runners' own CPU-offload path).
+ARENA_OFFLOAD_PARTS = (REASONER_OFFLOAD_PART, GENERATOR_OFFLOAD_PART, VAE_OFFLOAD_PART)
+
+
+def build_omni_offload_parts(model: OmniMoTModel) -> dict[str, "torch.nn.Module"]:
+    """Group the in-tree Cosmos3 model into offloadable module groups.
+
+    The MoT packs both pathways into a single network with dual projections per
+    layer, so rather than pointing at whole towers we gather the per-layer
+    understanding (``reasoner``) vs generation (``generator``) submodules into two
+    synthetic ``nn.ModuleList`` groups (referencing the existing submodule objects,
+    so the module tree, checkpoint loading, and the joint forward are unaffected).
+    The generation-side diffusion heads (``vae2llm`` / ``llm2vae`` / ``time_embedder``
+    and the optional action/sound heads) join the generator group because the
+    reasoner prefill skips vision/action/sound encode and decode entirely. The
+    ``vae`` group is the vision tokenizer network (staged around encode/decode).
+
+    All groups are disjoint (distinct module instances), as ``ModuleOffloadManager``
+    requires. Returns every known group; callers stage only the requested subset.
+    Modules left out of every staged group stay GPU-resident.
+    """
+    import torch.nn as nn
+
+    net = model.net
+    lm = net.language_model.model  # Qwen3VL(Moe)TextModel: embed_tokens, layers, norm, norm_moe_gen
+
+    reasoner_modules: list[nn.Module] = [lm.embed_tokens, lm.norm]
+    generator_modules: list[nn.Module] = [lm.norm_moe_gen]
+
+    # Generation-side diffusion encode/decode heads (only present when enabled).
+    for head_name in ("time_embedder", "vae2llm", "llm2vae", "action2llm", "llm2action", "sound2llm", "llm2sound"):
+        head = getattr(net, head_name, None)
+        if isinstance(head, nn.Module):
+            generator_modules.append(head)
+
+    for raw_layer in lm.layers:
+        # Unwrap torch.compile's OptimizedModule so we reference the real submodules
+        # (the manager rebinds their parameters between stages, which the compiled
+        # graph picks up as graph inputs on the next call).
+        layer = getattr(raw_layer, "_orig_mod", raw_layer)
+        attn = layer.self_attn
+        reasoner_modules += [
+            attn.q_proj,
+            attn.k_proj,
+            attn.v_proj,
+            attn.o_proj,
+            attn.q_norm,
+            attn.k_norm,
+            layer.input_layernorm,
+            layer.post_attention_layernorm,
+            layer.mlp,
+        ]
+        generator_modules += [
+            attn.q_proj_moe_gen,
+            attn.k_proj_moe_gen,
+            attn.v_proj_moe_gen,
+            attn.o_proj_moe_gen,
+            attn.q_norm_moe_gen,
+            attn.k_norm_moe_gen,
+            layer.input_layernorm_moe_gen,
+            layer.post_attention_layernorm_moe_gen,
+            layer.mlp_moe_gen,
+        ]
+
+    parts: dict[str, nn.Module] = {
+        REASONER_OFFLOAD_PART: nn.ModuleList(reasoner_modules),
+        GENERATOR_OFFLOAD_PART: nn.ModuleList(generator_modules),
+    }
+    # The vision tokenizer's underlying network (the offloadable VAE weights).
+    # ``tokenizer_vision_gen.model`` may itself be a thin (non-Module) wrapper whose
+    # ``.model`` is the actual nn.Module (e.g. Wan2pt2VAEInterface -> WanVAE -> WanVAE_).
+    vae = getattr(model.tokenizer_vision_gen, "model", None)
+    if vae is not None and not isinstance(vae, nn.Module):
+        vae = getattr(vae, "model", None)
+    if isinstance(vae, nn.Module):
+        parts[VAE_OFFLOAD_PART] = vae
+    return parts
+
 
 _BatchItem = TypeVar("_BatchItem")
 
@@ -1071,6 +1157,44 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self:
                 model.set_up_scheduler_and_sampler()
                 log.debug(f"Sampler overridden to: {sampler_override}")
 
+        # Single-GPU CPU offloading (opt-in via --offload-stages). The diffusion parts
+        # (reasoner / generator / vae) time-share one reusable GPU arena; 'guardrails'
+        # is handled separately by the guardrail runners. Default-off; when off the
+        # model's ``memory`` stays None (joint path, unchanged).
+        arena_stages = tuple(s for s in setup_args.offload_stages if s in ARENA_OFFLOAD_PARTS)
+        if arena_stages:
+            if compile_config.use_cuda_graphs:
+                raise NotImplementedError(
+                    "CPU offloading is incompatible with CUDA graphs (staging rebinds weight "
+                    "tensors between calls, which breaks captured static addresses). "
+                    "Disable --use-cuda-graphs or --offload-stages."
+                )
+            world = setup_args.dp_shard_size * setup_args.cp_size * setup_args.cfgp_size
+            if world != 1:
+                raise NotImplementedError(
+                    f"CPU offloading is single-GPU only (dp_shard*cp*cfgp must be 1, got {world})."
+                )
+            available_parts = build_omni_offload_parts(model)
+            missing = [s for s in arena_stages if s not in available_parts]
+            if missing:
+                raise ValueError(f"Requested offload stage(s) {missing} are unavailable for this model.")
+            offloader = OffloadPipeline(
+                stages=arena_stages,
+                parts={s: available_parts[s] for s in arena_stages},
+                device=torch.device("cuda", torch.cuda.current_device()),
+                pin_memory=True,
+            )
+            offloader.initialize()
+            # Drive staging from the model/decode sites without leaking inference-side
+            # machinery into the model: callers invoke ``_offload_stage_fn(part)``, which
+            # is a no-op for parts that aren't being offloaded (so call sites are uniform).
+            model._offloader = offloader  # keep alive
+            model._offload_stage_fn = lambda part: offloader.context(part) if offloader.has_part(part) else None
+            model._reasoner_generator_split = bool(
+                {REASONER_OFFLOAD_PART, GENERATOR_OFFLOAD_PART} & set(arena_stages)
+            )
+            log.info(f"Enabled single-GPU CPU offloading for stages: {', '.join(arena_stages)}.")
+
         vae_decode_stream: torch.cuda.Stream | None = None
         if setup_args.use_separate_pipeline_vision_decode_gpu:
             # The CP/CFGP ranks are partitioned into replica-local groups of size
@@ -1468,6 +1592,12 @@ def decode_vision(vision_latent: torch.Tensor) -> torch.Tensor:
 
             with self._get_timer(f"{self.model.__class__.__name__}.decode"):
                 output_vision = outputs.pop("vision")
+                # Stage the VAE onto the GPU arena for decode (no-op unless 'vae' is an
+                # offloaded stage). The denoise loop left the generator staged; this
+                # rebinds it back to CPU and brings the VAE in.
+                _stage_fn = getattr(self.model, "_offload_stage_fn", None)
+                if _stage_fn is not None:
+                    _stage_fn(VAE_OFFLOAD_PART)
                 decoded_vision = [decode_vision(vision) for vision in output_vision]
                 outputs["vision"] = [cast(torch.Tensor, vision) for vision in decoded_vision]
                 if self.vae_decode_stream is not None: