Skip to content

fix: serialize torch module construction to prevent cross-thread meta-tensor leak - #50

Open
lstein wants to merge 2 commits into
mainfrom
fix/model-load-init-race
Open

fix: serialize torch module construction to prevent cross-thread meta-tensor leak#50
lstein wants to merge 2 commits into
mainfrom
fix/model-load-init-race

Conversation

@lstein

@lstein lstein commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

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_pretrained constructs models under accelerate's init_empty_weights, which monkeypatches nn.Module.register_parameter as 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.

main already serializes the loader paths under MODEL_LOAD_LOCK (an earlier revision of this PR added equivalent locking for the v7 branch — superseded). Three sites still construct torch modules outside it:

  • SafetyChecker (session thread, after every generation when the NSFW checker is on). Its first-run download now happens via snapshot_download before taking the lock, so a multi-GB transfer can't stall every other model load.
  • Quantized Qwen2.5-VL encoder load.
  • ESRGAN RRDBNet constructions.

All three take MODEL_LOAD_LOCK.write_lock() around construction only, per the lock-ordering contract documented on _ModelLoadReadWriteLock (this lock before any ModelCache lock; none of these sites touch the cache).

Testing

  • ruff + format clean; modules import cleanly under main's locked dependency set.
  • The leak mechanism and fix shape were validated empirically on the v7 revision of this PR (isolated two-thread repro of the accelerate patch leak, plus a threading regression test of the equivalent lock).

🤖 Generated with Claude Code

@joshistoast

Copy link
Copy Markdown
Collaborator

change merge destination to main

…_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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants