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

[rwagwani]

Summary

Checkpoint loading copies each tensor from mmap-backed safetensors storage directly into a CUDA target, so the host-to-device path handles a memory-mapped page fault per tensor. On Grace (GB300) this dominates wall-clock: a Cosmos3-Super reasoner load takes 1,184 s with zero disk I/O — the data is already in page cache and ~1.3 of 72 cores are busy.

Materialising each tensor into ordinary (anonymous) host memory before the copy fixes it:

GB300, single process	before	after
Cosmos3-Nano reasoner load	285 s	8 s	35.6×
Cosmos3-Super reasoner load	1,184 s	67 s	17.7× †
Full 14-command sweep	4 h 49 m	2 h 51 m	1.68×

† Re-measured on a later boot of the same node as 1,211 s → 83 s (14.6×). The effect reproduces in shape but the exact ratio moves ~20% between boots — read the Super figure as order 15×, not a constant. See "Run-to-run spread" below.

Output is byte-identical on both architectures and all models tested. Decode times are unchanged, confirming only the load path was affected.

The mechanism is one clone(); the rest is the gate helper and comments. The copy is applied conditionally, for the reason set out under "Conditional application" below.

---

Root cause

torch.distributed.checkpoint.hf_storage.HuggingFaceStorageReader._process_read_request, reached from cosmos_framework/inference/model.py::from_pretrained_dcp:

tensor = f.get_slice(req.storage_index.fqn)[slices]   # mmap-backed
target_tensor = planner.resolve_tensor(req).detach()  # already on CUDA
target_tensor.copy_(tensor)                           # per-tensor H2D from file-backed pages

The model is moved to GPU before weights are loaded, so each tensor is an individual H2D copy whose source is a memory-mapped file.

Evidence:

• py-spy: every sample lands in _process_read_request — one hot frame, not diffuse cost
• /proc/<pid>/io: read_bytes: 0 while rchar climbs past 2.59 GB — served entirely from page cache, no disk I/O
• ~1.3 of 72 cores busy, GPU ~3% — serial, GPU idle and waiting

---

Why anonymous host memory, not merely pre-faulting

Micro-benchmark, 4.61 GiB Cosmos3-Super shard, 48 tensors, page-fault cost measured separately and pre-paid, each regime run both first and last to control for ordering:

regime	GB300	A100
From mmap storage → GPU (first)	17.53 ms/t	11.91 ms/t
From mmap storage → GPU (last)	9.00 ms/t	11.59 ms/t
Via pinned staging → GPU	1.73 ms/t	9.99 ms/t
From ordinary RAM → GPU	1.61 ms/t	10.14 ms/t
Via pinned staging, in-RAM source	1.69 ms/t	10.79 ms/t

Two results drive the design:

1. Pre-faulting in place is not sufficient. Even after faults are pre-paid, the mmap source still costs 9.00 ms/t versus 1.61 ms/t from heap. The source must be anonymous memory, not file-backed.
2. Pinning is unnecessary. With a heap source, pageable (1.61) and pinned (1.69) are equivalent — pinning adds cost without benefit.

---

Alternatives measured and rejected

approach	measured result	why rejected
thread_count=16 on the reader	285 s → 246 s (14%)	Parallelism caps at shard count (the queue holds files) and the loop is Python-level, so workers serialise on the GIL — py-spy showed 1–3 of 10 workers active
Bulk load_file() per shard	>575 s, worse than baseline	Adds a full shard materialisation without removing a single H2D transfer; reading was never the cost
Pinned staging buffer	1,184 s → 111 s on Super	Works, but slower than this PR (111 s vs 67 s) and needs a per-thread buffer sized to the largest tensor (1.45 GiB), which multiplies across ranks
Heap materialisation (this PR)	1,184 s → 67 s	Simplest, fastest, no buffer management

The pinned-vs-heap pair on Super is the cleanest evidence, since the two differ in exactly one respect and heap wins 67 s to 111 s.

---

Conditional application — implemented

The first revision of this PR applied the copy unconditionally, and disclosed a cost on x86. On configurations where the mmap path is not slow the copy is pure overhead — measured on 8× A100 (Cosmos3-Super, 8 torchrun ranks), alternating runs: stock 431/433 s vs unconditional 489/489 s, ~13.2% slower. The cause is straightforward: on x86 the per-tensor mmap H2D is already fast, so the copy buys nothing, and with 8 ranks copying concurrently they contend for host memory bandwidth. Single-process Edge and Nano on that node were unaffected.

Rather than leave the policy open, the copy is now conditional:

COSMOS_MATERIALIZE_CHECKPOINT unset  -> on for aarch64/arm64, off elsewhere
COSMOS_MATERIALIZE_CHECKPOINT=1      -> force on
COSMOS_MATERIALIZE_CHECKPOINT=0      -> force off

Why arch is the default rather than the condition. The real condition is "does this platform have a slow file-backed mmap H2D path". aarch64 is a proxy for that, and a proxy hardcoded with no escape hatch is an unfalsifiable claim in code — wrong for any future x86 with a slow path, or any ARM without one. Making it a default keeps the heuristic correctable without a code change.

Verified on both architectures

Gate resolution, same build on both nodes:

node	platform.machine()	unset	=1	=0
A100	x86_64	off	on	off
GB300	aarch64	on	on	off

x86 — 8× A100, Cosmos3-Super reasoner, 8 ranks, alternating runs, same harness that produced the original 431/433 vs 489/489 measurement:

arm	wall
stock	434 s / 431 s
gated, default (off)	429 s / 430 s
gated, forced on	490 s

The regression is gone at the default, and the forced arm reproduces it — which also confirms the gate is live code rather than something that never fires.

aarch64 — 1× GB300, Cosmos3-Super reasoner, single process:

arm	wall	load window
gated, default (on)	124 s	83 s
gated, forced off	1252 s	1211 s

14.6× on the load window, with the default firing automatically and no configuration required.

All outputs byte-identical within each node: A100 9c81833184f7… across all five runs, GB300 d520e852… across both, each matching that node's own stock baseline.

Measurement scope: all timings above come from the diffusers reader path (_DiffusersHuggingFaceStorageReader), which is what the Cosmos3 checkpoints load through. The plain HF reader inherits the same mixin and therefore the same behaviour by construction, but was not separately benchmarked. Both readers were verified to resolve the gate identically (unset/=1/=0 → off/on/off on x86_64).

On .contiguous().clone(): a bare .clone() is very likely equivalent here and would avoid a redundant second copy when a slice is non-contiguous. We propose the .contiguous() form only because it is what every number above was measured against; switching would mean quoting timings for code we did not run. Happy to change it.

Run-to-run spread, disclosed: the headline 1,184 s → 67 s (17.7×) and this re-verification's 1,211 s → 83 s (14.6×) are the same measurement taken on different boots of the same node. The effect reproduces in shape and order of magnitude, but the exact ratio moves by roughly 20% between boots — the headline figure should be read as "order 15×", not as a precise constant.

Alternatives considered for the condition itself

condition	why not
Gate on world_size == 1	Rank contention is the actual mechanism, but this would silently drop the benefit on multi-GPU Grace (GH200, multi-node), a configuration we have not measured
Runtime probe of the copy path	Most principled, but first-tensor timing is noisy and needs a warm-up, roughly doubling the size of a patch whose main virtue is being reviewable at a glance
Hardcoded arch with no override	Same default, but no escape hatch and no way to correct the heuristic without a code change

Happy to drop the env var, rename it, invert the default, or move to a probe — whichever is preferred. The measurements above should make any of those a small edit rather than a re-investigation.

---

Correctness

Byte-identical output against stock, every configuration tested:

configuration	md5	matches stock
GB300, Cosmos3-Nano reasoner	560ce35de53c6c4a6435db615d861d40	✅
GB300, Cosmos3-Super reasoner	d520e852a059ed52c6d42953f787e70a	✅
A100, Cosmos3-Nano reasoner	560ce35de53c6c4a6435db615d861d40	✅
A100, Cosmos3-Super reasoner, 8 ranks	9c81833184f7a7067c7f59b3326a91b2	✅

Full-sweep comparison on GB300: 62 of 63 output files bit-identical. The single exception (transfer_multi_control) is under investigation and is not attributable to this change — all six transfer specs come from one command and therefore one model load, and the other five are bit-identical; had the patch altered any weight, all six would have moved. It has since differed a third time under a second, independent patch sharing no code with this one, which points to nondeterminism inherent to that spec rather than to either change.

Peak host memory checked under the heaviest configuration available (Cosmos3-Super, 30 shards, 8 ranks): 73 GB of 1771 GB, no pressure.

---

Upstream note

The same pattern exists in torch/distributed/checkpoint/hf_storage.py::_process_read_request and affects any project loading a DCP/HF checkpoint onto GPU. This PR fixes it within cosmos-framework's subclass so the benefit does not depend on a torch release; a corresponding upstream report is worth filing separately.

---

Environment

	GB300 node	A100 node
GPU	1× GB300 (sm_100)	8× A100-SXM4-80GB (sm_80)
CPU	Grace Neoverse-V2, 72c @ 3.3 GHz	AMD EPYC 7V12, 96c @ 2.4 GHz
RAM	486 GB	1771 GB
Storage	ext4, local NVMe, 5.6 GB/s O_DIRECT	overlay, local NVMe, 1.2 GB/s
torch	2.10.0+cu130	2.10.0+cu130

Benchmarks pinned to cosmos-framework 5e67049 (v1.2.2) and cosmos 0299468, identical checkpoint SHAs on both nodes. All comparison results were produced on unmodified code; the patch was measured separately and reverted.

Rebase note. This branch has since been rebased onto main at ee58e41 to satisfy the up-to-date-branch requirement. The rebase was clean and the three intervening commits (#148, #153, #159) touch no file this PR modifies. The one nearby change is model/generator/utils/safetensors_loader.py (#148), which is the VLM safetensors loader — a different code path from the DCP HuggingFaceStorageReader patched here, and its changes concern shard/replica semantics, not the H2D copy.

The x86 gate A/B was nonetheless re-run in full on the rebased branch, since the 13.2% regression and its removal are the load-bearing claims for the gate. Arms selected by git checkout of the real commits (stock ee58e41, patched 3d8c469), same 8-rank harness as before:

arm	wall	md5
stock	434 s / 431 s	9c81833184f7
gated, default (off)	434 s / 433 s	9c81833184f7
gated, forced on	492 s (+13.4%)	9c81833184f7

Unchanged from the pre-rebase measurement: the default is indistinguishable from stock, the forced arm reproduces the regression, and all five runs are byte-identical to the stock baseline. The GB300 numbers were not re-run post-rebase (separate node); the patched code path is untouched by the rebase, but say the word if you want them repeated.

Storage, RAM speed, thread configuration, aarch64 vectorisation and the bulk safetensors path were each ruled out by measurement before arriving at this diagnosis.

---

Thanks to @NVIDIA research (Liang Feng) for reviewing an earlier version of this analysis and correcting the mechanism — an initial benchmark had measured a cold page cache and misattributed the cost to pageable-vs-pinned transfers.

[lfengad]

Should be all default off? Since not necessarily would benefit in all aarch64/arm64 case? Still think this is highly related with the storage system.

[rwagwani]

Fair challenge — and the published table genuinely could not answer it, since the two nodes differ
in architecture, filesystem and host-GPU interconnect at once.

I ran the controlled version: one node, one shard, filesystem as the only variable, then the same
sweep on x86 for a 2×2. Full numbers in
#150 (comment).

Short version: on aarch64 the penalty holds on every filesystem — 9.03×–13.42× — including
tmpfs, which has no block-device fault handler behind it, and overlay, which is what the x86 node
uses. It never collapses. On x86 it is absent everywhere (0.21×–0.28×; mmap is ~4× faster).

So architecture flips the sign of the effect and the filesystem only scales its magnitude — a
1.49× spread across five filesystems against a ~40× gap between architectures. Storage is
second-order, not causal.

On "not necessarily would benefit in all aarch64/arm64 case" — you're right, and the sweep does not
fix that. It cannot separate architecture from host-GPU interconnect (C2C vs PCIe), because both
boxes differ in both. aarch64 is a better-supported proxy now, but still a proxy. That is why it
is a default with an env override rather than a hard condition, and I'd argue for keeping that
shape precisely because of the uncertainty you're pointing at.

[lfengad]

Could have default 0 in the env.get?

[rwagwani]

That would make it opt-in everywhere — with a "0" default, override is never None, so the
platform.machine() branch below becomes unreachable. Worth being explicit that the change deletes
the arch heuristic rather than just adjusting it.

I'd argue against, but only on the strength of the new measurements
(#150 (comment)): on aarch64 the mmap
path is 9–13× slower than staging on every filesystem tested, and a Cosmos3-Super reasoner load
goes 1211 s → 83 s. Default-off means every Grace user pays a ~20-minute load with nothing
indicating a fix exists — the failure mode is silent, and the people most affected are least likely
to know the env var is there.

Default-on for aarch64 with COSMOS_MATERIALIZE_CHECKPOINT=0 as the escape hatch keeps the
regression risk bounded (x86 is unaffected: the gate resolves off, verified at 434/433 s against
434/431 s stock) while making the fix reachable without configuration.

That said — this is a one-line change and I'm happy to take it if you'd rather ship opt-in. The
measurements make it reversible in either direction, so it's a policy call, not a technical one.
Your call and I'll follow it.

[rwagwani]

Filesystem sweep: the gate condition, measured

@lfengad — your objection that this is "highly related with the storage system" was fair, and the
published table could not answer it. GB300 and the A100 node differ in architecture, filesystem
and host-GPU interconnect simultaneously (aarch64/ext4-on-NVMe/C2C vs x86_64/overlayfs/PCIe), so
nothing in it separated the three. read_bytes: 0 ruled out disk bandwidth, but not the
filesystem's mmap fault path — minor faults on page-cache-resident pages still run through it, and
overlayfs and ext4 do not behave identically there.

So I ran the controlled version: one node, one shard, filesystem as the only variable — then
the same sweep on x86 for a 2×2. (Attribution, since it matters for how much weight to give this:
the aarch64 arm was executed on the GB300 node and the x86 arm on the A100 node; I have since
regenerated both tables from the raw per-repeat JSON rather than from either node's rendered
summary, and every figure below is reproduced from that raw data.)

Q1 — filesystem, or architecture?

The discriminating quantity is the penalty ratio: mmap source vs heap source measured on the
same filesystem, so storage is held constant within each comparison.

GB300 — aarch64, 5 filesystems

filesystem	mmap pre-faulted	heap .clone()	ratio (median)	ratio envelope
ext4 (native NVMe)	3.57 (2.59–4.76)	0.34 (0.28–0.42)	10.55×	6.2×–16.9×
tmpfs	3.41 (2.87–3.86)	0.33 (0.27–0.35)	10.29×	8.2×–14.2×
xfs (loopback)	4.02 (3.08–4.49)	0.30 (0.24–0.30)	13.42×	10.2×–18.9×
overlay, lower layer	3.00 (2.26–4.71)	0.33 (0.26–0.34)	9.03×	6.7×–17.8×
overlay, upper layer	3.34 (3.26–3.90)	0.33 (0.27–0.34)	10.06×	9.6×–14.5×

Medians alone would overstate what 5 repeats support, so the envelope column takes the per-repeat
extremes. Across all five filesystems the ratio envelope is 6.2×–18.9×.

A100 — x86_64, 3 filesystems (every writable path on that pod is overlay, so the layered
overlay variants could not be constructed there)

filesystem	mmap pre-faulted	heap .clone()	ratio (median)	ratio envelope
overlay (native)	1.06 (0.97–1.06)	5.00 (4.85–5.10)	0.21×	0.19×–0.22×
tmpfs	1.38 (1.38–1.42)	4.91 (4.74–5.12)	0.28×	0.27×–0.30×
xfs (loopback)	1.33 (1.30–1.34)	5.28 (5.15–5.34)	0.25×	0.24×–0.26×

The penalty never collapses on aarch64 — not on tmpfs, which has no block-device fault handler
behind it at all, and not on overlay, which is precisely what the x86 node uses. And on x86 it is
absent on every filesystem: the mmap source is ~4× faster there, which independently
reproduces the 13.2% 8-rank regression this PR already discloses, now visible per tensor.

Architecture flips the sign of the effect; the filesystem only scales its magnitude — 1.49×
spread across the five filesystems on aarch64, 1.33× on x86, against a ~40× gap between the two
architectures. Storage is second-order by roughly 25×, not causal.

What this does not settle, and I would rather say so than let the result be read as more than
it is: architecture and host-GPU interconnect remain confounded. Both boxes differ in both, so
this cannot distinguish "Grace CPU" from "C2C rather than PCIe". aarch64 is still a proxy — a
much better supported one, but a proxy. That is exactly why it is a default with an env override
rather than a hard condition, and I would keep it that way.

Q2 — is pre-faulting in place sufficient?

No. On aarch64, after every fault is pre-paid, the mmap source still costs 3.0–4.0 ms/tensor
against 0.30–0.34 ms/tensor from the heap — a ~10× gap that survives on all five filesystems. The
source has to be anonymous memory; paying the faults is not enough.

Q3 — .contiguous().clone() vs bare .clone()

Equivalent — within ±2% on both architectures, on every filesystem. The PR body hedged that a
bare .clone() was "very likely equivalent"; it is now measured. I'll switch to the bare
.clone(), which drops a redundant copy when a slice is non-contiguous and invalidates none of
the numbers above.

Q4 — a correction to this PR

Volunteering this because it cuts against what the PR currently says.

Measured per tensor, pinned staging beats the heap copy on both architectures — 15–23% faster
on GB300, 20–31% on A100, tight and reproducible across repeats. That contradicts this line in the
PR body:

Pinning is unnecessary. With a heap source, pageable (1.61) and pinned (1.69) are equivalent
— pinning adds cost without benefit.

That comparison started from a source already in RAM, so it excluded the staging copy. Comparing
like with like — both staging from the mmap source, which is what the code actually does — pinned
wins. I'm withdrawing that claim from the body.

The decision to reject pinned staging still stands, but on the remaining evidence rather than that
one: the end-to-end Super load runs the other way (111 s pinned vs 67 s heap), and a pinned buffer
must be sized to the largest tensor (1.45 GiB) per reader thread and per rank. I want to be
precise about the limits of the new measurement, though: the harness allocates one pinned buffer,
once, outside the timed region, so it excludes allocation cost entirely. Buffer allocation and
sizing plausibly explain the end-to-end reversal, but this sweep does not decompose it and I am not
going to assert it as measured. What is established: pinned wins the steady-state transfer, heap
wins the full load, and the two measure different things. If allocation can be amortised across a
whole load, pinned staging is worth revisiting as a follow-up.

Method and provenance

• Checkpoint: Cosmos3-Edge transformer/diffusion_pytorch_model-00001-of-00002.safetensors,
  snapshot ff48d221, 4.66 GiB, 425 tensors — not Super, which is not present on the GB300
  node. Byte-identical on both boxes (HF blobs are content-addressed: sha256 f74b228d…), so the
  cross-architecture rows are genuinely comparable. The mechanism is checkpoint-independent, but
  the PR's headline load figures were taken on Super and these were not.
• Every file is read end to end before timing, so all regimes run cache-warm; read_bytes deltas
  confirm no timed window touched the block device.
• Source and staged mappings are classified per regime from /proc/self/maps, verifying the mmap
  regimes really read file-backed memory and the heap regimes anonymous memory — the premise the
  original write-up asserted rather than checked.
• CUDA destinations pre-allocated; warm-up pass before the first measured regime; regime order
  alternated between repeats; median of 5.
• One presentation caveat, disclosed rather than smoothed: mmap, first touch on the two overlay
  variants is bimodal, not noisy (overlay_upper: 74.9, 77.7 │ 135.1, 138.2, 138.8 ms/tensor).
  A min–max range there would read as a centred estimate and imply a filesystem effect that is a
  sampling artifact. mmap, pre-faulted on overlay_lower is flagged by the same test
  (2.26–4.71). That is why the Q1 table above carries a per-repeat envelope rather than medians
  alone: on the most conservative reading the aarch64 ratio floor is 6.2×, not 9.03×. The
  conclusion is unchanged — 6.2× against an x86 ceiling of 0.30× is still a ~20× separation — but
  the medians should not be read as precise.

Happy to publish the raw JSON or the harness if useful.

---

Edited after posting: added per-repeat ratio envelopes to both Q1 tables, corrected a statement
that the mmap_prefaulted ranges were uniformly tight (one is not — overlay_lower spans
2.26–4.71), and made explicit which node executed which arm. The Q1–Q4 conclusions are unchanged;
all figures were re-derived from the raw JSON. Shard identity independently confirmed by
sha256sum on both nodes: f74b228d29f844a58bef266f3afc2d695fdc7e00f0d18b618f5889966586891b.