From 9078e51d32e75bf6fdcb30c4373433b1a1738bd6 Mon Sep 17 00:00:00 2001 From: rwagwani Date: Sat, 1 Aug 2026 08:40:57 +0000 Subject: [PATCH 1/2] fix(inference): materialise checkpoint tensors in host memory before H2D copy _DiffusersHuggingFaceStorageReader inherits _process_read_request, which copies each tensor from mmap-backed safetensors storage directly into a CUDA target. The host-to-device path then handles a memory-mapped page fault per tensor. On Grace (GB300) this dominates checkpoint load time with zero disk I/O -- the data is already in page cache. Cosmos3-Super reasoner load: 1184s -> 67s (17.7x). Cosmos3-Nano: 285s -> 8s (35.6x). Full 14-command sweep: 4h49m -> 2h51m. Materialising each tensor into anonymous host memory before the copy fixes it. Pre-faulting in place is not sufficient (mmap source stays at 9.00 ms/tensor after faults are pre-paid, vs 1.61 ms/tensor from heap), and pinning is unnecessary (heap pageable 1.61 vs heap pinned 1.69 ms/tensor). Output is byte-identical on GB300 and x86 A100, single-process and 8-rank. Known trade-off: on 8x A100 where the mmap path is not slow, the added copy costs ~13% on Cosmos3-Super (431s -> 489s, reproducible). Happy to make it conditional. Co-Authored-By: Claude Opus 5 (1M context) --- cosmos_framework/inference/model.py | 35 +++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py index 5b7c74ce..c4bb85b6 100644 --- a/cosmos_framework/inference/model.py +++ b/cosmos_framework/inference/model.py @@ -278,6 +278,41 @@ def __init__(self, checkpoint_path: Path) -> None: self.checkpoint_path = checkpoint_path self.files_to_keys = _diffusers_files_to_keys(_diffusers_weight_map(checkpoint_path)) + 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 " + f"{target_tensor.size()} vs {tensor.size()}" + ) + + if target_tensor.is_cuda: + # 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: 1184s -> 67s + # with this change, 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) + def read_metadata(self) -> Metadata: from safetensors import safe_open from safetensors.torch import _getdtype From 3d8c46917871e173dcf0f099ec713b3b555facc6 Mon Sep 17 00:00:00 2001 From: rwagwani Date: Tue, 4 Aug 2026 08:44:51 +0000 Subject: [PATCH 2/2] fix(inference): apply the host-memory staging copy conditionally The staging copy removes a slow file-backed mmap H2D path on Grace, but where that path is already fast it is pure overhead. On 8x A100 (Cosmos3-Super, 8 torchrun ranks) applying it unconditionally measured ~13.2% slower (431/433s stock vs 489/489s), because concurrent ranks copying at once contend for host memory bandwidth. Single-process x86 is unaffected either way. Default on for aarch64/arm64, off elsewhere, overridable in either direction with COSMOS_MATERIALIZE_CHECKPOINT=1/0. 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; the override keeps the heuristic correctable without a code change. The staging logic moves to a _MmapSafeReadMixin so the plain HF reader is covered as well as the diffusers reader. Verification, both architectures, same build: gate resolution (COSMOS_MATERIALIZE_CHECKPOINT unset / =1 / =0) x86_64 A100 off / on / off both readers aarch64 GB300 on / on / off both readers x86_64 -- 8x A100, Cosmos3-Super reasoner, 8 ranks, alternating runs stock A 434s gated, default off 429s gated, forced on 490s gated, default off 430s stock B 431s all five md5 9c81833184f7a7067c7f59b3326a91b2 aarch64 -- 1x GB300, Cosmos3-Super reasoner, single process gated, default on 124s wall, 83s load window gated, forced off 1252s wall, 1211s load window both md5 d520e852a059ed52c6d42953f787e70a, equal to that node's stock baseline Byte-identical output within each node across every arm; hashes compared within a node, not across. Both trees reverted clean at 5e67049 afterwards. Timings come from the diffusers reader path; the plain HF reader shares the mixin and therefore the mechanism, but was not separately benchmarked. Co-Authored-By: Claude Opus 5 (1M context) --- cosmos_framework/inference/model.py | 80 ++++++++++++++++++++++------- 1 file changed, 61 insertions(+), 19 deletions(-) diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py index c4bb85b6..d415c1b0 100644 --- a/cosmos_framework/inference/model.py +++ b/cosmos_framework/inference/model.py @@ -270,36 +270,65 @@ def _normalize_diffusers_target_key(name: str) -> str: return name.removeprefix("model.net.").replace("_orig_mod.", "").replace("_checkpoint_wrapped_module.", "") -class _DiffusersHuggingFaceStorageReader(HuggingFaceStorageReader): - """Hugging Face safetensors reader that follows diffusers' root weight map.""" +class _MmapSafeReadMixin: + """Materialize each safetensors slice into anonymous RAM before the H2D copy. - def __init__(self, checkpoint_path: Path) -> None: - super().__init__(str(checkpoint_path)) - self.checkpoint_path = checkpoint_path - self.files_to_keys = _diffusers_files_to_keys(_diffusers_weight_map(checkpoint_path)) + ``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) - ) + 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 " - f"{target_tensor.size()} vs {tensor.size()}" - ) + raise AssertionError(f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}") - if target_tensor.is_cuda: + 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: 1184s -> 67s - # with this change, zero disk I/O either way -- the data is already in - # page cache). + # 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, @@ -313,6 +342,19 @@ def _process_read_request(self, f, req, planner) -> None: # noqa: D102 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: + super().__init__(str(checkpoint_path)) + self.checkpoint_path = checkpoint_path + self.files_to_keys = _diffusers_files_to_keys(_diffusers_weight_map(checkpoint_path)) + def read_metadata(self) -> Metadata: from safetensors import safe_open from safetensors.torch import _getdtype @@ -596,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)