From 1040711119264d0c0aa18b22d213d9e1443bc688 Mon Sep 17 00:00:00 2001 From: nvdang Date: Tue, 6 Oct 2026 11:12:26 +0700 Subject: [PATCH 1/3] Avoid unsafe pickle fallback when loading PyTorch checkpoints Remove the weights_only=False fallback for legacy tar checkpoints and pass weights_only=True to the remaining torch.load calls on model files, so a malicious checkpoint cannot execute code via pickle. Legacy tar checkpoints should be converted to safetensors offline. Co-Authored-By: Claude Sonnet 5.5 --- .../runtime/loader/weight_utils.py | 19 ++----------------- .../runtime/postprocess/rife_interpolator.py | 2 +- .../sglang/srt/model_loader/weight_utils.py | 19 ++----------------- python/sglang/srt/models/deepseek_ocr.py | 4 ++-- python/sglang/srt/models/roberta.py | 2 +- 5 files changed, 8 insertions(+), 38 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index 5cfa32016911..f0e87efd053a 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -335,23 +335,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( diff --git a/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py b/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py index 04c9831c0087..b82b527cd8ae 100644 --- a/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py +++ b/python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py @@ -300,7 +300,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) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 82951f9f199e..98f988dccf19 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1296,23 +1296,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( diff --git a/python/sglang/srt/models/deepseek_ocr.py b/python/sglang/srt/models/deepseek_ocr.py index 8e242e715633..b72cb4f6573d 100644 --- a/python/sglang/srt/models/deepseek_ocr.py +++ b/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 diff --git a/python/sglang/srt/models/roberta.py b/python/sglang/srt/models/roberta.py index 0fbf1c2b479d..8efcaafbb40d 100644 --- a/python/sglang/srt/models/roberta.py +++ b/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 From 600dd5ef3296825a2260bad934ff0e05d9121f72 Mon Sep 17 00:00:00 2001 From: nvdang Date: Tue, 6 Oct 2026 19:23:19 +0700 Subject: [PATCH 2/3] Add regression test for weights_only checkpoint loading Co-Authored-By: Claude Sonnet 5.5 --- .../unit/model_loader/test_weight_utils.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/test/registered/unit/model_loader/test_weight_utils.py b/test/registered/unit/model_loader/test_weight_utils.py index 9ebaa72037f1..a7f1bd5d1d6f 100644 --- a/test/registered/unit/model_loader/test_weight_utils.py +++ b/test/registered/unit/model_loader/test_weight_utils.py @@ -7,7 +7,10 @@ from types import SimpleNamespace from unittest.mock import patch +import torch + from sglang.srt.model_loader.weight_utils import ( + _load_pt_file, filter_duplicate_safetensors_files, maybe_add_mtp_safetensors, ) @@ -199,5 +202,36 @@ 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")) + + +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_pickle_payload_is_not_executed(self): + marker = os.path.join(self.folder, "pwned") + path = os.path.join(self.folder, "pytorch_model.bin") + torch.save({"w": _OpenOnUnpickle(marker)}, path) + with self.assertRaises(Exception): + _load_pt_file(path) + self.assertFalse(os.path.exists(marker)) + + if __name__ == "__main__": unittest.main() From b15d28d9dea28f8ed6e10e34c818d79d7bf474f6 Mon Sep 17 00:00:00 2001 From: nvdang Date: Tue, 13 Oct 2026 15:40:21 +0700 Subject: [PATCH 3/3] Test the legacy tar fallback with a real tar checkpoint torch.save writes a zip file, so the previous payload test also passed against the old fallback. Build a legacy tar checkpoint instead so the test fails if the weights_only=False fallback is reintroduced. Co-Authored-By: Claude Sonnet 5.5 --- .../unit/model_loader/test_weight_utils.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/test/registered/unit/model_loader/test_weight_utils.py b/test/registered/unit/model_loader/test_weight_utils.py index a7f1bd5d1d6f..bc0c6eeb2716 100644 --- a/test/registered/unit/model_loader/test_weight_utils.py +++ b/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 @@ -210,6 +213,19 @@ 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() @@ -224,11 +240,11 @@ def test_tensor_state_dict_loads(self): loaded = _load_pt_file(path) self.assertTrue(torch.equal(loaded["w"], torch.arange(4))) - def test_pickle_payload_is_not_executed(self): + 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") - torch.save({"w": _OpenOnUnpickle(marker)}, path) - with self.assertRaises(Exception): + _write_legacy_tar_checkpoint(path, _OpenOnUnpickle(marker)) + with self.assertRaises(RuntimeError): _load_pt_file(path) self.assertFalse(os.path.exists(marker))