Stream safetensors checkpoints into models

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

@@ -157,6 +157,16 @@ 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."""
+
+    init_device: str | None = None
+    """Optional device used for initial network parameter allocation.
+
+    Large streaming-checkpoint models can set this to the final runtime device
+    so the module is not first materialized as fp32 CPU tensors.
+    """
+
     batch_shape: tuple[int, ...] = (1,)
     """Batch dims of the latent (excluding the L, D dims)."""
 
@@ -273,19 +283,29 @@ def __init__(self, config: Wan21TransformerConfig) -> None:
         self._output_height: int | None = None
         self._output_width: int | None = None
 
-        self.network = config.network.setup()
-        self.network = self.network.to(dtype=config.dtype)
+        self.network = self._setup_network(config)
         self.network.eval()
         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:
@@ -308,6 +328,23 @@ def __init__(self, config: Wan21TransformerConfig) -> None:
             self._cuda_graph_dispatch.uncond_call or self.network
         )
 
+    @staticmethod
+    def _setup_network(config: Wan21TransformerConfig) -> WanDiTNetwork:
+        init_device = (
+            None if config.init_device is None else torch.device(config.init_device)
+        )
+        if init_device is None:
+            return config.network.setup().to(dtype=config.dtype)
+
+        previous_dtype = torch.get_default_dtype()
+        try:
+            torch.set_default_dtype(config.dtype)
+            with torch.device(init_device):
+                network = config.network.setup()
+        finally:
+            torch.set_default_dtype(previous_dtype)
+        return network.to(device=init_device, dtype=config.dtype)
+
     @property
     def latent_shape(self) -> tuple[int, ...]:
         """Per-rank post-patchify latent shape ``[*batch_shape, L/cp, D]``.

flashdreams/tests/test_checkpoint_loading.py

@@ -0,0 +1,113 @@
+# 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 json
+
+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)
+
+
+def test_sharded_safetensors_model_load_streams_without_merged_state_dict(
+    monkeypatch, tmp_path
+) -> None:
+    """Stream indexed safetensors shards into a model without merging first."""
+    checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load")
+    shard_a = tmp_path / "model-00001-of-00002.safetensors"
+    shard_b = tmp_path / "model-00002-of-00002.safetensors"
+    index_path = tmp_path / "model.safetensors.index.json"
+    expected_weight = torch.arange(6, dtype=torch.float32).view(2, 3)
+    expected_bias = torch.tensor([3.0, 4.0], dtype=torch.float32)
+    save_safetensors_file({"weight": expected_weight}, shard_a)
+    save_safetensors_file({"bias": expected_bias}, shard_b)
+    index_path.write_text(
+        json.dumps(
+            {
+                "metadata": {"total_size": 0},
+                "weight_map": {
+                    "weight": shard_a.name,
+                    "bias": shard_b.name,
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+    model = torch.nn.Linear(3, 2)
+
+    def reject_merge(*_args, **_kwargs) -> None:
+        pytest.fail("sharded model loads must not materialize a merged state dict")
+
+    monkeypatch.setattr(
+        checkpoint_load,
+        "_load_sharded_safetensors_index_checkpoint",
+        reject_merge,
+    )
+
+    actual = checkpoint_load.load_checkpoint(str(index_path), model=model)
+
+    assert actual is model
+    torch.testing.assert_close(model.weight, expected_weight)
+    torch.testing.assert_close(model.bias, expected_bias)

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/lingbot/webrtc/server.py

@@ -89,12 +89,14 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int, default=8080)
     parser.add_argument(
+        "--config-name",
         "--config_name",
         type=str,
-        default="lingbot-world-fast",
+        default=LingbotRuntimeConfig().config_name,
         help="LingBot-World config preset from PIPELINE_CONFIGS.",
     )
     parser.add_argument(
+        "--no-compile",
         "--no_compile",
         action="store_true",
         help="Disable torch.compile when building the Lingbot pipeline.",
@@ -106,12 +108,14 @@ def parse_args() -> argparse.Namespace:
         help="Torch device used for the Lingbot runtime.",
     )
     parser.add_argument(
+        "--warmup-chunks",
         "--warmup_chunks",
         type=int,
         default=10,
         help="Number of synthetic startup chunks to generate for kernel autotuning.",
     )
     parser.add_argument(
+        "--warmup-timeout-s",
         "--warmup_timeout_s",
         type=float,
         default=600.0,

integrations/lingbot/lingbot/webrtc/session.py

@@ -800,7 +800,10 @@ def _initialize_sync(self) -> None:
             enable_sync_and_profile=True,
             diffusion_model=dict(
                 seed=rollout_seed,
-                transformer=dict(compile_network=self.config.compile_network),
+                transformer=dict(
+                    compile_network=self.config.compile_network,
+                    init_device=str(self._device),
+                ),
             ),
         )
         self._pipeline = pipeline_config.setup().to(device=self._device)

integrations/lingbot/tests/test_distributed_server_main.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import sys
 from argparse import Namespace
 
 import pytest
@@ -53,6 +54,27 @@ def _args(device: str = "cuda:0") -> Namespace:
     )
 
 
+def test_parse_args_defaults_to_webrtc_preset_and_accepts_kebab_options(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "lingbot-webrtc",
+            "--warmup-chunks",
+            "0",
+            "--no-compile",
+        ],
+    )
+
+    args = server.parse_args()
+
+    assert args.config_name == "lingbot-world-fast-taehv-window15-sink3"
+    assert args.warmup_chunks == 0
+    assert args.no_compile is True
+
+
 def test_initialize_distributed_single_process_honors_default_device(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:

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(