Stream safetensors checkpoints into models

flashdreams/flashdreams/core/checkpoint/load.py

@@ -28,6 +28,7 @@
 import torch
 from huggingface_hub import hf_hub_download, try_to_load_from_cache
 from loguru import logger
+from safetensors import safe_open
 from safetensors.torch import load as load_safetensors
 from safetensors.torch import load_file as load_safetensors_file
 from safetensors.torch import save_file as save_safetensors
@@ -737,12 +738,123 @@ def _load_checkpoint_from_local(
 ) -> dict[str, torch.Tensor]:
     """Load checkpoint from local filesystem."""
     if ext == ".safetensors":
-        with open(path, "rb") as f:
-            return load_safetensors(f.read())
+        return load_safetensors_file(path, device=_safetensors_device(map_location))
     else:
         return torch.load(path, map_location=map_location, weights_only=False)
 
 
+def _stream_safetensors_into_model(
+    model: torch.nn.Module,
+    path: str,
+) -> torch.nn.Module:
+    """Copy a safetensors checkpoint into a model with bounded host residency.
+
+    Args:
+        model: Materialized destination model.
+        path: Local safetensors checkpoint path.
+
+    Returns:
+        The destination model with checkpoint weights loaded.
+
+    Raises:
+        RuntimeError: Checkpoint keys or tensor shapes do not match the model.
+    """
+    model_state = model.state_dict()
+    checkpoint_fd = os.open(path, os.O_RDONLY)
+    evict_interval = 512 * 1024**2
+    bytes_since_evict = 0
+
+    def evict_checkpoint_pages() -> None:
+        if not hasattr(os, "posix_fadvise") or not hasattr(os, "POSIX_FADV_DONTNEED"):
+            return
+        try:
+            os.posix_fadvise(checkpoint_fd, 0, 0, os.POSIX_FADV_DONTNEED)
+        except OSError as exc:
+            logger.warning(f"Could not evict checkpoint page cache for {path}: {exc}")
+
+    try:
+        with safe_open(path, framework="pt", device="cpu", backend="mmap") as source:
+            checkpoint_keys = set(source.keys())
+            model_keys = set(model_state)
+            missing = sorted(model_keys - checkpoint_keys)
+            unexpected = sorted(checkpoint_keys - model_keys)
+            if missing or unexpected:
+                details = []
+                if missing:
+                    details.append(f"Missing key(s): {', '.join(missing[:20])}")
+                if unexpected:
+                    details.append(f"Unexpected key(s): {', '.join(unexpected[:20])}")
+                raise RuntimeError(
+                    f"Checkpoint does not match {type(model).__name__}: "
+                    + "; ".join(details)
+                )
+
+            for name, destination in model_state.items():
+                source_shape = tuple(source.get_slice(name).get_shape())
+                if source_shape != tuple(destination.shape):
+                    raise RuntimeError(
+                        f"Checkpoint tensor {name!r} has shape {source_shape}, "
+                        f"expected {tuple(destination.shape)}"
+                    )
+
+            with torch.no_grad():
+                for name, destination in model_state.items():
+                    tensor = source.get_tensor(name)
+                    destination.copy_(tensor)
+                    bytes_since_evict += tensor.numel() * tensor.element_size()
+                    del tensor
+                    if bytes_since_evict >= evict_interval:
+                        evict_checkpoint_pages()
+                        bytes_since_evict = 0
+    finally:
+        evict_checkpoint_pages()
+        os.close(checkpoint_fd)
+
+    return model
+
+
+def _resolve_streamable_safetensors_path(
+    checkpoint_path: str,
+    *,
+    local_cache_dir: str,
+    checkpoint_min_free_gb: float | None,
+) -> str | None:
+    """Resolve a locally available safetensors file for streaming model loads.
+
+    Args:
+        checkpoint_path: Local path, S3 URI, or Hugging Face URL.
+        local_cache_dir: Directory for S3 and merged-safetensors caches.
+        checkpoint_min_free_gb: Optional Hugging Face cache-space requirement.
+
+    Returns:
+        Local safetensors path, or ``None`` when materialization is still required.
+    """
+    if _is_sharded_safetensors_index_checkpoint(checkpoint_path):
+        if checkpoint_path.startswith("s3://"):
+            return None
+        cache_path = _sharded_safetensors_merge_cache_path(
+            checkpoint_path, local_cache_dir
+        )
+        if os.path.exists(cache_path):
+            logger.info(f"Streaming merged sharded checkpoint from cache: {cache_path}")
+            return cache_path
+        return None
+
+    if _get_checkpoint_extension(checkpoint_path) != ".safetensors":
+        return None
+    if _is_huggingface_checkpoint_url(checkpoint_path):
+        return _download_checkpoint_from_huggingface_url(
+            checkpoint_path,
+            checkpoint_min_free_gb=checkpoint_min_free_gb,
+        )
+    if checkpoint_path.startswith("s3://"):
+        cache_path = os.path.join(
+            local_cache_dir, checkpoint_path.removeprefix("s3://")
+        )
+        return cache_path if os.path.exists(cache_path) else None
+    return checkpoint_path
+
+
 def _load_checkpoint_from_s3(
     s3_path: str,
     ext: str,
@@ -837,8 +949,9 @@ def load_checkpoint(
     Args:
         checkpoint_path: ``s3://`` URI, local path, or HF URL. Single-file or
             DCP directory.
-        model: Model to load weights into. Required for DCP. Optional for
-            single-file: when provided, ``load_state_dict`` is called.
+        model: Model to load weights into. Required for DCP. Cached
+            safetensors are streamed into a provided model; other single-file
+            formats use ``load_state_dict``.
         checkpoint_type: ``"auto"``, ``"single"``, or ``"distributed"``.
         local_cache_dir: Directory for caches.
         credential_path: S3 credentials path.
@@ -873,6 +986,16 @@ def load_checkpoint(
                 checkpoint_type = "distributed"
 
     if checkpoint_type == "single":
+        if model is not None:
+            stream_path = _resolve_streamable_safetensors_path(
+                checkpoint_path,
+                local_cache_dir=local_cache_dir,
+                checkpoint_min_free_gb=checkpoint_min_free_gb,
+            )
+            if stream_path is not None:
+                _stream_safetensors_into_model(model, stream_path)
+                logger.info(f"Streamed checkpoint into model: {checkpoint_path}")
+                return model
         state_dict = load_single_checkpoint(
             checkpoint_path=checkpoint_path,
             local_cache_dir=local_cache_dir,

flashdreams/flashdreams/recipes/wan/transformer/wan21.py

@@ -157,6 +157,9 @@ class Wan21TransformerConfig(TransformerConfig):
     """Pre-load state-dict remap (e.g. Self-Forcing's
     ``generator_ema.model.…`` layout)."""
 
+    stream_checkpoint: bool = False
+    """Load cached safetensors directly into the model with bounded host residency."""
+
     batch_shape: tuple[int, ...] = (1,)
     """Batch dims of the latent (excluding the L, D dims)."""
 
@@ -279,13 +282,24 @@ def __init__(self, config: Wan21TransformerConfig) -> None:
         self.network.set_context_parallel_group(cp_group=self._cp_group)
 
         if config.checkpoint_path is not None:
-            state_dict = load_checkpoint(
-                config.checkpoint_path,
-                checkpoint_min_free_gb=config.checkpoint_min_free_gb,
-            )
-            if config.state_dict_transform is not None:
-                state_dict = config.state_dict_transform(state_dict)
-            self.network.load_state_dict(state_dict)
+            if config.stream_checkpoint:
+                if config.state_dict_transform is not None:
+                    raise ValueError(
+                        "stream_checkpoint does not support state_dict_transform"
+                    )
+                load_checkpoint(
+                    config.checkpoint_path,
+                    model=self.network,
+                    checkpoint_min_free_gb=config.checkpoint_min_free_gb,
+                )
+            else:
+                state_dict = load_checkpoint(
+                    config.checkpoint_path,
+                    checkpoint_min_free_gb=config.checkpoint_min_free_gb,
+                )
+                if config.state_dict_transform is not None:
+                    state_dict = config.state_dict_transform(state_dict)
+                self.network.load_state_dict(state_dict)
         self.network.update_parameters_after_loading_checkpoint()
 
         if config.compile_network:

flashdreams/tests/test_checkpoint_loading.py

@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Checkpoint loading behavior tests."""
+
+import importlib
+
+import pytest
+import torch
+from safetensors.torch import save_file as save_safetensors_file
+
+pytestmark = pytest.mark.ci_cpu
+
+
+def test_local_safetensors_uses_file_backed_loader(monkeypatch, tmp_path) -> None:
+    """Load local safetensors without materializing the file as bytes."""
+    checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load")
+    checkpoint_path = tmp_path / "weights.safetensors"
+    expected = {"weight": torch.ones(2)}
+    calls: list[tuple[str, str]] = []
+
+    def fake_load_file(path: str, *, device: str) -> dict[str, torch.Tensor]:
+        calls.append((path, device))
+        return expected
+
+    def reject_bytes_load(_data: bytes) -> dict[str, torch.Tensor]:
+        pytest.fail("safetensors checkpoints must use the file-backed loader")
+
+    monkeypatch.setattr(checkpoint_load, "load_safetensors_file", fake_load_file)
+    monkeypatch.setattr(checkpoint_load, "load_safetensors", reject_bytes_load)
+
+    actual = checkpoint_load.load_single_checkpoint(
+        str(checkpoint_path), map_location=torch.device("cpu")
+    )
+
+    assert actual is expected
+    assert calls == [(str(checkpoint_path), "cpu")]
+
+
+def test_safetensors_model_load_streams_without_full_state_dict(
+    monkeypatch, tmp_path
+) -> None:
+    """Stream safetensors tensors directly into a materialized model."""
+    checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load")
+    checkpoint_path = tmp_path / "weights.safetensors"
+    expected = torch.arange(6, dtype=torch.float32).view(2, 3)
+    save_safetensors_file({"weight": expected}, checkpoint_path)
+    model = torch.nn.Linear(3, 2, bias=False)
+
+    def reject_full_load(*_args, **_kwargs) -> None:
+        pytest.fail("model loads must not materialize the complete state dict")
+
+    monkeypatch.setattr(checkpoint_load, "load_safetensors_file", reject_full_load)
+
+    actual = checkpoint_load.load_checkpoint(str(checkpoint_path), model=model)
+
+    assert actual is model
+    torch.testing.assert_close(model.weight, expected)

integrations/lingbot/lingbot/config.py

@@ -76,6 +76,7 @@
                 in_dim=16 + 4 + 16,
             ),
             checkpoint_path=LINGBOT_WORLD_V1_CHECKPOINT_PATH,
+            stream_checkpoint=True,
             # Single-rollout layout: tensors flow through the stack as
             # ``[T, C, H, W]`` (or ``[T, ...]``) with no leading batch/view dim.
             batch_shape=(),
@@ -136,7 +137,8 @@
 )
 
 # LingBot-World v2 uses the same architecture and runtime as v1. The
-# transformer checkpoint is the only model-level substitution.
+# transformer checkpoint is the only model-level substitution; it inherits
+# the bounded checkpoint loader from the v1 base config.
 PIPELINE_LINGBOT_WORLD_V2_14B_CAUSAL_FAST = derive_config(
     PIPELINE_LINGBOT_WORLD_FAST,
     name="lingbot-world-v2-14b-causal-fast",

integrations/lingbot/tests/test_smoke.py

@@ -195,6 +195,14 @@ def test_lingbot_configs_carry_documented_checkpoint_disk_requirement() -> None:
         )
 
 
+def test_lingbot_configs_enable_streaming_checkpoint_load() -> None:
+    """Use bounded checkpoint loading for every LingBot model preset."""
+    for cfg in RUNNER_CONFIGS.values():
+        transformer = cfg.pipeline.diffusion_model.transformer
+        assert isinstance(transformer, LingbotWorldTransformerConfig)
+        assert transformer.stream_checkpoint
+
+
 def test_v2_only_replaces_the_v1_checkpoint() -> None:
     """Derive the v2 model by replacing only the v1 checkpoint and slug."""
     expected = derive_config(