PyTorch #182948 — Why torch.cuda._lazy_call() Deadlocked When Its Callback Re-entered It

PyTorch Lazy Initialization, Part 1 of 3 — A non-reentrant initialization lock, callback execution inside the critical section, and a thread waiting for itself

A GPU initialization failure does not always produce an exception.

Sometimes the process simply stops.

No CUDA error appears.

No Python traceback is printed.

The program does not crash, but it never makes progress again.

PyTorch PR #182948 fixed one such failure in torch.cuda._lazy_call().

The issue appeared when a callback passed to _lazy_call() invoked another operation that re-entered _lazy_call().

Examples included:

torch.cuda.set_rng_state()

torch.cuda.manual_seed()

The outer call already held PyTorch’s CUDA initialization lock.

The nested call attempted to acquire the same lock again.

That lock was an ordinary:

threading.Lock()

It was not reentrant.

The same thread could not acquire it twice.

The execution therefore became:

Acquire the initialization lock
        ↓
Invoke a callback while still holding the lock
        ↓
The callback re-enters _lazy_call()
        ↓
Try to acquire the same lock again
        ↓
Wait forever

No second thread was required.

The thread deadlocked against itself.

The problem was originally encountered on ROCm while larger models were running through torch.compile, even though the implementation lives under PyTorch’s shared torch.cuda namespace.


Why PyTorch has lazy GPU initialization

Importing PyTorch does not necessarily initialize the GPU runtime immediately.

That would make even simple operations such as:

import torch

perform device discovery, initialize a driver context, create runtime state, and establish other GPU resources before the program had actually requested GPU work.

Instead, PyTorch uses lazy initialization.

A simplified lifecycle is:

Import torch
        ↓
CUDA or ROCm runtime remains uninitialized
        ↓
GPU-related requests are recorded or deferred
        ↓
The first operation requiring the backend triggers initialization
        ↓
Queued callbacks are executed
        ↓
The backend becomes fully initialized

This design has several practical benefits.

A Python program can import PyTorch on a machine without using a GPU.

Startup avoids unnecessary runtime work.

Code can queue operations such as random-number-generator setup before the device context exists.

The mechanism requires a boundary between two states:

Not initialized

and:

Initialized

_lazy_call() helps manage work that can arrive on either side of that boundary.


What _lazy_call() was responsible for

Before the fix, the central logic was effectively:

def _lazy_call(callback, **kwargs):
    with _initialization_lock:
        if is_initialized():
            callback()
        else:
            queue(callback)

The intended behavior was straightforward.

Before initialization

Backend not initialized
        ↓
Store the callback
        ↓
Run it later during _lazy_init()

After initialization

Backend already initialized
        ↓
Run the callback immediately

The lock serialized the decision.

Without synchronization, one thread could be queueing a callback while another thread completed initialization and drained the queue.

The lock itself was not the error.

The error was that the callback was executed while the lock was still held.


The lock protected more than the state decision

A critical section should normally contain only the state that requires mutual exclusion.

In this case, the important decision was:

Should this callback be queued?

or

Should it run now?

But the old implementation extended the critical section beyond that decision.

Acquire lock
        ↓
Check initialized state
        ↓
Run the entire callback
        ↓
Release lock

The lock remained held for whatever the callback decided to do.

That callback could:

  • set an RNG state,

  • seed generators,

  • call another PyTorch API,

  • allocate a tensor,

  • trigger another lazy operation,

  • raise an exception,

  • block on another resource,

  • or re-enter _lazy_call() itself.

The initialization lock therefore protected not only PyTorch’s internal state transition.

It also enclosed arbitrary callback behavior.

That made the critical section depend on code outside the lock’s own implementation.


A non-reentrant lock cannot be acquired twice by the same thread

Python’s ordinary threading.Lock does not track recursive ownership in a way that permits the owner to acquire it again.

Consider:

lock.acquire()

lock.acquire()

The second acquisition waits for someone to release the lock.

But the only thread that can reach the corresponding release is the same thread now blocked inside the second acquisition.

Thread owns lock
        ↓
Thread waits to acquire lock
        ↓
Thread cannot reach release
        ↓
No progress is possible

This differs from threading.RLock, which allows the owning thread to acquire the lock recursively and maintains an acquisition count.

PyTorch was using a normal Lock.

Once a callback re-entered _lazy_call(), the second acquisition could never complete.


The already-initialized path could deadlock immediately

One of the two affected paths occurred after GPU initialization had already completed.

Suppose the backend was initialized.

A caller invoked:

torch.cuda._lazy_call(callback)

The old function entered the initialization lock:

_initialization_lock acquired

It then checked:

is_initialized() == True

and executed the callback without releasing the lock.

lock held
        ↓
callback begins

Now suppose that callback called:

torch.cuda.set_rng_state(state)

That operation could eventually use _lazy_call() again.

The nested execution became:

Outer _lazy_call()
        ↓
Lock acquired
        ↓
callback()
        ↓
set_rng_state()
        ↓
Inner _lazy_call()
        ↓
Attempts same lock
        ↓
Deadlock

The outer invocation could not finish until the callback returned.

The callback could not return until the nested _lazy_call() acquired the lock.

The lock could not be released until the outer invocation finished.

The dependency formed a closed cycle.


The callback did not have to call _lazy_call() explicitly

Most users never call the internal _lazy_call() function directly.

That does not prevent reentrancy.

The callback only needs to invoke another API whose implementation eventually uses _lazy_call().

For example:

User callback
        ↓
torch.cuda.set_rng_state()
        ↓
internal lazy callback registration
        ↓
_lazy_call()

or:

Queued initialization callback
        ↓
torch.cuda.manual_seed()
        ↓
another lazy callback
        ↓
_lazy_call()

The recursive edge may be several functions away from the original callback.

This is why callback-under-lock patterns are dangerous.

The lock owner cannot determine reentrancy merely by inspecting the callback’s first line.


The GIL did not prevent this deadlock

Python’s Global Interpreter Lock does not make every Python lock reentrant.

The GIL and _initialization_lock serve different purposes.

GIL
→ coordinates Python bytecode execution
_initialization_lock
→ protects PyTorch’s GPU initialization state

A thread can hold the GIL and still block while trying to acquire a non-reentrant threading.Lock that it already owns.

The problem was not simultaneous execution of Python bytecode by two threads.

It was a circular wait involving one explicit lock.

Same Python thread

holds _initialization_lock

and

waits for _initialization_lock

The hang provided little information about its origin

A Python exception leaves a traceback.

A segmentation fault may produce a core dump.

A self-deadlock often produces neither.

The process remains alive.

The thread waits inside lock acquisition.

GPU utilization may fall to zero.

The user may observe only:

Model execution stopped.

No error appeared.

The process did not exit.

Because the issue was originally seen while running larger models under torch.compile, the apparent cause could have been mistaken for:

  • a generated GPU kernel hang,

  • a compiler deadlock,

  • a distributed synchronization problem,

  • a stream dependency,

  • an allocator stall,

  • or a device-runtime failure.

The actual cycle existed in Python-side lazy initialization.


Why the problem could remain hidden in ordinary use

Many callbacks do not re-enter lazy initialization.

For example:

Acquire lock
        ↓
Run simple callback
        ↓
Callback returns
        ↓
Release lock

works.

The defect requires a callback with a specific property:

callback
→ invokes another lazy GPU operation

A simple callback that mutates ordinary Python state may never expose it.

A callback that performs device-related RNG setup can.

The old implementation was therefore not guaranteed to deadlock.

It was structurally capable of deadlocking when the callback graph contained a recursive lazy call.


Reentrancy changes the shape of a call graph

Without reentrancy, the call graph looked like this:

_lazy_call()
        ↓
callback()
        ↓
return

With reentrancy:

_lazy_call()
        ↓
callback()
        ↓
GPU API
        ↓
_lazy_call()
        ↓
callback()

The second _lazy_call() is not a separate independent operation.

It occurs while the first invocation still owns the initialization lock.

The same function therefore appears twice on the active stack under one lock acquisition.

This is the boundary the old implementation had not preserved.


The first repair: run the obvious fast path before taking the lock

The new implementation begins with:

if is_initialized() or hasattr(_tls, "is_initializing"):
    callback()
    return

The already-initialized path no longer needs to acquire the initialization lock merely to run a callback.

Backend initialized
        ↓
Run callback immediately
        ↓
No initialization lock held

If the callback re-enters _lazy_call():

Outer callback
        ↓
Nested _lazy_call()
        ↓
Backend still initialized
        ↓
Nested callback runs

There is no recursive lock acquisition.

The self-deadlock disappears.

The callback can still recurse incorrectly for other reasons, but it no longer waits for PyTorch’s initialization lock merely because of that recursion.


Why a pre-lock fast path was not enough by itself

Consider two threads.

Thread A

Checks is_initialized()
→ False

Before Thread A acquires the lock, Thread B initializes the backend.

Thread B

Acquires initialization lock
        ↓
Initializes backend
        ↓
Releases lock

Thread A now acquires the lock.

Inside the critical section it sees:

is_initialized()
→ True

The old code would invoke the callback immediately while still holding the lock.

The first fast-path check did not protect this race because initialization changed between:

the check before the lock

and:

the check inside the lock

The fixed implementation therefore also changes the inside-lock branch.


The lock now decides, but does not execute

The new structure is effectively:

if is_initialized() or hasattr(_tls, "is_initializing"):
    callback()
    return

run_now = False

with _initialization_lock:
    if is_initialized():
        run_now = True
    else:
        queue(callback)

if run_now:
    callback()

Inside the lock, PyTorch decides only:

Queue later

or

Run after releasing the lock

The callback itself is outside the critical section.

Acquire lock
        ↓
Inspect initialization state
        ↓
Record decision
        ↓
Release lock
        ↓
Invoke callback if needed

The lock protects the transition.

It no longer owns the callback’s complete execution.


run_now closes the initialization race

The run_now variable handles this timeline:

Thread A:
initial fast-path check says “not initialized”

Thread B:
finishes initialization

Thread A:
acquires the lock and sees “initialized”

Thread A cannot queue the callback.

The queue may already have been drained.

But it also must not execute the callback while holding the lock.

The solution is:

Inside lock:
run_now = True

Outside lock:
callback()

This preserves both requirements.

The callback is not lost.

The callback is not run under the lock.

The fix narrowed the lock’s ownership boundary

Before:

Initialization lock owns:

state inspection

queue mutation

callback execution

After:

Initialization lock owns:

state inspection

queue mutation

run-now decision

Callback execution belongs to the caller again.

This is a more precise ownership model.

Internal synchronization
→ protected internally

External callback behavior
→ executed after internal lock release

Why changing the lock to RLock was not the chosen repair

One possible response to recursive acquisition would be:

_initialization_lock =
    threading.RLock()

That could allow the same thread to acquire the lock recursively.

PR #182948 did not take that approach.

It retained the ordinary lock and moved callback execution outside the critical section.

This distinction matters.

An RLock could make this one recursive acquisition legal while still allowing arbitrary callback work to occur under the initialization lock.

Callback runs under lock

Callback blocks elsewhere

Initialization state remains locked

It could also preserve long critical sections and more complicated lock-order interactions.

The landed design removes the structural cause instead:

The initialization lock should not be held while an external callback is running.

The reasoning about broader lock design is an architectural interpretation; the source directly confirms only that the patch retained the ordinary lock and moved callback execution outside it.


There was a second deadlock path during initialization

The already-initialized path was not the only failure.

_lazy_init() also drained queued callbacks while holding _initialization_lock.

At that point:

torch._C._cuda_init()
→ already completed

_initialized
→ still False

_tls.is_initializing
→ True

If a queued callback called _lazy_call() again, the old nested function saw:

is_initialized()
→ False

and attempted to acquire the lock already held by _lazy_init().

That produced another self-deadlock.

The new fast path includes:

hasattr(_tls, "is_initializing")

so a callback running on the initialization thread can execute immediately without reacquiring the lock.

Part 2 examines these two paths separately and explains why is_initialized() alone could not represent the transitional state.


A Boolean initialized flag did not describe every safe state

The backend lifecycle contains more than two relevant conditions.

A simplified state machine is:

1. Not initialized

2. Native backend initialization has completed,
   but queued Python callbacks are still being drained

3. Fully initialized

During state 2:

_initialized == False

because the complete lazy-initialization procedure has not finished.

But the native backend is already ready enough for reentrant initialization-thread callbacks to execute.

That transitional condition is represented by:

_tls.is_initializing

The fix therefore uses:

initialized
OR
current thread is the initializer

as the immediate-execution condition.


Thread-local state prevents other threads from bypassing the lock

The marker is stored in:

threading.local()

It applies only to the thread currently draining the initialization queue.

Initializer thread
→ sees is_initializing
→ may run reentrant callback immediately
Other thread
→ does not see that marker
→ still waits for initialization coordination

A global Boolean such as:

initialization_in_progress = True

could allow unrelated threads to incorrectly bypass the lock while initialization was incomplete.

The thread-local marker distinguishes:

the thread that already owns the initialization process

from:

other threads that still need synchronization

The first regression test reproduces the initialized-path deadlock

The patch adds a subprocess test shaped like this:

import torch

torch.cuda.init()

state = torch.cuda.get_rng_state()

torch.cuda._lazy_call(
    lambda: torch.cuda.set_rng_state(state)
)

print("done")

The sequence is:

Initialize backend
        ↓
Enter _lazy_call()
        ↓
Callback invokes set_rng_state()
        ↓
set_rng_state() re-enters _lazy_call()

Before the fix:

Outer _lazy_call holds lock
        ↓
Inner _lazy_call waits for same lock
        ↓
Subprocess times out

After the fix:

Callback runs without lock
        ↓
Nested _lazy_call completes
        ↓
“done” is printed

The test runs in another process because the old code can permanently block the interpreter executing that scenario.


The second regression test reproduces queued reentrancy

The second script is conceptually:

import torch

torch.cuda._lazy_call(
    lambda: torch.cuda.manual_seed(42)
)

torch.cuda.init()

print("done")

The first _lazy_call() occurs before initialization, so the callback is queued.

Then:

torch.cuda.init()
        ↓
_lazy_init() acquires initialization lock
        ↓
native backend initialization completes
        ↓
queued callback begins
        ↓
manual_seed() re-enters _lazy_call()

Before the fix, the nested call waited for the lock held by _lazy_init().

After the fix, the initializer thread sees _tls.is_initializing and executes the nested callback immediately.

Part 2 examines why this route required a different condition from the already-initialized fast path.


The regression tests use a 60-second process timeout

Both scenarios are run through a common helper that starts a fresh Python subprocess.

The parent waits for a maximum of:

60 seconds

The successful route prints:

done

and exits quickly.

The timeout exists only to bound the failure case.

If the subprocess does not return, the parent reports that _lazy_call() likely deadlocked.

The commit states that both tests timed out before the fix and passed after it on a ROCm gfx90a host.

Part 3 examines why a separate process was a stronger boundary than a thread for testing a lock cycle that can freeze the interpreter.


The fix changes lock duration as well as reentrancy

Removing callback execution from the critical section also shortens the time other threads can be blocked by _initialization_lock.

Before:

Lock duration
=
state decision
+
entire callback duration

After:

Lock duration
=
state decision
+
queue update

A callback that performs slow work no longer keeps unrelated lazy calls waiting on the initialization lock for the full callback duration.

The PR was aimed at correctness rather than a published performance improvement, so no general latency claim should be made.

But the ownership duration is visibly smaller in the new structure.


Callback exceptions also occur after lock release

Under the old initialized path, a callback exception propagated while the with block was unwinding.

Python would still release the lock as the context manager exited.

The new path goes further:

Lock already released
        ↓
Callback begins
        ↓
Callback may succeed or raise

Internal initialization-state synchronization is no longer active while external callback behavior is being evaluated.

This reduces the amount of lock state that must be reasoned about during callback failure.


The patch fixed the torch.cuda implementation only

The landed commit modifies:

torch/cuda/__init__.py

test/test_cuda.py

It does not modify:

torch/xpu/__init__.py

torch/mtia/__init__.py

The PR explicitly noted that XPU and MTIA had the same callback-under-lock structure and would require follow-up treatment.

The current main snapshots checked for this article still show:

with _initialization_lock:
    if is_initialized():
        callback()

in both XPU and MTIA.

Their lazy-initialization loops also drain queued callbacks while holding the same initialization lock.

This confirms a remaining structural pattern.

It does not by itself prove that a real XPU or MTIA workload currently triggers the same deadlock.

That distinction will be addressed in Part 3.


A duplicated pattern is not automatically a reproduced bug

The source supports:

XPU and MTIA contain similar callback-under-lock code.

It does not yet establish:

A specific XPU callback re-enters _lazy_call() and deadlocks.

A specific MTIA callback does the same.

To turn structural similarity into a concrete backend finding, each backend needs a callback that:

enters its _lazy_call()

runs under the initialization lock

re-enters the same lazy-call path

The CUDA/ROCm patch provided two minimal reproductions for its path.

Equivalent backend-specific reproductions have not been established by the source reviewed here.


What the patch directly changed

PyTorch commit cec50d89 directly changed:

  • callback execution after CUDA/ROCm initialization,

  • callback execution when initialization finishes during a lock race,

  • reentrant callbacks on the initialization thread,

  • the critical section around queue-or-run decisions,

  • regression coverage for both deadlock routes.

The central code transition was:

Before

Acquire initialization lock
→ decide
→ execute callback
→ release lock
After

Check immediate safe path
→ acquire lock if necessary
→ decide
→ release lock
→ execute callback

What the patch does not prove

The public source does not establish that:

  • every historical torch.compile hang on ROCm came from this bug,

  • all callbacks are now safe from every form of recursion,

  • callbacks cannot deadlock on unrelated locks,

  • XPU or MTIA has reproduced the same failure,

  • replacing callback-under-lock removes every lazy-initialization race,

  • the change provides a measurable model-performance improvement,

  • the 60-second timeout represents the actual deadlock-detection latency required on every CI machine.

The supported conclusion is narrower:

Two concrete torch.cuda._lazy_call() reentrancy paths could acquire the same non-reentrant initialization lock twice on one thread, and the patch removed callback execution from that lock boundary.


The deadlock was created by a correct lock in the wrong place

The initialization lock had a legitimate purpose.

It prevented multiple threads from racing through backend initialization and callback-queue state.

The callback also had a legitimate purpose.

It applied GPU state once the runtime was ready.

The failure emerged from their composition:

Correct lock

+

Valid callback

+

Callback executed while lock remains held

+

Callback re-enters lock owner

=

Self-deadlock

No isolated component needed to be obviously broken.

The contract between them was.


The safest critical section is the smallest one that preserves the state transition

The final design can be summarized as:

Lock protects:

whether to queue

whether initialization completed during the race

which callback list owns the request

It does not protect:

what the callback does after the decision

This is a reusable concurrency principle.

Do not call code capable of re-entering the synchronization boundary while holding a non-reentrant lock that protects that boundary.

The callback does not need to be formally user-provided for this principle to apply.

Internal framework callbacks can re-enter through other internal APIs just as easily.


Part 2: the two reentrancy paths were not identical

The next article separates:

Path 1

Backend already initialized
→ callback was executed under lock

from:

Path 2

Native initialization complete
but Python queue still draining
→ _initialized still False
→ initializer thread re-enters under the same lock

It will examine why the final implementation needed both:

is_initialized()

and:

hasattr(_tls, "is_initializing")

as well as the run_now decision made inside the lock.


Part 3: a deadlock test must survive the deadlock

The final article examines the regression strategy.

Potentially permanent lock wait
        ↓
fresh subprocess
        ↓
finite timeout
        ↓
parent process remains in control

It also examines the remaining backend boundary:

CUDA/ROCm
→ fixed and tested

XPU/MTIA
→ similar source pattern still visible

runtime reproduction
→ not yet established here

Related material


Patch status: Landed in PyTorch main as commit cec50d89
Affected function: torch.cuda._lazy_call()
Lock type: Non-reentrant threading.Lock
Deadlock trigger: A callback re-entered _lazy_call() while the outer call or _lazy_init() held the initialization lock
Primary repair: Decide under the lock, execute callbacks after releasing it
Transitional-state repair: Run immediately when the current thread has _tls.is_initializing
Regression coverage: Two subprocess tests with a 60-second timeout
Validation recorded by the patch: Both old failure paths timed out before the fix and completed after it on a ROCm gfx90a host
Remaining source boundary: Similar callback-under-lock structures remain visible in current XPU and MTIA initialization code

This is Part 1 of a three-part series on PyTorch’s lazy-initialization reentrancy deadlock.

Part 2 examines the already-initialized path and the queued-during-initialization path separately, and explains why run_now and the thread-local initialization marker solve different races.

Part 3 examines why the regression tests run in subprocesses and how the similar XPU and MTIA source patterns should be distinguished from a reproduced backend failure.

#PyTorch #CUDA #ROCm #Python #Concurrency #Deadlock #LazyInitialization #ThreadingLock #TorchCompile #CodeAnalysis


Popular posts from this blog

AMD Is Trying to Turn Kernel Diversity Into One Software Asset

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?

NVIDIA CUTLASS #3017 — Why load() and store() Mapped the Same Tile Coordinate to Different Addresses