fix: serialize torch module construction to prevent cross-thread meta-tensor leak - #50
Open
lstein wants to merge 2 commits into
Open
fix: serialize torch module construction to prevent cross-thread meta-tensor leak#50lstein wants to merge 2 commits into
lstein wants to merge 2 commits into
Conversation
lstein
requested review from
JPPhoto,
Pfannkuchensack and
blessedcoolant
as code owners
August 4, 2026 00:24
Collaborator
|
change merge destination to |
…_LOAD_LOCK
Model construction is not thread-safe process-wide: transformers'
from_pretrained enters accelerate's init_empty_weights, which monkeypatches
nn.Module.register_parameter as a class attribute and restores it on exit.
Two overlapping constructions on different threads leak a patch permanently,
after which every later-constructed module gets meta-device parameters
("Cannot copy out of meta tensor; no data!") until the process restarts.
The model-load machinery already serializes its own loads under
MODEL_LOAD_LOCK, but three sites construct torch modules outside it:
- SafetyChecker (runs on the session thread after every generation when the
NSFW checker is enabled). Its first-run download now happens via
snapshot_download before the lock is taken, so a multi-GB transfer cannot
stall other loads.
- The quantized Qwen2.5-VL encoder load.
- ESRGAN's RRDBNet constructions.
All three now take MODEL_LOAD_LOCK.write_lock() around construction,
following the lock-ordering contract documented on _ModelLoadReadWriteLock
(this lock before any ModelCache lock; none of these sites touch the cache).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
lstein
force-pushed
the
fix/model-load-init-race
branch
from
August 4, 2026 02:35
d5d25c3 to
bd182df
Compare
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
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.
Summary
Serializes the three remaining direct torch model constructions under
MODEL_LOAD_LOCK, closing the cross-thread meta-tensor leak for code paths outside the model-load machinery.Background. transformers'
from_pretrainedconstructs models under accelerate'sinit_empty_weights, which monkeypatchesnn.Module.register_parameteras a class attribute and restores the saved original on exit. Two overlapping constructions on different threads each save the other's patched attribute as "the original", and one patch leaks permanently — after which every module constructed in the process gets meta-device parameters (NotImplementedError: Cannot copy out of meta tensor; no data!) until restart. Reproduced empirically; hit in the field when image-index worker/request threads began loading models concurrently with generation.mainalready serializes the loader paths underMODEL_LOAD_LOCK(an earlier revision of this PR added equivalent locking for the v7 branch — superseded). Three sites still construct torch modules outside it:snapshot_downloadbefore taking the lock, so a multi-GB transfer can't stall every other model load.RRDBNetconstructions.All three take
MODEL_LOAD_LOCK.write_lock()around construction only, per the lock-ordering contract documented on_ModelLoadReadWriteLock(this lock before anyModelCachelock; none of these sites touch the cache).Testing
🤖 Generated with Claude Code