fix(inference): materialise checkpoint tensors in host memory before H2D copy (17.7x faster load on GB300)

cosmos_framework/inference/model.py

@@ -270,7 +270,84 @@ def _normalize_diffusers_target_key(name: str) -> str:
     return name.removeprefix("model.net.").replace("_orig_mod.", "").replace("_checkpoint_wrapped_module.", "")
 
 
-class _DiffusersHuggingFaceStorageReader(HuggingFaceStorageReader):
+class _MmapSafeReadMixin:
+    """Materialize each safetensors slice into anonymous RAM before the H2D copy.
+
+    ``HuggingFaceStorageReader._process_read_request`` copies tensors straight from the
+    ``mmap``-backed safetensors slice onto the GPU, so the transfer handles a page fault
+    per tensor. On Grace that dominates checkpoint load: a Cosmos3-Super reasoner load
+    spends ~1200s here with zero disk I/O, the data already resident in page cache.
+
+    The copy is applied conditionally -- see ``_materialize_enabled`` for why.
+    """
+
+    _materialize_cache = None
+
+    @classmethod
+    def _materialize_enabled(cls) -> bool:
+        """Whether to stage tensors through anonymous host memory before the H2D copy.
+
+        Defaults on for aarch64 and off elsewhere, overridable in either direction with
+        COSMOS_MATERIALIZE_CHECKPOINT=1/0.
+
+        The staging copy removes a slow file-backed mmap H2D path on Grace, but where that
+        path is already fast the copy is pure overhead: on 8x A100 (Cosmos3-Super, 8 ranks)
+        applying it unconditionally measured ~13.2% slower, because concurrent ranks contend
+        for host memory bandwidth. The architecture is only a proxy for "is the mmap H2D path
+        slow here", so it is used as a default rather than as a hard condition.
+        """
+        if cls._materialize_cache is None:
+            import os
+            import platform
+
+            override = os.environ.get("COSMOS_MATERIALIZE_CHECKPOINT")
+            if override is not None:
+                cls._materialize_cache = override.strip().lower() not in (
+                    "0",
+                    "false",
+                    "no",
+                    "off",
+                    "",
+                )
+            else:
+                cls._materialize_cache = platform.machine().lower() in ("aarch64", "arm64")
+        return cls._materialize_cache
+
+    def _process_read_request(self, f, req, planner) -> None:  # noqa: D102
+        slices = tuple(slice(offset, offset + length) for offset, length in zip(req.storage_offsets, req.lengths))
+        tensor = f.get_slice(req.storage_index.fqn)[slices]
+        target_tensor = planner.resolve_tensor(req).detach()
+
+        if target_tensor.size() != tensor.size():
+            raise AssertionError(f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}")
+
+        if target_tensor.is_cuda and self._materialize_enabled():
+            # Materialise into anonymous host memory before the H2D copy.
+            #
+            # The base implementation copies straight from mmap-backed safetensors
+            # storage into a CUDA tensor, so the transfer handles a page fault per
+            # tensor. On Grace this dominates load time (Cosmos3-Super load: ~1200s
+            # down to ~70-85s with this change, measured across two boots; zero disk
+            # I/O either way -- the data is already in page cache).
+            #
+            # Pre-faulting in place is not sufficient: measured on a 4.61 GiB shard,
+            # an mmap source still costs 9.00 ms/tensor after faults are pre-paid,
+            # versus 1.61 ms/tensor from the heap. Pinning the staging buffer is not
+            # necessary either -- with a heap source, pageable (1.61) and pinned
+            # (1.69) are equivalent. A single clone() does both jobs: the memcpy
+            # faults the pages sequentially and leaves the data resident and
+            # anonymous. It is freed as soon as the H2D copy completes.
+            tensor = tensor.contiguous().clone()
+
+        target_tensor.copy_(tensor)
+        planner.commit_tensor(req, target_tensor)
+
+
+class _MmapSafeHuggingFaceStorageReader(_MmapSafeReadMixin, HuggingFaceStorageReader):
+    """Plain HF safetensors reader with the mmap-H2D staging copy."""
+
+
+class _DiffusersHuggingFaceStorageReader(_MmapSafeReadMixin, HuggingFaceStorageReader):
     """Hugging Face safetensors reader that follows diffusers' root weight map."""
 
     def __init__(self, checkpoint_path: Path) -> None:
@@ -561,7 +638,7 @@ def from_pretrained_dcp(
                     return model
                 state_dict = get_model_state_dict(model)
                 _raise_on_missing_vision_keys(checkpoint_path, state_dict)
-                storage_reader = HuggingFaceStorageReader(str(checkpoint_path))
+                storage_reader = _MmapSafeHuggingFaceStorageReader(str(checkpoint_path))
             case _:
                 assert_never(checkpoint_type)
         dcp.load(state_dict=state_dict, storage_reader=storage_reader)