fix(security): set DataLoader pin_memory explicitly (automatic memory pinning) - #378
Merged
Merged
Conversation
* Merge pull request #371 from LiudengZhang/fix/save-load-fine-tuning-head Fix save_model to include fine-tuning head weights * Bump version from 2.0.1 to 2.0.2 * Update accelerate package version to 1.13.0 (#375) * Update pyproject.toml --------- Co-authored-by: LiudengZhang <99156394+LiudengZhang@users.noreply.github.com>
…d pinning
Trail of Bits' Semgrep rule `automatic-memory-pinning` (CWE-676) flags
`torch.utils.data.DataLoader` calls that leave `pin_memory` unset, since
the pinning behaviour is then implicit. Explicitly gate pinning on GPU
availability across all 27 DataLoader call sites:
pin_memory=torch.cuda.is_available()
This pins host memory only when a CUDA device is present (faster, async
host->GPU copies) and avoids wasting scarce page-locked memory on
CPU-only machines. Previously-hardcoded `pin_memory=True` sites
(scgpt, tahoe, transcriptformer) are unified to the same predicate so
they no longer pin unconditionally on CPU-only runs.
Refs: trailofbits.python.automatic-memory-pinning
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01R2VtZqWHS8Dv4V88SbjftT
bputzeys
approved these changes
Jul 8, 2026
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01E81Hd78n4wNdihG3hT9Wqp
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What & why
Trail of Bits' Semgrep rule
automatic-memory-pinning(CWE-676) flags everytorch.utils.data.DataLoaderthat leavespin_memoryunset — the pinning behaviour is then implicit/undefined.This PR sets the choice explicitly on all 27 DataLoader call sites (15 files):
Scope
pin_memorykwarg.pin_memory=True(scgpt ×4, tahoe ×1, transcriptformer ×1): unified to the same predicate so they no longer pin unconditionally on CPU-only runs.Validation
pin_memory=torch.cuda.is_available(); 0 leftoverpin_memory=True.py_compile) and importtorch.pin_memory.Review note
A high-effort code review flagged that
torch.cuda.is_available()keys on global CUDA presence rather than the model's configured device — so on a GPU host where a user explicitly selectsdevice="cpu", pinned memory is still allocated. This is a minor, edge-case performance concern (no correctness bug). We deliberately kepttorch.cuda.is_available()as it is the ToB-recommended idiom and satisfies the scan; a per-model configured-device check was deemed more invasive/risky for the benefit.Refs:
trailofbits.python.automatic-memory-pinning🤖 Generated with Claude Code