fix: load checkpoints in weights-only mode

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
         )
 
 

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,