From aae515daec1c3b022c061ce37280a0dd8c12e9eb Mon Sep 17 00:00:00 2001 From: Jonathan McCaffrey Date: Fri, 4 Sep 2026 13:56:55 -0700 Subject: [PATCH] fix: load checkpoints in weights-only mode Use explicit weights_only=True for .pt, .pth, and .ckpt reads from local paths, Hugging Face cache entries, S3, and the distributed-checkpoint cache. This keeps loading behavior consistent across supported PyTorch versions and matches the state-dict return contract. Keep safetensors handling and tensor state-dict compatibility unchanged. Add coverage for non-weight objects and normal tensor-only checkpoints across the supported legacy extensions. Signed-off-by: Jonathan McCaffrey --- .../flashdreams/core/checkpoint/load.py | 8 +- flashdreams/tests/test_checkpoint_loading.py | 94 +++++++++++++++++++ 2 files changed, 99 insertions(+), 3 deletions(-) diff --git a/flashdreams/flashdreams/core/checkpoint/load.py b/flashdreams/flashdreams/core/checkpoint/load.py index 4432e1be5..951177617 100644 --- a/flashdreams/flashdreams/core/checkpoint/load.py +++ b/flashdreams/flashdreams/core/checkpoint/load.py @@ -584,7 +584,9 @@ def load_distributed_checkpoint( if local_cache_checkpoint_path is not None and os.path.exists( local_cache_checkpoint_path ): - state_dict = torch.load(local_cache_checkpoint_path, map_location="cpu") + state_dict = torch.load( + local_cache_checkpoint_path, map_location="cpu", weights_only=True + ) model.load_state_dict(state_dict) logger.info( f"Loaded successfully from the local cache: {local_cache_checkpoint_path}" @@ -741,7 +743,7 @@ def _load_checkpoint_from_local( if ext == ".safetensors": return load_safetensors_file(path, device=_safetensors_device(map_location)) else: - return torch.load(path, map_location=map_location, weights_only=False) + return torch.load(path, map_location=map_location, weights_only=True) def _copy_checkpoint_tensor(destination: torch.Tensor, source: torch.Tensor) -> int: @@ -1083,7 +1085,7 @@ def _load_checkpoint_from_s3( return load_safetensors(data_bytes) else: return torch.load( - io.BytesIO(data_bytes), map_location=map_location, weights_only=False + io.BytesIO(data_bytes), map_location=map_location, weights_only=True ) diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py index f3b129bd1..2d19d7768 100644 --- a/flashdreams/tests/test_checkpoint_loading.py +++ b/flashdreams/tests/test_checkpoint_loading.py @@ -6,7 +6,9 @@ from __future__ import annotations import importlib +import io import json +import pickle from pathlib import Path from typing import Any @@ -17,6 +19,98 @@ pytestmark = pytest.mark.ci_cpu +def _record_object_load(marker_path: str) -> dict[str, torch.Tensor]: + Path(marker_path).write_text("loaded", encoding="utf-8") + return {"weight": torch.ones(1, 1)} + + +class _NonWeightCheckpointObject: + def __init__(self, marker_path: str) -> None: + self.marker_path = marker_path + + def __reduce__(self) -> tuple[Any, tuple[str]]: + return _record_object_load, (self.marker_path,) + + +@pytest.mark.parametrize("source", ["huggingface", "s3", "distributed-cache"]) +def test_checkpoint_loads_reject_non_weight_objects( + source: str, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Use weights-only mode for every core checkpoint source.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + marker_path = tmp_path / "loaded" + serialized = io.BytesIO() + torch.save(_NonWeightCheckpointObject(str(marker_path)), serialized) + checkpoint_bytes = serialized.getvalue() + + if source == "huggingface": + checkpoint_path = tmp_path / "weights.pt" + checkpoint_path.write_bytes(checkpoint_bytes) + monkeypatch.setattr( + checkpoint_load, + "_download_checkpoint_from_huggingface_url", + lambda *_args, **_kwargs: str(checkpoint_path), + ) + + def load() -> object: + return checkpoint_load.load_single_checkpoint( + "https://huggingface.co/org/model/resolve/main/weights.pt" + ) + + elif source == "s3": + + class FakeS3FileSystem: + def __init__(self, credential_path: str) -> None: + assert credential_path == "credentials" + + def create_stream(self, path: str, mode: str) -> io.BytesIO: + assert (path, mode) == ("s3://bucket/weights.pt", "rb") + return io.BytesIO(checkpoint_bytes) + + monkeypatch.setattr(checkpoint_load, "S3FileSystem", FakeS3FileSystem) + + def load() -> object: + return checkpoint_load._load_checkpoint_from_s3( + "s3://bucket/weights.pt", ".pt", "credentials" + ) + + else: + checkpoint_path = tmp_path / "bucket" / "model.pt" + checkpoint_path.parent.mkdir() + checkpoint_path.write_bytes(checkpoint_bytes) + model = torch.nn.Linear(1, 1, bias=False) + + def load() -> object: + return checkpoint_load.load_distributed_checkpoint( + model, + "s3://bucket/model", + local_cache_dir=str(tmp_path), + ) + + with pytest.raises(pickle.UnpicklingError): + load() + + assert not marker_path.exists() + + +@pytest.mark.parametrize("extension", [".pt", ".pth", ".ckpt"]) +def test_pickle_checkpoint_formats_still_load_tensor_state_dicts( + extension: str, + tmp_path: Path, +) -> None: + """Keep tensor-only legacy checkpoint formats working in safe mode.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + checkpoint_path = tmp_path / f"weights{extension}" + expected = {"weight": torch.arange(3)} + torch.save(expected, checkpoint_path) + + actual = checkpoint_load.load_single_checkpoint(str(checkpoint_path)) + + torch.testing.assert_close(actual["weight"], expected["weight"]) + + def test_local_safetensors_uses_file_backed_loader( monkeypatch: pytest.MonkeyPatch, tmp_path: Path,