[Fix] model_loader: Avoid unsafe pickle fallback when loading PyTorch checkpoints

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -302,23 +302,8 @@ def safetensors_weights_iterator(
 
 
 def _load_pt_file(bin_file: str, device: str) -> dict:
-    """Load a PyTorch checkpoint file, handling legacy tar format.
-
-    PyTorch 2.6 changed the default of weights_only from False to True.
-    Legacy tar format files cannot be loaded with weights_only=True.
-    This function tries weights_only=True first, then falls back to False
-    for legacy tar format files from trusted sources (HuggingFace Hub).
-    """
-    try:
-        return torch.load(bin_file, map_location=device, weights_only=True)
-    except RuntimeError as e:
-        if "legacy .tar format" in str(e):
-            logger.warning(
-                "Loading %s with weights_only=False (legacy tar format)",
-                os.path.basename(bin_file),
-            )
-            return torch.load(bin_file, map_location=device, weights_only=False)
-        raise
+    """Load a PyTorch checkpoint file with weights_only=True (no pickle execution)."""
+    return torch.load(bin_file, map_location=device, weights_only=True)
 
 
 def pt_weights_iterator(

python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py

@@ -303,7 +303,7 @@ def convert(param):
             else:
                 return {k: v for k, v in param.items() if "module." not in k}
 
-        state = torch.load(flownet_path, map_location="cpu", weights_only=False)
+        state = torch.load(flownet_path, map_location="cpu", weights_only=True)
         self.flownet.load_state_dict(convert(state), strict=False)
         logger.info("Loaded RIFE weights from %s", flownet_path)
 

python/sglang/srt/model_loader/weight_utils.py

@@ -1396,23 +1396,8 @@ def _load_file(st_file: str):
 
 
 def _load_pt_file(bin_file: str) -> dict:
-    """Load a PyTorch checkpoint file, handling legacy tar format.
-
-    PyTorch 2.6 changed the default of weights_only from False to True.
-    Legacy tar format files cannot be loaded with weights_only=True.
-    This function tries weights_only=True first, then falls back to False
-    for legacy tar format files from trusted sources (HuggingFace Hub).
-    """
-    try:
-        return torch.load(bin_file, map_location="cpu", weights_only=True)
-    except RuntimeError as e:
-        if "legacy .tar format" in str(e):
-            logger.warning(
-                "Loading %s with weights_only=False (legacy tar format)",
-                os.path.basename(bin_file),
-            )
-            return torch.load(bin_file, map_location="cpu", weights_only=False)
-        raise
+    """Load a PyTorch checkpoint file with weights_only=True (no pickle execution)."""
+    return torch.load(bin_file, map_location="cpu", weights_only=True)
 
 
 def pt_weights_iterator(

python/sglang/srt/models/deepseek_ocr.py

@@ -832,7 +832,7 @@ def _build_sam(
     )
     image_encoder.eval()
     if checkpoint is not None:
-        state_dict = torch.load(checkpoint)
+        state_dict = torch.load(checkpoint, map_location="cpu", weights_only=True)
         image_encoder.load_state_dict(
             {k[30:]: v for k, v in state_dict.items() if "vision_tower_high" in k},
             strict=True,
@@ -1416,7 +1416,7 @@ def build_qwen2_decoder_as_encoder(
         max_query=max_query,
     )
     if checkpoint is not None:
-        state_dict = torch.load(checkpoint)
+        state_dict = torch.load(checkpoint, map_location="cpu", weights_only=True)
         decoder_as_encoder.load_state_dict(state_dict, strict=True)
     return decoder_as_encoder
 

python/sglang/srt/models/roberta.py

@@ -276,7 +276,7 @@ def _load_sparse_linear(model_path_or_dir: str, sparse_head: str) -> dict:
             local_dir = download_from_hf(model_path_or_dir, allow_patterns=sparse_head)
             path = os.path.join(local_dir, sparse_head)
 
-        state_dict = torch.load(path)
+        state_dict = torch.load(path, map_location="cpu", weights_only=True)
         return state_dict
 
 

test/registered/unit/model_loader/test_weight_utils.py

@@ -1,7 +1,10 @@
 """Unit tests for srt/model_loader/weight_utils.py shard-index consistency."""
 
+import io
 import json
 import os
+import pickle
+import tarfile
 import tempfile
 import unittest
 from types import SimpleNamespace
@@ -10,6 +13,7 @@
 import torch
 
 from sglang.srt.model_loader.weight_utils import (
+    _load_pt_file,
     filter_duplicate_safetensors_files,
     initialize_dummy_weights,
     maybe_add_mtp_safetensors,
@@ -202,6 +206,50 @@ def test_local_unindexed_mtp(self):
                 listing.assert_not_called()
 
 
+class _OpenOnUnpickle:
+    def __init__(self, marker):
+        self.marker = marker
+
+    def __reduce__(self):
+        return (open, (self.marker, "w"))
+
+
+def _write_legacy_tar_checkpoint(path, payload):
+    members = {
+        "storages": pickle.dumps(0) + pickle.dumps([]),
+        "tensors": pickle.dumps(0),
+        "pickle": pickle.dumps(payload),
+    }
+    with tarfile.open(path, "w") as tar:
+        for name, data in members.items():
+            info = tarfile.TarInfo(name)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+
+class TestLoadPtFile(CustomTestCase):
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.folder = self._tmp.name
+
+    def tearDown(self):
+        self._tmp.cleanup()
+
+    def test_tensor_state_dict_loads(self):
+        path = os.path.join(self.folder, "pytorch_model.bin")
+        torch.save({"w": torch.arange(4)}, path)
+        loaded = _load_pt_file(path)
+        self.assertTrue(torch.equal(loaded["w"], torch.arange(4)))
+
+    def test_legacy_tar_pickle_payload_is_not_executed(self):
+        marker = os.path.join(self.folder, "pwned")
+        path = os.path.join(self.folder, "pytorch_model.bin")
+        _write_legacy_tar_checkpoint(path, _OpenOnUnpickle(marker))
+        with self.assertRaises(RuntimeError):
+            _load_pt_file(path)
+        self.assertFalse(os.path.exists(marker))
+
+
 class TestInitializeDummyWeights(CustomTestCase):
     def test_nvfp4_dummy_scales_are_neutral(self):
         model = torch.nn.Linear(8, 4, bias=False)