JAX hits swap OOM during parallel training on dual A100—even with 400 GB free RAM

[tudorcebere]

Hi JAX community,

First, thank you for the fantastic library!

System specs

2 × NVIDIA A100 (80 GB each)
500 GB + system RAM
Only 2 GB swap (cannot be increased on this machine)

What I’m doing

• Training a large custom architecture. One process with only one GPU used works flawlessly.
• For hyper-parameter tuning I try to launch two independent JAX processes, one per GPU (tested with Python multiprocessing/spawn, Ray workers, and plain Bash-spawned processes).

The issue

During jax.jit/XLA compile each process briefly allocates into swap. When both models compile at the same time, the 2 GB swap fills and the OS kills the job for OOM—even though ~400 GB of physical RAM remains unused. I verified this with htop that RAM stays mostly free while swap grows.

Is there a way to tell JAX/XLA never to touch swap (or to keep host allocations in RAM) during compilation? Alternatively, is there a flag or best practice for running parallel JAX compilations that avoids swap exhaustion?

Any pointers or heuristics would be much appreciated!

Thanks,
Tudor