JAX multi-threading race condition

[ammar1510]

I get a race condition when I try to run the following snippet on 2 disjoint shards.
Would help if someone can provide a mental model on thinking multi-threading setup in JAX & what kinds of pitfalls to avoid.

import threading
import time
import sys
import jax
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from functools import partial
from jax import lax, shard_map
import numpy as np

import jax.profiler

jax.profiler.start_server(9999)

devices = jax.devices()
num_devices = len(devices)
if jax.process_index() == 0:
    print(f"Total devices: {num_devices}")

# Mesh A: 1D (all devices in one group)
devices_1 = np.array(devices[: num_devices // 2])
mesh1d = Mesh(devices_1, ("i",))

# Mesh B: 2D (grid). We'll assume at least 4 devices.
# We reshape the devices into (2, num_devices // 2)

devices_2 = np.array(devices[num_devices // 2 :])
mesh2d = Mesh(devices_2.reshape(-1, 2), ("rows", "cols"))

def mesh_sharding(
    pspec: P,
    mesh: Mesh,
) -> NamedSharding:
    return NamedSharding(mesh, pspec)


@partial(shard_map, mesh=mesh1d, in_specs=(P("i"), P("i")), out_specs=P("i"))
def op_1(x, y):
    return x + y


@partial(
    shard_map,
    mesh=mesh2d,
    in_specs=(P("rows", None), P(None, "cols")),
    out_specs=P("rows", "cols"),
)
def op_2(x, y):
    return x @ y


def run_1(x, y):
    while True:
        if jax.process_index() == 0:
            print("starting op-1")
        result = op_1(x, y).block_until_ready()
        if jax.process_index() == 0:
            print("=========OP-1=========")
            for i, shard in enumerate(result.addressable_shards):
                print(f"  Shard {i} (Device {shard.device.id}):\n{shard.data.shape}")
            sys.stdout.flush()
        time.sleep(5)


def run_2(x, y):
    while True:
        if jax.process_index() == 0:
            print("starting op-2")
        result = op_2(x, y).block_until_ready()
        if jax.process_index() == 0:
            print("=========OP-2=========")
            for i, shard in enumerate(result.addressable_shards):
                print(f"  Shard {i} (Device {shard.device.id}):\n{shard.data.shape}")
            sys.stdout.flush()
        time.sleep(8)


def main():

    x_1 = jax.device_put(
        jax.random.normal(jax.random.key(0), (8192, 8192)),
        mesh_sharding(P("i"), mesh1d),
    )
    y_1 = jax.device_put(
        jax.random.normal(jax.random.key(42), (8192, 8192)),
        mesh_sharding(P("i"), mesh1d),
    )

    x_2 = jax.device_put(
        jax.random.normal(jax.random.key(0), (8192, 8192)),
        mesh_sharding(P("rows", None), mesh2d),
    )
    y_2 = jax.device_put(
        jax.random.normal(jax.random.key(42), (8192, 8192)),
        mesh_sharding(P(None, "cols"), mesh2d),
    )

    thread_1 = threading.Thread(target=run_1, args=(x_1, y_1), daemon=False)
    thread_2 = threading.Thread(target=run_2, args=(x_2, y_2), daemon=False)
    thread_1.start()
    thread_2.start()
    thread_1.join()
    thread_2.join()


if __name__ == "__main__":
    main()

https://symbolize.stripped_domain/r/?trace=189f590,7f5757a4251f&map=
*** SIGSEGV, si_code=128, see go/general-protection-fault received by PID 1826938 (TID 1828145) on cpu 222; stack trace: ***
PC: @          0x189f590  (unknown)  take_gil.llvm.17674134442370007079
    @     0x7f56a1dae682       1888  FailureSignalHandler()
    @     0x7f5757a42520  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=189f590,7f56a1dae681,7f5757a4251f&map=
E1220 21:42:23.568534 1828145 coredump_hook.cc:302] RAW: Remote crash data gathering hook invoked.
E1220 21:42:23.568549 1828145 coredump_hook.cc:341] RAW: Skipping coredump since rlimit was 0 at process start.
E1220 21:42:23.568557 1828145 client.cc:270] RAW: Coroner client retries enabled, will retry for up to 30 sec.
E1220 21:42:23.568559 1828145 coredump_hook.cc:397] RAW: Sending fingerprint to remote end.
E1220 21:42:23.568574 1828145 coredump_hook.cc:406] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] stat failed on crash reporting socket /var/google/services/logmanagerd/remote_coredump.socket (Is the listener running?): No such file or directory
E1220 21:42:23.568578 1828145 coredump_hook.cc:458] RAW: Dumping core locally.
E1220 21:42:24.205592 1828145 process_state.cc:778] RAW: Raising signal 11 with default behavior
./learn.sh: line 16: 1826938 Segmentation fault      (core dumped) python learn_mesh.py
##### Command execution on worker 0 failed with exit status 139. Continuing.

[YigitElma]

@ammar1510 I couldn't reproduce the error with 2 GPUs (the only difference I had was the second mesh, and I removed the profiler) or 8 CPUs. Do you get the error consistently or from time to time?

I am also not super familiar with the multithreading with JAX, but I know that certain operations in JAX have multithreaded C++ implementations which releases Pythons GIL. I assume that when your CPU dispatches the call to run_1 and run_2 or the compilation of them, the GIL is released for a short time and then locked again (then your GPU does the main computation). If both threads try to release it at the same time, you might get the error. Probably on my system, those calls never coincided, hence I didn't get the error.

If my understanding is correct, threading.Lock() might help to prevent those rare instances. I would also appreciate an explanation from someone who has a deeper understanding of the topic.

[ammar1510]

The above error is from a multi-processing run.

• I get the error consistenly after leaving it running for about 5-10 mins on a multi-node, multi-process setup.
• can you share some resources where I can read how JAX interacts and releases python GIL?