fix(security): harden CWE-502 torch.load sites (weights_only + nosemgrep) - #385
Conversation
…semgrep Follow-up to the annotation-only PR #384. Mitigates the CWE-502 (unsafe deserialization) scan findings from helicalAI/dashboard#1154 for the torch.load sites, and suppresses the Trail of Bits scanner on the sites that are now safe or unavoidably require full unpickling. Groups A + B (torch.load) addressed here; group C (raw pickle.load of gene dictionaries) is intentionally deferred to a follow-up. - weights_only=True (real mitigation) added to the loads whose payloads are tensors / plain state dicts: scgpt best_model.pt, uce all_tokens.torch, uce gene-embedding dicts, and the uce model-weights load (the last was missed by the Bastion scan but hardened for consistency). Under the pinned torch 2.7.0 this matches the existing default, so it is a no-op at runtime and only hardens intent. - HyenaDNA keeps weights_only=False: its Lightning-style .ckpt carries non-tensor objects the safe loader rejects. Documented + suppressed with a "trusted sources" justification (a try-safe-first attempt was rejected in review: it fired a misleading CWE-502 warning and double-read the file on every legitimate load). - base_models legacy fallback keeps weights_only=False (pre-v2.0.0 full-model pickles); documented + suppressed, reached only when the safe load fails. - Every torch.load carries a `# nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch` on the line immediately preceding the call (Semgrep only honors adjacent suppressions) plus a per-site CWE-502 justification. Validated: scGPT and HyenaDNA load real downloaded checkpoints end-to-end; py_compile passes on all touched files. CI exercises these load paths against real (non-mocked) upstream artifacts. Refs helicalAI/dashboard#1154 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CYrvyUo92dDqbcfZSUw7Bp
| all_pe = torch.load(token_file) | ||
| # weights_only=True restricts torch.load to tensors/plain types (all_tokens.torch is a | ||
| # tensor), so a tampered file cannot execute arbitrary code during unpickling | ||
| # (CWE-502, helicalAI/dashboard#1154). | ||
| # nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch | ||
| all_pe = torch.load(token_file, weights_only=True) |
There was a problem hiding this comment.
Do we need the #nosemgrep here if we are only loading with weights_only=True?
There was a problem hiding this comment.
With weights_only=True, this is essentially a false positive — that argument is the mitigation, not the vulnerability. The Semgrep rule (trailofbits.python.pickles-in-pytorch) flags every torch.load call regardless of arguments, because it can't always statically prove the flag's value. Here you've already applied the recommended fix.
Why torch.load is flagged in general
torch.load historically defaults to Python's pickle module for deserialization. Pickle can execute arbitrary code during unpickling (via reduce / GLOBAL opcodes), so loading an untrusted checkpoint = arbitrary code execution on your machine. That's CWE-502, and it's a real, exploited attack vector (malicious models on Hugging Face, etc.).
Why weights_only=True defuses it
When you pass weights_only=True, PyTorch does not use the general pickle machinery. Instead it uses a restricted unpickler (_load with a weights_only_unpickler) that:
- Only allows a hardcoded allowlist of safe types (tensors, Storage, primitive types, basic containers like dict/list, and a small set of explicitly registered classes).
- Refuses arbitrary GLOBAL opcodes — so it cannot import and call arbitrary functions/classes, which is exactly the mechanism pickle RCE relies on.
So a crafted checkpoint that tries to smuggle in a reduce-based payload will raise an UnpicklingError rather than execute. This is why weights_only=True became the default in PyTorch 2.6.
So, summary:
- Semgrep rule always flags every torch.load call regardless of the arguments
- Since, there's no risk because we are using weights_only=True, we mark the code as false positive by adding the # nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch and documenting why it is a false positive in the code itself
| # weights_only=True restricts torch.load to tensors/plain types (the UCE checkpoint is | ||
| # a plain state dict), so a tampered checkpoint cannot execute arbitrary code during | ||
| # unpickling (CWE-502, helicalAI/dashboard#1154). Not in the Bastion CSV, but hardened | ||
| # here for consistency with the token-file load above. | ||
| # nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch | ||
| torch.load(model_path, map_location=model_config["device"], weights_only=True), |
There was a problem hiding this comment.
Do we need the #nosemgrep here if we are only loading with weights_only=True?
| # Safe: weights_only=True restricts torch.load to tensors/plain types, so a | ||
| # tampered checkpoint cannot execute arbitrary code (CWE-502, helicalAI/dashboard#1154). | ||
| # nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch | ||
| state_dict = torch.load(path, weights_only=True) |
There was a problem hiding this comment.
Do we need the #nosemgrep here if we are only loading with weights_only=True?
2 participants
What & why
Follow-up to the annotation-only PR #384. This PR actually mitigates the CWE-502 (unsafe deserialization) scan findings from helicalAI/dashboard#1154 for the
torch.loadsites, and suppresses the Trail of Bits scanner on the sites that are now safe (or unavoidably require full unpickling).Scope: Groups A + B (
torch.load). Group C (rawpickle.loadof gene dictionaries — geneformer/uce/transcriptformer/tahoe) is intentionally deferred to a follow-up PR and is left untouched here.Changes
scgpt_utils.py,uce_utils.py(×2),uce/gene_embeddings.py(×2)weights_only=True+# nosemgreptranscriptformer/model.py,tahoe/.../blocks.py,base_models.py:319# nosemgrep+ justificationweights_only=True; just clears the scannerNotes:
weights_onlyalready defaults toTrue, so the added kwargs are no-ops at runtime and only harden intent.uce_utils.py:167(UCE model-weights load) was missed by the Bastion scan but hardened here for consistency with the token-file load in the same file.# nosemgrep: trailofbits.python.pickles-in-pytorch.pickles-in-pytorchis placed on the line immediately preceding itstorch.load(Semgrep only honors adjacent suppressions).Verification
test_scgpt_utils.py/test_hyena_dna_model.pypass).python -m py_compilepasses on all 7 touched files.