[Fix] model_loader: Avoid unsafe pickle fallback when loading PyTorch checkpoints by nguyendvk · Pull Request #42760 · sgl-project/sglang

nguyendvk · 2026-10-06T12:23:27Z

Motivation

torch.load(weights_only=False) deserializes with pickle, so a crafted checkpoint can execute code when it is loaded. _load_pt_file fell back to it for legacy tar checkpoints, and several loaders called torch.load on model files without weights_only=True. This is the same class of issue as #35653 (RIFE) and is consistent with the contribution guide's rule against unpickling untrusted data.

Modifications

Remove the weights_only=False fallback in _load_pt_file (srt/model_loader/weight_utils.py and multimodal_gen/runtime/loader/weight_utils.py). Legacy tar checkpoints now fail with torch's own error and should be converted to safetensors or a zip-format .pt offline.
Pass weights_only=True to torch.load in deepseek_ocr.py (x2), roberta.py and rife_interpolator.py.
Add unit tests for _load_pt_file: a tensor state dict still loads, and a pickle payload is not executed.
srt/debug_utils is unchanged: it reads developer-written tensor dumps, not model weights.

Accuracy Tests

Not applicable: only the checkpoint deserialization path changes; no model forward code is touched.

Speed Tests and Profiling

Not applicable.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ✅ Run #39535409763
Latest PR Test (Extra): ❌ Run #39535409518
Latest PR Test (AMD ROCm 10): ❌ Run #39535409559

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 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

thanhhao98 · 2026-10-13T02:06:43Z

/rerun-test test/registered/unit/model_loader/test_weight_utils.py python/sglang/multimodal_gen/test/unit/test_weight_utils.py test/registered/vlm/test_vision_openai_server_a.py::TestDeepseekOCRServer test/registered/unit/multimodal/test_deepseek_ocr_geometry.py python/sglang/multimodal_gen/test/server/test_server_1_gpu.py::TestDiffusionServerOneGpu::test_diffusion_generation[wan2_1_t2v_1.3b_frame_interp_2x]

github-actions · 2026-10-13T02:07:07Z

🚀 ubuntu-latest (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/model_loader/test_weight_utils.py
cd test/ && python3 registered/unit/multimodal/test_deepseek_ocr_geometry.py

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

python3 -m pytest python/sglang/multimodal_gen/test/unit/test_weight_utils.py -x
python3 -m pytest python/sglang/multimodal_gen/test/server/test_server_1_gpu.py::TestDiffusionServerOneGpu::test_diffusion_generation[wan2_1_t2v_1.3b_frame_interp_2x] -x

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/vlm/test_vision_openai_server_a.py TestDeepseekOCRServer

thanhhao98 · 2026-10-13T02:42:55Z

@nvpohanh pls help to review, thanks!

nvpohanh

[by Claude Code] Reviewed head 600dd5e. One inline comment on test coverage.

nvpohanh · 2026-10-13T03:21:22Z

+        loaded = _load_pt_file(path)
+        self.assertTrue(torch.equal(loaded["w"], torch.arange(4)))
+
+    def test_pickle_payload_is_not_executed(self):


[by Claude Code] Severity: style | Confidence: High

Issue

test_pickle_payload_is_not_executed does not cover the removed fallback. torch.save writes a zip file, so torch.load(weights_only=True) raises UnpicklingError on main too, before the legacy .tar format branch runs. I ran TestLoadPtFile against main's _load_pt_file, and both tests pass.

Example

A legacy tar file with members storages (pickle.dumps(0) + pickle.dumps([])), tensors (pickle.dumps(0)), and pickle (the _OpenOnUnpickle payload) creates the marker on main and raises RuntimeError on this PR (checked with torch 2.14).

Fix

Build the payload as that legacy tar file with tarfile instead of torch.save, so the test fails if the fallback comes back.

Good catch, thanks. Replaced the torch.save payload with a hand-built legacy tar checkpoint (storages/tensors/pickle members). Checked with torch 2.14: on main's _load_pt_file the payload runs and the marker is created; on this branch it raises RuntimeError and nothing runs. The tensor state-dict test is unchanged. Also thanks for the history on the fallback.

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 <noreply@anthropic.com>

nvpohanh · 2026-10-13T10:35:51Z

/tag-and-rerun-ci

nvpohanh · 2026-10-13T10:37:21Z

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.

nvpohanh · 2026-10-13T12:23:06Z

[by Claude Code] The relevant CI failure (https://github.com/sgl-project/sglang/actions/runs/38384431344/job/115259669168) is not caused by this PR. It is a known breakage on main, tracked in #44054. Proposed fix: #44055. CI for this PR is blocked until that is fixed.

Picks up sgl-project#44055, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/38384431344/job/115259669168 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-13T12:46:07Z

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.

nvpohanh · 2026-10-13T17:54:54Z

/rerun-failed-ci

…hts-only-upstream # Conflicts: # test/registered/unit/model_loader/test_weight_utils.py

nvpohanh · 2026-10-19T03:53:18Z

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, fix the PR title if the title check fails, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.

nvpohanh · 2026-10-19T03:53:29Z

/rerun-failed-ci

nvpohanh · 2026-10-19T03:54:37Z

[by Claude Code] @nguyendvk, the PR Title check failed, so I renamed this PR from Avoid unsafe pickle fallback when loading PyTorch checkpoints to [Fix] model_loader: Avoid unsafe pickle fallback when loading PyTorch checkpoints to match the required [Type] scope: title format. Feel free to change the type or scope if I picked the wrong one.

Picks up sgl-project#45038, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/39122132737/job/117963435935 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-19T16:31:48Z

/rerun-failed-ci

nvpohanh · 2026-10-19T17:31:52Z

/rerun-failed-ci

Starts a fresh PR Test Base run: the GitHub jobs API keeps returning 502 for run 39239074876 (https://github.com/sgl-project/sglang/actions/runs/39239074876/job/118205011160), so rerunning that run cannot pass its wait-for-base-b gate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-20T04:01:54Z

/rerun-failed-ci

nvpohanh · 2026-10-20T05:32:59Z

/rerun-failed-ci

nvpohanh · 2026-10-20T12:54:09Z

[by Claude Code] The relevant CI failure (base-c-test-8-gpu-b300 (0)) is not caused by this PR. It is a known CI breakage from main that affects every PR from a fork, tracked in #45392: the gated Muse Spark test needs a repository secret that GitHub does not pass to fork PRs. Proposed fix: #45393. CI for this PR is blocked until that is fixed.

Picks up sgl-project#45428, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/39348980625/job/118556303062 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-21T02:22:40Z

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, fix the PR title if the title check fails, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.

nvpohanh · 2026-10-21T02:41:03Z

[by Claude Code] The relevant CI failure (base-a-test-cpu (11)) is not caused by this PR. It is a known breakage on main that #45476 will fix. CI for this PR is blocked until that PR merges.

Picks up sgl-project#45476, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/39532544294/job/119001829089 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh · 2026-10-21T07:05:55Z

/rerun-failed-ci

nvpohanh · 2026-10-21T20:36:31Z

[by Claude Code] All NV pipelines have passed.

nvpohanh · 2026-10-22T03:17:44Z

@mmangkad Could you review this? This is SGLang security issue. Thanks!

nvpohanh · 2026-10-23T12:39:21Z

@b8zhong @mickqian could you review this security change? Thanks!

nvdang and others added 2 commits October 6, 2026 13:11

Add regression test for weights_only checkpoint loading

600dd5e

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

nguyendvk requested review from AgainstEntropy, BBuf, HaiShaw, b8zhong, kevin-mii, mickqian, mmangkad, niehen6174, ping1jing2 and yichiche as code owners October 6, 2026 12:23

github-actions Bot added deepseek diffusion SGLang Diffusion labels Oct 6, 2026

nvpohanh reviewed Oct 13, 2026

View reviewed changes

Comment thread python/sglang/srt/model_loader/weight_utils.py Resolved

nvdang and others added 2 commits October 13, 2026 15:40

Merge branch 'main' into fix/torch-load-weights-only-upstream

3bbfc8f

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 13, 2026

This was referenced Oct 13, 2026

[CI] test_pi05_modelopt_fp8_cuda.py fails on main: No module named 'imageio' #44054

Closed

[CI] Run test_pi05_modelopt_fp8_cuda in the kernel-unit lane #44055

Merged

Merge branch 'main' into fix/torch-load-weights-only-upstream

d115352

Picks up sgl-project#44055, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/38384431344/job/115259669168 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh mentioned this pull request Oct 15, 2026

[CI] Runner 5090-vast-no-0 is out of disk: every job fails in Install dependencies with 'No space left on device' #44479

Closed

Merge remote-tracking branch 'upstream/main' into fix/torch-load-weig…

6963ece

…hts-only-upstream # Conflicts: # test/registered/unit/model_loader/test_weight_utils.py

This comment has been minimized.

Sign in to view

Merge branch 'main' into fix/torch-load-weights-only-upstream

9dd8934

This comment has been minimized.

Sign in to view

nvpohanh changed the title ~~Avoid unsafe pickle fallback when loading PyTorch checkpoints~~ [Fix] model_loader: Avoid unsafe pickle fallback when loading PyTorch checkpoints Oct 19, 2026

Merge branch 'main' into fix/torch-load-weights-only-upstream

b462c1c

Picks up sgl-project#45038, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/39122132737/job/117963435935 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh mentioned this pull request Oct 20, 2026

[CI] test_muse_spark_12_b300_pd fails on fork PRs: GatedRepoError 403 (HF_MUSE_SPARK_TOKEN unset) #45392

Closed

Merge branch 'main' into fix/torch-load-weights-only-upstream

7454093

Picks up sgl-project#45428, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/39348980625/job/118556303062 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge branch 'main' into fix/torch-load-weights-only-upstream

5a378cf

Picks up sgl-project#45476, which fixes the unrelated CI failure https://github.com/sgl-project/sglang/actions/runs/39532544294/job/119001829089 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

nvpohanh approved these changes Oct 22, 2026

View reviewed changes

mickqian approved these changes Oct 27, 2026

View reviewed changes

nvpohanh merged commit f82f8a9 into sgl-project:main Oct 27, 2026
274 of 317 checks passed

leejnau mentioned this pull request Nov 3, 2026

[CI] diffusion: Wan VAE decode time spikes on the 1-gpu-5090 sm* runners #48196

Open

Conversation

nguyendvk commented Oct 6, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

thanhhao98 commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

thanhhao98 commented Oct 13, 2026

Uh oh!

nvpohanh left a comment

Choose a reason for hiding this comment

Uh oh!

nvpohanh Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

nguyendvk Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nvpohanh commented Oct 13, 2026

Uh oh!

nvpohanh commented Oct 13, 2026

Uh oh!

nvpohanh commented Oct 13, 2026

Uh oh!

nvpohanh commented Oct 13, 2026

Uh oh!

nvpohanh commented Oct 13, 2026

Uh oh!

This comment has been minimized.

This comment has been minimized.

nvpohanh commented Oct 19, 2026

Uh oh!

nvpohanh commented Oct 19, 2026

Uh oh!

nvpohanh commented Oct 19, 2026

Uh oh!

nvpohanh commented Oct 19, 2026

Uh oh!

nvpohanh commented Oct 19, 2026

Uh oh!

nvpohanh commented Oct 20, 2026

Uh oh!

nvpohanh commented Oct 20, 2026

Uh oh!

nvpohanh commented Oct 20, 2026

Uh oh!

nvpohanh commented Oct 21, 2026

Uh oh!

nvpohanh commented Oct 21, 2026

Uh oh!

nvpohanh commented Oct 21, 2026

Uh oh!

nvpohanh commented Oct 21, 2026

Uh oh!

nvpohanh commented Oct 22, 2026

Uh oh!

nvpohanh commented Oct 23, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

nguyendvk commented Oct 6, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading