jaunatisblue
72f95599bb
Some checks failed
Build wheels / build (ubuntu-latest, 3.11) (push) Has been cancelled
Build wheels / build (ubuntu-latest, 3.12) (push) Has been cancelled
Build wheels / build (ubuntu-latest, 3.13) (push) Has been cancelled
Tests / check (push) Has been cancelled
Tests / build (ubuntu-latest, 3.11) (push) Has been cancelled
Tests / build (ubuntu-latest, 3.12) (push) Has been cancelled
Tests / build (ubuntu-latest, 3.13) (push) Has been cancelled
128 lines
3.9 KiB
Bash
Executable File
128 lines
3.9 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
# Large CPU expectation benchmarks for two-server runs.
|
|
#
|
|
# Defaults assume two Intel Xeon Platinum 8558P servers with about 500 GiB RAM
|
|
# each. Override HOSTFILE, PYTHON_BIN, MPIEXEC, or the per-case knobs below as
|
|
# needed.
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
cd "$ROOT_DIR"
|
|
|
|
PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
|
|
MPIEXEC="${MPIEXEC:-mpiexec}"
|
|
HOSTFILE="${HOSTFILE:-hostfile}"
|
|
|
|
MPS_RANKS="${MPS_RANKS:-8}"
|
|
MPS_THREADS="${MPS_THREADS:-12}"
|
|
TN_RANKS="${TN_RANKS:-12}"
|
|
TN_THREADS="${TN_THREADS:-8}"
|
|
|
|
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
|
|
export MKL_NUM_THREADS="${MKL_NUM_THREADS:-1}"
|
|
|
|
run_mpi() {
|
|
local ranks="$1"
|
|
shift
|
|
"$MPIEXEC" -hostfile "$HOSTFILE" -n "$ranks" "$PYTHON_BIN" "$@"
|
|
}
|
|
|
|
run_case() {
|
|
local title="$1"
|
|
shift
|
|
echo
|
|
echo "================================================================================"
|
|
echo "$title"
|
|
echo "================================================================================"
|
|
echo "HOSTFILE=$HOSTFILE PYTHON_BIN=$PYTHON_BIN MPIEXEC=$MPIEXEC"
|
|
echo "OMP_NUM_THREADS=$OMP_NUM_THREADS MKL_NUM_THREADS=$MKL_NUM_THREADS"
|
|
echo "$*"
|
|
"$@"
|
|
}
|
|
|
|
case "${1:-help}" in
|
|
smoke)
|
|
run_case "MPS MPI smoke: n=40 layers=30 bond=2048" \
|
|
run_mpi "$MPS_RANKS" benchmark_cpu_expectation.py \
|
|
--mpi --mps \
|
|
--nqubits "${MPS_SMOKE_NQ:-40}" \
|
|
--nlayers "${MPS_SMOKE_LAYERS:-30}" \
|
|
--bond "${MPS_SMOKE_BOND:-2048}" \
|
|
--torch-threads "$MPS_THREADS" \
|
|
--circuits brickwall_cnot reversed_cnot shifted_cz \
|
|
--observables ring_xz open_zz range2_xx
|
|
|
|
run_case "TN MPI smoke: n=32 layers=16 target_slices=12" \
|
|
run_mpi "$TN_RANKS" benchmark_cpu_expectation.py \
|
|
--mpi \
|
|
--nqubits "${TN_SMOKE_NQ:-32}" \
|
|
--nlayers "${TN_SMOKE_LAYERS:-16}" \
|
|
--torch-threads "$TN_THREADS" \
|
|
--circuits brickwall_cnot shifted_cz rxx_rzz \
|
|
--observables ring_xz open_zz range2_xx \
|
|
--tn-target-slices "${TN_SMOKE_SLICES:-12}"
|
|
;;
|
|
|
|
mps-long)
|
|
run_case "MPS MPI long: n=64 layers=48 bond=4096" \
|
|
run_mpi "$MPS_RANKS" benchmark_cpu_expectation.py \
|
|
--mpi --mps \
|
|
--nqubits "${MPS_LONG_NQ:-64}" \
|
|
--nlayers "${MPS_LONG_LAYERS:-48}" \
|
|
--bond "${MPS_LONG_BOND:-4096}" \
|
|
--torch-threads "$MPS_THREADS" \
|
|
--circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz \
|
|
--observables ring_xz open_zz mixed_local range2_xx
|
|
;;
|
|
|
|
mps-pressure)
|
|
run_case "MPS MPI pressure: n=80 layers=64 bond=4096" \
|
|
run_mpi "$MPS_RANKS" benchmark_cpu_expectation.py \
|
|
--mpi --mps \
|
|
--nqubits "${MPS_PRESSURE_NQ:-80}" \
|
|
--nlayers "${MPS_PRESSURE_LAYERS:-64}" \
|
|
--bond "${MPS_PRESSURE_BOND:-4096}" \
|
|
--torch-threads "$MPS_THREADS" \
|
|
--circuits brickwall_cnot reversed_cnot shifted_cz rxx_rzz swap_scramble \
|
|
--observables ring_xz open_zz mixed_local range2_xx long_z_string
|
|
;;
|
|
|
|
tn-long)
|
|
run_case "TN MPI long: n=36 layers=20 target_slices=24" \
|
|
run_mpi "$TN_RANKS" benchmark_cpu_expectation.py \
|
|
--mpi \
|
|
--nqubits "${TN_LONG_NQ:-36}" \
|
|
--nlayers "${TN_LONG_LAYERS:-20}" \
|
|
--torch-threads "$TN_THREADS" \
|
|
--circuits brickwall_cnot shifted_cz rxx_rzz \
|
|
--observables ring_xz open_zz range2_xx \
|
|
--tn-target-slices "${TN_LONG_SLICES:-24}"
|
|
;;
|
|
|
|
all)
|
|
"$0" smoke
|
|
"$0" mps-long
|
|
"$0" tn-long
|
|
;;
|
|
|
|
help|*)
|
|
cat >&2 <<'EOF'
|
|
Usage: tools/run_cpu_large_cases.sh [smoke|mps-long|mps-pressure|tn-long|all]
|
|
|
|
Common overrides:
|
|
HOSTFILE=hostfile
|
|
PYTHON_BIN=.venv/bin/python
|
|
MPIEXEC=mpiexec
|
|
MPS_RANKS=8 MPS_THREADS=12
|
|
TN_RANKS=12 TN_THREADS=8
|
|
|
|
Scale overrides:
|
|
MPS_LONG_NQ=64 MPS_LONG_LAYERS=48 MPS_LONG_BOND=4096
|
|
MPS_PRESSURE_NQ=80 MPS_PRESSURE_LAYERS=64 MPS_PRESSURE_BOND=4096
|
|
TN_LONG_NQ=36 TN_LONG_LAYERS=20 TN_LONG_SLICES=24
|
|
EOF
|
|
exit 2
|
|
;;
|
|
esac
|