From b0c0d2f5b9480426bbc3b5a4abfef02085f1ea81 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Tue, 4 Aug 2026 16:33:24 +0000 Subject: [PATCH] Stream safetensors checkpoints into models Signed-off-by: Gangzheng Tong --- .../flashdreams/core/checkpoint/load.py | 131 +++++++++++++++++- .../recipes/wan/transformer/wan21.py | 28 +++- flashdreams/tests/test_checkpoint_loading.py | 70 ++++++++++ integrations/lingbot/lingbot/config.py | 4 +- integrations/lingbot/tests/test_smoke.py | 8 ++ 5 files changed, 229 insertions(+), 12 deletions(-) create mode 100644 flashdreams/tests/test_checkpoint_loading.py diff --git a/flashdreams/flashdreams/core/checkpoint/load.py b/flashdreams/flashdreams/core/checkpoint/load.py index 4d11f764..dfb580d7 100644 --- a/flashdreams/flashdreams/core/checkpoint/load.py +++ b/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, diff --git a/flashdreams/flashdreams/recipes/wan/transformer/wan21.py b/flashdreams/flashdreams/recipes/wan/transformer/wan21.py index c1ae0ab0..4606dcf9 100644 --- a/flashdreams/flashdreams/recipes/wan/transformer/wan21.py +++ b/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: diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py new file mode 100644 index 00000000..b48ffc49 --- /dev/null +++ b/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) diff --git a/integrations/lingbot/lingbot/config.py b/integrations/lingbot/lingbot/config.py index 72c647f5..ef8530d4 100644 --- a/integrations/lingbot/lingbot/config.py +++ b/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", diff --git a/integrations/lingbot/tests/test_smoke.py b/integrations/lingbot/tests/test_smoke.py index bf9c1ab1..b6bf596a 100644 --- a/integrations/lingbot/tests/test_smoke.py +++ b/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(