From f612759ccc5ab63cb3708a9f3222a6e1e194ba60 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 12:19:33 -0400 Subject: [PATCH 1/5] examples/README: note JAX/MPI conflict for enlarge-subspace driver, unify GPU example ranks enlarge_batch_from_transitions (qiskit-addon-sqd, JAX-based) has no MPI awareness -- every rank redundantly runs it, and if JAX is set up for GPU, every rank tries to grab one at once, which doesn't work. Note it: use a single MPI rank, or JAX_PLATFORMS=cpu to keep this step on CPU regardless of rank count. Also made the two --device gpu example commands consistent with each other: 8 MPI ranks, --adet_comm_size 4 --bdet_comm_size 2, no --task_comm_size (previously 2/2/2). --- python/examples/README.md | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/python/examples/README.md b/python/examples/README.md index da59ac4..b5e1d52 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -27,7 +27,7 @@ mpirun -np 8 python -u run_sbd_diag.py \ --device gpu \ --fcidump ../../vendor/sbd-upstream/data/n2/fcidump.txt \ --adetfile ../../vendor/sbd-upstream/data/n2/1em3-alpha.txt \ - --adet_comm_size 2 --bdet_comm_size 2 --task_comm_size 2 + --adet_comm_size 4 --bdet_comm_size 2 # Retrieve the 1-/2-particle RDMs and save them to a file mpirun -np 2 python -u run_sbd_diag.py --rdm_output /tmp/h2o_rdms.npz @@ -79,7 +79,7 @@ mpirun -np 8 python -u run_sqd_sbd.py \ --fcidump /path/to/fci_dump.txt \ --samples_per_batch 800 --num_batches 3 --max_iterations 10 \ --device gpu \ - --adet_comm_size 2 --bdet_comm_size 2 --task_comm_size 2 + --adet_comm_size 4 --bdet_comm_size 2 ``` **count_dict.json format:** A JSON object mapping bitstrings to shot counts, as @@ -154,6 +154,14 @@ See [SQD Parameters](#sqd-parameters) below for the flags it shares with of `--sqd_carryover_threshold`, and `--max_dim`'s risk profile is sharper here). +The excitation-expansion step between rounds runs through qiskit-addon-sqd's +own JAX-based `enlarge_batch_from_transitions`, which has no MPI awareness — +every rank redundantly runs it on CPU, or (if JAX is set up for GPU in your +environment) every rank tries to grab a GPU for it at once, which doesn't +work. Run this driver with a single MPI rank, or set `JAX_PLATFORMS=cpu` to +keep this step off the GPU regardless of rank count (SBD's own `--device +gpu` diagonalization is unaffected either way). + ### 4. run_sqd_sbd.ipynb — Jupyter walkthrough (serial) Interactive single-rank companion to `run_sqd_sbd.py`. Same SQD self-consistent From 930df2ca7714f2ede49fe8d2f195ff5cb1ed99b1 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 12:30:24 -0400 Subject: [PATCH 2/5] Update README.md shorten the comments --- python/examples/README.md | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/python/examples/README.md b/python/examples/README.md index b5e1d52..447a907 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -15,14 +15,14 @@ Runs a single TPB diagonalization from an FCIDUMP file and alpha determinant file. No SQD loop, no Qiskit dependency. ```bash -# H2O with 2 MPI ranks +# H2O with 2 MPI ranks (--device cpu) mpirun -np 2 python -u run_sbd_diag.py \ --device cpu \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --adetfile ../../vendor/sbd-upstream/data/h2o/h2o-1em3-alpha.txt \ --adet_comm_size 2 -# N2 with GPU +# N2 with GPU (--device gpu) mpirun -np 8 python -u run_sbd_diag.py \ --device gpu \ --fcidump ../../vendor/sbd-upstream/data/n2/fcidump.txt \ @@ -68,11 +68,11 @@ eigensolver backend. Supports two bitstring input modes: ```bash # H2O with the bundled counts file (275 bitstrings -> ~ -76.236 Ha) -mpirun -np 4 python -u run_sqd_sbd.py \ +mpirun -np 8 python -u run_sqd_sbd.py \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --counts count_dict_h2o.json \ - --device cpu \ - --adet_comm_size 2 --bdet_comm_size 2 + --device gpu \ + --adet_comm_size 4 --bdet_comm_size 2 # Custom system with random bitstrings mpirun -np 8 python -u run_sqd_sbd.py \ @@ -138,12 +138,14 @@ both stop moving (`--energy_tol`/`--occupancies_tol`) -- `--max_iterations` is a safety cap, not the expected stopping mechanism. ```bash -mpirun -np 4 python -u run_sqd_enlarge_subspace_sbd.py \ +JAX_PLATFORMS=cpu` mpirun -np 8 python -u run_sqd_enlarge_subspace_sbd.py \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --counts count_dict_h2o.json \ - --device cpu \ - --adet_comm_size 2 --enlarge_threshold 1e-4 + --device gpu \ + --adet_comm_size 4 --bdet_comm_size 2 --enlarge_threshold 1e-4 ``` +Note that qiskit-addon-sqd's own JAX-based `enlarge_batch_from_transitions` has no MPI awareness. Set `JAX_PLATFORMS=cpu` +when running a MPI job using more than 1 rank. Same bundled 275-bitstring H2O pool as `run_sqd_sbd.py`'s own example above: plain SQD reaches **≈ -76.236 Ha** and stops there; this driver keeps going @@ -154,13 +156,7 @@ See [SQD Parameters](#sqd-parameters) below for the flags it shares with of `--sqd_carryover_threshold`, and `--max_dim`'s risk profile is sharper here). -The excitation-expansion step between rounds runs through qiskit-addon-sqd's -own JAX-based `enlarge_batch_from_transitions`, which has no MPI awareness — -every rank redundantly runs it on CPU, or (if JAX is set up for GPU in your -environment) every rank tries to grab a GPU for it at once, which doesn't -work. Run this driver with a single MPI rank, or set `JAX_PLATFORMS=cpu` to -keep this step off the GPU regardless of rank count (SBD's own `--device -gpu` diagonalization is unaffected either way). + ### 4. run_sqd_sbd.ipynb — Jupyter walkthrough (serial) From 1e25c159a09726e079b2a633681407788a8cceba Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 12:44:14 -0400 Subject: [PATCH 3/5] examples/README: fix stray backtick in the enlarge-subspace command The JAX_PLATFORMS=cpu example carried a stray backtick (`JAX_PLATFORMS=cpu` mpirun ...`), which bash reads as an unterminated command substitution -- copy-pasting it fails. Also a blank line before the note so it reads as its own paragraph, 'an MPI job', and drop a double blank line left behind by the earlier trim. Verified all three documented commands on h100 (8 ranks, --device gpu, 4x2 grid): run_sqd_sbd.py H2O gives -76.2359466308 and run_sqd_enlarge_subspace_sbd.py gives -76.2421767512, both matching the documented values. Also confirmed the note's premise directly -- with GPU-enabled JAX and 8 ranks, omitting JAX_PLATFORMS=cpu fails with JaxRuntimeError RESOURCE_EXHAUSTED / CUDA_ERROR_OUT_OF_MEMORY. --- python/examples/README.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/examples/README.md b/python/examples/README.md index 447a907..db5f6f4 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -138,14 +138,15 @@ both stop moving (`--energy_tol`/`--occupancies_tol`) -- `--max_iterations` is a safety cap, not the expected stopping mechanism. ```bash -JAX_PLATFORMS=cpu` mpirun -np 8 python -u run_sqd_enlarge_subspace_sbd.py \ +JAX_PLATFORMS=cpu mpirun -np 8 python -u run_sqd_enlarge_subspace_sbd.py \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --counts count_dict_h2o.json \ --device gpu \ --adet_comm_size 4 --bdet_comm_size 2 --enlarge_threshold 1e-4 ``` + Note that qiskit-addon-sqd's own JAX-based `enlarge_batch_from_transitions` has no MPI awareness. Set `JAX_PLATFORMS=cpu` -when running a MPI job using more than 1 rank. +when running an MPI job using more than 1 rank. Same bundled 275-bitstring H2O pool as `run_sqd_sbd.py`'s own example above: plain SQD reaches **≈ -76.236 Ha** and stops there; this driver keeps going @@ -156,8 +157,6 @@ See [SQD Parameters](#sqd-parameters) below for the flags it shares with of `--sqd_carryover_threshold`, and `--max_dim`'s risk profile is sharper here). - - ### 4. run_sqd_sbd.ipynb — Jupyter walkthrough (serial) Interactive single-rank companion to `run_sqd_sbd.py`. Same SQD self-consistent From 7d9e322ff8743d1e1c1bfe2eacb57f375d0fecd8 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 12:51:05 -0400 Subject: [PATCH 4/5] run_sqd_enlarge_subspace_sbd.py: sync docstring with the README's JAX/MPI note The module docstring still showed 'mpirun -np 4' with no mention of JAX, which the README now documents as a setup that dies with CUDA_ERROR_OUT_OF_MEMORY when JAX is GPU-enabled -- so anyone reading the source or --help instead of the README got the broken recipe. Docstring now carries the same JAX_PLATFORMS=cpu guidance and the same 8-rank / 4x2 GPU command the README uses. Verified the docstring's command verbatim on h100: converges to -76.2421767512, matching the documented value. --- python/examples/run_sqd_enlarge_subspace_sbd.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/examples/run_sqd_enlarge_subspace_sbd.py b/python/examples/run_sqd_enlarge_subspace_sbd.py index bc2b63c..cf9964c 100644 --- a/python/examples/run_sqd_enlarge_subspace_sbd.py +++ b/python/examples/run_sqd_enlarge_subspace_sbd.py @@ -33,10 +33,18 @@ run reaching it before either real criterion is a sign something needs tuning, not the expected happy path. +enlarge_batch_from_transitions is JAX-based and has no MPI awareness. Set +JAX_PLATFORMS=cpu when running on more than 1 rank: otherwise, if JAX is +set up for GPU, every rank tries to grab a GPU for that step at once and +the run dies with CUDA_ERROR_OUT_OF_MEMORY. SBD's own --device gpu +diagonalization is unaffected either way. + Usage (MPI required): - mpirun -np 4 python run_sqd_enlarge_subspace_sbd.py \ + JAX_PLATFORMS=cpu mpirun -np 8 python run_sqd_enlarge_subspace_sbd.py \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --counts count_dict_h2o.json \ + --device gpu \ + --adet_comm_size 4 --bdet_comm_size 2 \ --enlarge_threshold 1e-4 --max_iterations 10 """ From 44c3d07142c6c0681c6b54773225fd24808ef405 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 13:00:01 -0400 Subject: [PATCH 5/5] Update README.md add notes about AMD GPUs --- python/examples/README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/examples/README.md b/python/examples/README.md index db5f6f4..60e4c3e 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -5,7 +5,9 @@ Examples demonstrating SBD's capabilities for quantum chemistry calculations. ## Overview - **Communication:** MPI for distributed computing -- **Backends:** CPU (host OpenMP), GPU (NVHPC Thrust, NVIDIA only) and GPU (OpenMP target offload, NVIDIA and AMD), switchable at runtime via `device` parameter +- **Backends:** CPU (host OpenMP, `--device cpu`), GPU (NVHPC Thrust, NVIDIA only, `--device gpu`) and GPU (OpenMP target offload, NVIDIA and AMD, `--device gpu-omp`), switchable at runtime via `device` parameter + +Replace `--device gpu` in all the examples below with `--device gpu-omp` if you use AMD GPUs. ## Examples @@ -15,14 +17,14 @@ Runs a single TPB diagonalization from an FCIDUMP file and alpha determinant file. No SQD loop, no Qiskit dependency. ```bash -# H2O with 2 MPI ranks (--device cpu) +# H2O with 2 MPI ranks with CPU mpirun -np 2 python -u run_sbd_diag.py \ --device cpu \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --adetfile ../../vendor/sbd-upstream/data/h2o/h2o-1em3-alpha.txt \ --adet_comm_size 2 -# N2 with GPU (--device gpu) +# N2 with GPU mpirun -np 8 python -u run_sbd_diag.py \ --device gpu \ --fcidump ../../vendor/sbd-upstream/data/n2/fcidump.txt \