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

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:

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

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:

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:

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.