materialize reasoner and generator on cpu instead of gpu

rahul-steiger-nv · rahul-steiger-nv · commit 2bfe944671ba · 2026-06-03T08:25:24.000-07:00
Signed-off-by: Rahul Steiger &lt;rsteiger@nvidia.com&gt;
diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py
@@ -51,7 +51,7 @@
 from cosmos_framework.tools.visualize.video import save_img_or_video
 from cosmos_framework.configs.base.defaults.compile import CompileConfig
 from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
-from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel
+from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel, cpu_offload_materialization
 from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING
 from cosmos_framework.model.vfm.upsampler.prompts import is_upsampled_prompt
 
@@ -74,67 +74,19 @@
 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.
+    The reasoner/generator partition is owned by the network
+    (:meth:`Cosmos3VFMNetwork.offload_module_groups`) so the load-time CPU
+    materialization and this runtime ``OffloadPipeline`` share one source of truth.
+    Each group is wrapped in an ``nn.ModuleList`` referencing the existing submodule
+    objects (no module-tree changes), and the ``vae`` group is added for the vision
+    tokenizer network. All groups are disjoint; callers stage only the requested
+    subset, and 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),
-    }
+    groups = model.net.offload_module_groups()  # {"reasoner": [...], "generator": [...]}
+    parts: dict[str, nn.Module] = {name: nn.ModuleList(modules) for name, modules in groups.items()}
+
     # 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_).
@@ -1102,60 +1054,68 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self:
         sampler_override = setup_args.sampler
         parallelism_config = cls._get_parallelism_config(setup_args)
         compile_config = cls._get_compile_config(setup_args)
-        if setup_args.checkpoint_type == CheckpointType.DCP and setup_args.config_file_type == ConfigFileType.MODULE:
-            from cosmos_framework.inference.common.config import save_config
-            from cosmos_framework.utils.vfm.model_loader import load_model_from_checkpoint
-
-            if not setup_args.experiment:
-                raise ValueError("'experiment' is required")
-            if not setup_args.config_file:
-                raise ValueError("'config_file' is required")
-
-            Cosmos3OmniModel.before_load_model()
-            model, config = load_model_from_checkpoint(
-                experiment_name=setup_args.experiment,
-                config_file=setup_args.config_file,
-                checkpoint_path=setup_args.checkpoint_path,
-                credential_path=setup_args.credential_path or None,
-                parallelism_config=attrs.asdict(parallelism_config),
-                compile_config=attrs.asdict(compile_config),
-                load_ema_to_reg=setup_args.use_ema_weights,
-                experiment_opts=[
-                    *setup_args.experiment_overrides,
-                    f"model.config.rectified_flow_inference_config.scheduler_type={sampler_override}",
-                ],
-                use_cache_checkpoint=setup_args.checkpoint_cache_dir is not None,
-                cache_checkpoint_rootdir=str(setup_args.checkpoint_cache_dir or ""),
-            )
-            model = cast("OmniMoTModel", model)
-            Cosmos3OmniModel.after_load_model(model)
-            save_config(config, setup_args.output_dir)
-        else:
-            checkpoint_path = setup_args.download_checkpoint()
-            if setup_args.config_file_type == ConfigFileType.MODULE:
-                config = None
+        # Two-phase materialization for single-GPU CPU offloading: build the reasoner/
+        # generator towers directly on CPU during model construction so they never occupy
+        # GPU memory (the checkpoint shards load straight into CPU tensors). No-op when
+        # those stages aren't requested.
+        net_offload_parts = tuple(
+            s for s in setup_args.offload_stages if s in (REASONER_OFFLOAD_PART, GENERATOR_OFFLOAD_PART)
+        )
+        with cpu_offload_materialization(net_offload_parts):
+            if setup_args.checkpoint_type == CheckpointType.DCP and setup_args.config_file_type == ConfigFileType.MODULE:
+                from cosmos_framework.inference.common.config import save_config
+                from cosmos_framework.utils.vfm.model_loader import load_model_from_checkpoint
+
+                if not setup_args.experiment:
+                    raise ValueError("'experiment' is required")
+                if not setup_args.config_file:
+                    raise ValueError("'config_file' is required")
+
+                Cosmos3OmniModel.before_load_model()
+                model, config = load_model_from_checkpoint(
+                    experiment_name=setup_args.experiment,
+                    config_file=setup_args.config_file,
+                    checkpoint_path=setup_args.checkpoint_path,
+                    credential_path=setup_args.credential_path or None,
+                    parallelism_config=attrs.asdict(parallelism_config),
+                    compile_config=attrs.asdict(compile_config),
+                    load_ema_to_reg=setup_args.use_ema_weights,
+                    experiment_opts=[
+                        *setup_args.experiment_overrides,
+                        f"model.config.rectified_flow_inference_config.scheduler_type={sampler_override}",
+                    ],
+                    use_cache_checkpoint=setup_args.checkpoint_cache_dir is not None,
+                    cache_checkpoint_rootdir=str(setup_args.checkpoint_cache_dir or ""),
+                )
+                model = cast("OmniMoTModel", model)
+                Cosmos3OmniModel.after_load_model(model)
+                save_config(config, setup_args.output_dir)
             else:
-                model_dict = setup_args.load_model_config_dict()
-                if setup_args.vlm_processor_from_checkpoint:
-                    # Source the VLM processor from the loaded checkpoint's own
-                    # bundled files instead of the repository hardcoded in the
-                    # model config. Drops the redundant base-model download.
-                    tokenizer_cfg = model_dict["config"]["vlm_config"]["tokenizer"]
-                    tokenizer_cfg.pop("repository", None)
-                    tokenizer_cfg.pop("revision", None)
-                    tokenizer_cfg.pop("subdir", None)
-                    tokenizer_cfg["tokenizer_type"] = str(checkpoint_path)
-                config = Cosmos3OmniConfig(model=model_dict)
-            model = Cosmos3OmniModel.from_pretrained_dcp(
-                checkpoint_path,
-                config=config,
-                parallelism_config=parallelism_config,
-                compile_config=compile_config,
-            ).model
-            if model.config.rectified_flow_inference_config.scheduler_type != sampler_override:
-                model.config.rectified_flow_inference_config.scheduler_type = sampler_override
-                model.set_up_scheduler_and_sampler()
-                log.debug(f"Sampler overridden to: {sampler_override}")
+                checkpoint_path = setup_args.download_checkpoint()
+                if setup_args.config_file_type == ConfigFileType.MODULE:
+                    config = None
+                else:
+                    model_dict = setup_args.load_model_config_dict()
+                    if setup_args.vlm_processor_from_checkpoint:
+                        # Source the VLM processor from the loaded checkpoint's own
+                        # bundled files instead of the repository hardcoded in the
+                        # model config. Drops the redundant base-model download.
+                        tokenizer_cfg = model_dict["config"]["vlm_config"]["tokenizer"]
+                        tokenizer_cfg.pop("repository", None)
+                        tokenizer_cfg.pop("revision", None)
+                        tokenizer_cfg.pop("subdir", None)
+                        tokenizer_cfg["tokenizer_type"] = str(checkpoint_path)
+                    config = Cosmos3OmniConfig(model=model_dict)
+                model = Cosmos3OmniModel.from_pretrained_dcp(
+                    checkpoint_path,
+                    config=config,
+                    parallelism_config=parallelism_config,
+                    compile_config=compile_config,
+                ).model
+                if model.config.rectified_flow_inference_config.scheduler_type != sampler_override:
+                    model.config.rectified_flow_inference_config.scheduler_type = sampler_override
+                    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'
diff --git a/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py b/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py
@@ -269,6 +269,52 @@ def init_weights(self, buffer_device: torch.device | None):
 
         self.language_model.init_weights(buffer_device=buffer_device)
 
+    def offload_module_groups(self) -> dict[str, list[torch.nn.Module]]:
+        """Partition the network into the reasoner vs generator CPU-offload groups.
+
+        The MoT interleaves both pathways per layer, so each group gathers the
+        per-layer understanding (reasoner) vs generation submodules — plus the
+        generation-side diffusion encode/decode heads, which the reasoner prefill
+        never touches. Returns references to the existing submodule objects (no
+        module-tree changes); the two groups are disjoint, as the offload manager
+        and the two-phase materialization both require. Single source of truth for
+        both the load-time CPU materialization and the runtime ``OffloadPipeline``.
+        """
+        lm = self.language_model.model  # Qwen3VL(Moe)TextModel
+        reasoner: list[torch.nn.Module] = [lm.embed_tokens, lm.norm]
+        generator: list[torch.nn.Module] = [lm.norm_moe_gen]
+        for head_name in ("time_embedder", "vae2llm", "llm2vae", "action2llm", "llm2action", "sound2llm", "llm2sound"):
+            head = getattr(self, head_name, None)
+            if isinstance(head, torch.nn.Module):
+                generator.append(head)
+        for raw_layer in lm.layers:
+            # Unwrap torch.compile's OptimizedModule so we reference the real submodules.
+            layer = getattr(raw_layer, "_orig_mod", raw_layer)
+            attn = layer.self_attn
+            reasoner += [
+                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 += [
+                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,
+            ]
+        return {"reasoner": reasoner, "generator": generator}
+
     def generate_reasoner_text(
         self,
         input_ids: torch.Tensor,
diff --git a/cosmos_framework/model/vfm/omni_mot_model.py b/cosmos_framework/model/vfm/omni_mot_model.py
@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import collections
+import contextvars
 import time
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, Mapping, Optional, Tuple
@@ -59,6 +60,66 @@
 from cosmos_framework.utils.vfm.parallelism import ParallelDims
 
 
+# Names of the network offload groups (from ``Cosmos3VFMNetwork.offload_module_groups``)
+# whose weights should be materialized directly on CPU at load time (two-phase
+# materialization), so they never occupy GPU memory. Set by the inference layer via
+# ``cpu_offload_materialization`` around model construction; read inside ``build_net``.
+_CPU_OFFLOAD_NET_PARTS: contextvars.ContextVar[tuple[str, ...]] = contextvars.ContextVar(
+    "_cpu_offload_net_parts", default=()
+)
+
+
+@contextmanager
+def cpu_offload_materialization(net_parts: tuple[str, ...]):
+    """Materialize the named network offload groups on CPU during model construction.
+
+    Wrap model loading (``from_pretrained_dcp`` / ``load_model_from_checkpoint``) with
+    this so the listed groups (e.g. ``("reasoner", "generator")``) are built directly on
+    CPU and the checkpoint shards load into CPU tensors — they never touch the GPU. Empty
+    ``net_parts`` is a no-op (default joint materialization, unchanged).
+    """
+    token = _CPU_OFFLOAD_NET_PARTS.set(tuple(net_parts))
+    try:
+        yield
+    finally:
+        _CPU_OFFLOAD_NET_PARTS.reset(token)
+
+
+def _offloaded_tensor_ids(net: torch.nn.Module, net_parts: tuple[str, ...]) -> set[int]:
+    """Collect ``id()`` of every parameter/buffer belonging to the offloaded groups."""
+    groups = net.offload_module_groups()
+    ids: set[int] = set()
+    for part in net_parts:
+        for module in groups.get(part, []):
+            for tensor in (*module.parameters(recurse=True), *module.buffers(recurse=True)):
+                ids.add(id(tensor))
+    return ids
+
+
+def _materialize_meta_tensors(net: torch.nn.Module, device: torch.device, skip_ids: set[int] | None) -> None:
+    """Materialize every still-``meta`` parameter/buffer on ``device`` (skipping ``skip_ids``).
+
+    Allocates real empty tensors in place. ``skip_ids`` leaves the offloaded tensors on
+    ``meta`` for the first (GPU) pass so they sidestep the wasted random init, then a
+    second pass with ``skip_ids=None`` lands them on CPU.
+    """
+    for module in net.modules():
+        for name, param in list(module._parameters.items()):
+            if param is None or not param.is_meta:
+                continue
+            if skip_ids is not None and id(param) in skip_ids:
+                continue
+            module._parameters[name] = torch.nn.Parameter(
+                torch.empty_like(param, device=device), requires_grad=param.requires_grad
+            )
+        for name, buffer in list(module._buffers.items()):
+            if buffer is None or not buffer.is_meta:
+                continue
+            if skip_ids is not None and id(buffer) in skip_ids:
+                continue
+            module._buffers[name] = torch.empty_like(buffer, device=device)
+
+
 class OmniMoTModel(ImaginaireModel):
     """
     Mixture of Transformers (MoT) model to be trained with the flow matching objective
@@ -241,12 +302,27 @@ def build_net(self, dtype: torch.dtype):
 
         with misc.timer("meta to cuda and broadcast model states"):
             net = net.to(dtype=dtype)
-            net.to_empty(device=DEVICE)
+            cpu_offload_net_parts = _CPU_OFFLOAD_NET_PARTS.get() if DEVICE == Device.CUDA else ()
+            if cpu_offload_net_parts:
+                # Single-GPU CPU offloading: materialize only the non-offloaded modules on
+                # the GPU and keep the offloaded towers on ``meta``, so ``init_weights``
+                # skips their random init — pure waste here (they come entirely from the
+                # checkpoint and have no non-persistent buffers) and crippling on CPU.
+                _materialize_meta_tensors(
+                    net, torch.device("cuda"), skip_ids=_offloaded_tensor_ids(net, cpu_offload_net_parts)
+                )
+            else:
+                net.to_empty(device=DEVICE)
+
+            # Weight init is only needed on CUDA (CPU/meta are for checkpoint conversion and
+            # smoke tests). It initializes the non-offloaded modules and their non-persistent
+            # buffers (e.g. RoPE); the offloaded towers stay on ``meta`` (init no-ops there).
             if DEVICE == Device.CUDA:
-                # Weight initialization is not needed for other devices (cpu,
-                # meta), since they are only for checkpoint conversion and smoke
-                # tests.
                 net.init_weights(buffer_device=DEVICE)
+                if cpu_offload_net_parts:
+                    # Land the offloaded towers on CPU; ``dcp.load`` fills them next, so they
+                    # never occupy GPU memory.
+                    _materialize_meta_tensors(net, torch.device("cpu"), skip_ids=None)
                 if getattr(self.config, "lora_enabled", False):
                     self._init_lora_weights_post_materialization(net)
 
