[ltx-core] SafetensorsStateDictLoader uses safetensors' default mmap backend; `pread` is ~2x faster and far more consistent on MPS

[dmlsrc]

Summary

SafetensorsStateDictLoader.load opens every checkpoint shard with a bare
safetensors.safe_open(...), so it takes safetensors' default mmap backend on
every platform. safetensors 0.8 added a pread backend that skips the file
mapping. On Apple Silicon, switching only that argument measures about 2x faster
on a warm page cache, and the gap widens at real checkpoint sizes where the
mapping path also becomes very inconsistent run to run.

This is a throughput change only. I measured the allocation shape both ways and
it is identical, so this does not reduce memory use -- see "What this does not
change" below.

packages/ltx-core/src/ltx_core/loader/sft_loader.py:30

with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:

There is no backend= argument anywhere under ltx_core/loader/.

Why it is worth the one line

Every weight load in a generation goes through this method: the Gemma text
encoder shards, the transformer, the spatial upsampler, the video decoder, and
the audio decoder plus vocoder. And because DiffusionStage builds the
transformer on each call and frees it on exit, a two-stage generation reads the
43 GiB monolith twice.

On a full distilled two-stage run, changing only this argument took the
generation from 216 s to 82 s -- 2.6x, about 133 s saved per generation.
Details under "End-to-end" below.

Environment

• Repo: a95ab856bf29407b6b066ede0abe1846050db56c (Automated PR - 2026-08-25)
• Hardware: Apple M1 Max, 64 GB unified memory
• macOS 26.6.2, Python 3.14.7
• torch 2.13.0, safetensors 0.8.0
• Checkpoint: ltx-2.3-22b-distilled-1.1.safetensors (43 GiB)

Reproduction

Time the loader's own call shape, varying only backend. Each backend runs in a
fresh process so page-cache state and RSS high-water marks do not leak between
measurements, and the page cache is warmed first so this measures mapping and
API cost rather than disk.

import struct, json, time
from pathlib import Path
import torch
from safetensors import safe_open

BACKEND = "mmap"      # or "pread"
PATH = Path("/path/to/ltx-2.3-22b-distilled-1.1.safetensors")
TARGET = 6 * 2**30

with PATH.open("rb") as f:
    (n,) = struct.unpack("<Q", f.read(8))
    header = json.loads(f.read(n))
header.pop("__metadata__", None)

names, nbytes = [], 0
for name, meta in sorted(header.items(), key=lambda kv: kv[1]["data_offsets"][0]):
    if meta["dtype"] not in ("BF16", "F16", "F32"):
        continue
    if meta["data_offsets"][1] > TARGET:
        break
    names.append(name)
    nbytes = meta["data_offsets"][1]

with PATH.open("rb") as f:                      # warm the page cache
    left = nbytes
    while left > 0 and (b := f.read(min(64 << 20, left))):
        left -= len(b)

t0 = time.monotonic()
sd = {}
with safe_open(str(PATH), framework="pt", device="mps", backend=BACKEND) as f:
    for name in names:
        sd[name] = f.get_tensor(name)
torch.mps.synchronize()
dt = time.monotonic() - t0
print(BACKEND, f"{dt:.2f}s", f"{nbytes / dt / 2**30:.2f} GiB/s")

Results

Warm page cache, fresh process per backend, same 1062-tensor 5.99 GiB subset:

Backend	Time	Throughput
mmap (current)	1.92 s	3.12 GiB/s
pread	0.89 s	6.76 GiB/s

At 30 GiB, where the page cache cannot hold the data alongside the resident
tensors, two runs of each:

Backend	Run 1	Run 2
mmap (current)	16.00 s (1.88 GiB/s)	40.33 s (0.74 GiB/s)
pread	10.45 s (2.87 GiB/s)	10.26 s (2.92 GiB/s)

The mapping path is not only slower, it varied 2.5x between two identical runs.
pread was within 2% of itself.

End-to-end

A full distilled two-stage generation at 448x256x25, seed 42, same prompt, with
SafetensorsStateDictLoader.load patched to pass backend="pread" and nothing
else changed. The transformer is rebuilt per stage, which is stock
DiffusionStage behavior. Arms were interleaved stock, pread, stock, pread so
page-cache drift hits both equally:

Phase	stock (mmap)	pread	Delta
prompt_encode	49.5 / 42.0 s	15.9 / 15.8 s	-29.9 s
denoise_stage_1	87.4 / 86.7 s	33.8 / 34.1 s	-53.0 s
latent_upsample	2.2 / 2.4 s	0.9 / 1.0 s	-1.4 s
denoise_stage_2	75.4 / 76.8 s	27.2 / 27.5 s	-48.8 s
decode_and_encode	1.1 / 1.1 s	1.1 / 1.1 s	0.0 s
total	219.5 / 213.1 s	81.6 / 82.1 s	-134.5 s

The denoise phases include the transformer build, which is where most of the
saving is. prompt_encode moves because the Gemma shards load through the same
method -- one shard alone is 24.6 GiB, read at 5.19 GiB/s on the pread path.
decode_and_encode does no significant loading and is unchanged, which is the
control.

This was a deliberately small generation so that loading, not denoising,
dominates. At higher resolutions the ratio shrinks but the absolute saving does
not: it is a fixed cost per weight load.

What this does not change

I expected pread to also improve the allocation shape and it does not. Loading
the same subset both ways gives:

	mmap	pread
Tensors	1062	1062
Distinct storage pointers	1062	1062
Process RSS delta	6.00 GiB	6.00 GiB
MPS driver allocation delta	5.99 GiB	5.99 GiB
Storage writable in place	yes	yes

Every tensor gets its own MPS allocation either way. The end-to-end runs agree:
peak process footprint was 43.7 / 43.7 GB on pread against 43.8 / 44.1 GB on
the mapping path, and peak MPS driver allocation 41.9 / 43.0 GB against 43.0 /
43.1 GB. So this is purely a read-throughput change, not a memory one.

What I did not test

• CUDA and Windows. The measurements above are MPS only, which is why the
  suggested diff scopes the change to MPS rather than switching the default
  everywhere. pread may well help on other platforms, but I have not measured
  it and would not want to claim it.
• torch 2.14, which was not released when I measured. It reworks several MPS
  paths, though none that I would expect to touch this one.

Suggested change

--- a/packages/ltx-core/src/ltx_core/loader/sft_loader.py
+++ b/packages/ltx-core/src/ltx_core/loader/sft_loader.py
@@
         for shard_path in model_paths:
-            with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:
+            # safetensors >= 0.8 exposes a pread backend that skips the file mapping.
+            # On MPS the mapping path is about 2x slower on a warm cache and varies
+            # by 2.5x run to run at checkpoint scale; pread is stable. Left as the
+            # default elsewhere since only MPS was measured.
+            open_kwargs = {"backend": "pread"} if device.type == "mps" else {}
+            with safetensors.safe_open(
+                shard_path, framework="pt", device=str(device), **open_kwargs
+            ) as f:

backend is keyword-only and was added in safetensors 0.8, while ltx-core
currently declares safetensors unpinned. So this needs either a floor bump:

-    "safetensors",
+    "safetensors>=0.8",

or a feature check if you would rather not raise the floor:

_HAS_BACKEND = "backend" in inspect.signature(safetensors.safe_open).parameters

Happy to open a PR for whichever shape you prefer, or to leave it here if it is
easier to take internally.

Note for reviewers, unrelated to this diff

load accumulates the entire state dict in sd before returning it, so the
full checkpoint is resident as a dict before the model consumes it. On the 43
GiB monolith that is a large transient independent of which backend reads it. I
have not investigated whether the consumer then holds a second copy, and did not
want to bundle a speculative memory question into a measured throughput fix -- I
can open that separately if it would be useful.