examples: give each MPI rank its own JAX device for the subspace expansion

[Sophia Wen (hfwen0502)]

Summary

The subspace expansion in run_sqd_enlarge_subspace_sbd.py is JAX, and JAX has no
MPI awareness: every rank picks the first device it can see — the same device for
all of them. Since XLA reserves 75% of a device on first use, one rank claims it and
the rest die with RESOURCE_EXHAUSTED: CUDA_ERROR_OUT_OF_MEMORY. The documented
workaround was JAX_PLATFORMS=cpu, which sidesteps the problem rather than fixing it.

This gives each rank its own device by asking the backend which one it will use, so JAX
follows the same gpu_id = rank % num_gpus convention the README already documents for
SBD instead of inventing a second policy.

Why local_device_ids is passed explicitly

The driver does use cluster_detection_method="mpi4py" — that part works, since the
rendezvous fields come from comm.Get_rank()/Get_size(). What it does not let JAX
infer is the local device, because JAX's mpi4py path derives the local rank from
COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED), and on MPICH 5.0.0 (ch4:ucx) that puts
every rank in its own node — local_size=1 everywhere, and MPI.Win.Allocate_shared
over COMM_WORLD is refused outright ("processes are not in the same shared memory
domain"), so the locality table itself is wrong, not just the split. Left to JAX, every
rank got local_device_ids=[0] and ran on physical GPU 0.

That went unnoticed because Device.id is a global id in multi-process JAX: eight
processes with one device each report cuda:0…cuda:7 whichever card they are on.
device.local_hardware_id is the physical index, and it read 0 on every rank.

The change

sbd.get_device_id() — new public query, backed by a planned_device_id binding.
bindings.cpp already computes rank % device_count inside each diag entry point and
discards it; this returns it.

The driver passes it as local_device_ids, and additionally:

• sets XLA_PYTHON_CLIENT_PREALLOCATE=false so XLA takes memory on demand instead of
  reserving 75% away from SBD. Without it, SBD's Davidson basis runs out of room at
  large --max_dim and fails inside tpb_diag with std::bad_alloc — a memory error
  that reads as SBD's fault but is caused by JAX's reservation.
• When SBD solver is running on CPUs, force the subspace expansion to CPU via jax.config.update("jax_platforms", "cpu").
• honours an explicit JAX_PLATFORMS=cpu, and sets neither variable if the caller has.
• resolves --device auto on rank 0 and broadcasts it. _check_cuda() shells out to
  nvidia-smi with a 2 s timeout and measures ~6 s when 8 ranks probe at once, so it
  could time out on some ranks and not others. A diverged device_str sent the ranks
  that resolved to a GPU into jax.distributed.initialize() — a COMM_WORLD rendezvous
  — while the rest skipped it, hanging the job. Broadcasting makes agreement structural,
  and runs one probe instead of eight.

Also dropped: the startup hardware report. print_device_info() re-runs the same
nvidia-smi probe and printed "none detected" on a fully populated node when it timed
out; the per-rank device assignment is now reported instead.

Testing

8 ranks, one node (x86_64 / H100, MPICH) — bundled H2O pool, all five paths exit 0
and reproduce the same energies (round 1 -76.2359466308, round 2 -76.2421767512,
final subspace 1742 x 1742, stopping on closure at round 2):

--device
gpu (Thrust)	✓
gpu-omp (OMP offload)	✓
cpu	✓ (CPU fallback, warns per rank)
auto	✓ (resolves to gpu, all ranks agree)
gpu + preset JAX_PLATFORMS=cpu	✓ (honoured)

• Physical assignment verified via device.local_hardware_id: 0..7 across the
  eight ranks on both GPU backends, unset on cpu. This is the check that read a
  flat 0 before.
• Oversubscription: 16 ranks on 8 GPUs pairs correctly (rank r → r % 8).
• Large subspace: a long multi-round run at a high --max_dim completed with JAX
  and SBD sharing each card and headroom remaining.
• Runs with no launcher wrapper and no JAX_PLATFORMS — the configuration that
  previously failed.
• run_sqd_sbd.py with --device auto on 8 ranks: -76.2359466308.

8 ranks, two nodes (aarch64 / GB200, NVHPC 26.1, MPICH, Slurm) — 4 ranks and 4 GPUs
per node. Each rank lands on a distinct physical GPU within its own node:

rank	node	get_device_id	local_hardware_id
0-3	A	0, 1, 2, 3	0, 1, 2, 3
4-7	B	0, 1, 2, 3	0, 1, 2, 3

CUDA_VISIBLE_DEVICES=0,1,2,3 on every rank there, so the placement comes from this
change rather than the launcher. Measured with the real apply_excitations kernel: the
expansion output reports platform=gpu and SingleDeviceSharding on the assigned card.
The driver also reproduces the H2O reference above across nodes.

Notes for review

• run_sqd_sbd.py gets the same --device auto broadcast, in its own commit. That bug
  is pre-existing and unrelated to JAX — that driver has no
  jax.distributed.initialize — but a diverged device_str leaves ranks on different
  backends (the cpu and gpu extension modules are separate builds) and they then call
  solve_sci_batch, which is collective. Happy to split it out if you'd rather.
• bindings.cpp changed, so all three backends need rebuilding.