{ "cells": [ { "cell_type": "markdown", "id": "44f669e4", "metadata": {}, "source": [ "# 🔪 The Sharp Bits 🔪" ] }, { "cell_type": "markdown", "id": "cd9c59fa", "metadata": {}, "source": [ "## Compatible devices\n", "\n", "- **NVIDIA GPU**: Heavily tested, best performance\n", "- **AMD GPU**: To be supported in the future\n", "- **CPU**: Supported but not heavily tested\n", "- **Google TPU**: Not supported" ] }, { "cell_type": "markdown", "id": "cd83762e", "metadata": {}, "source": [ "## Global system\n", "\n", "Quantax stores the system geometry and Hilbert space — the {py:class}`~quantax.sites.Sites` (or {py:class}`~quantax.sites.Lattice`) — as a **process-global constant**, retrieved anywhere via {py:func}`quantax.get_sites` and {py:func}`quantax.get_lattice`. This differs from most NQS packages, where the system is an ordinary object passed around explicitly.\n", "\n", "The consequence is that **only one `Sites` should exist per Python process.** Constructing a second one does not raise an error; it merely warns and silently overwrites the first, which can lead to hard-to-debug behavior. This is especially easy to trigger in notebooks by re-running a cell that builds the system.\n", "\n", "To switch to a different system, restart the Python interpreter (or the notebook kernel)." ] }, { "cell_type": "markdown", "id": "4f4464ee", "metadata": {}, "source": [ "## Randomness\n", "\n", "Quantax keeps a single global PRNG key, replicated across all devices. Set it with {py:func}`quantax.set_random_seed` (the default seed is 42) for reproducible runs, and draw fresh subkeys with {py:func}`quantax.get_subkeys`.\n", "\n", "Because `get_subkeys` reads and updates this global key, it is **not jittable**. Call it outside jitted functions and pass the resulting keys in as arguments." ] }, { "cell_type": "markdown", "id": "594ce3ca", "metadata": {}, "source": [ "## Data precision\n", "\n", "By default, JAX\n", "\n", "- disables double precision, so `float64`/`complex128` silently fall back to `float32`/`complex64`;\n", "- uses TF32, a reduced-precision format, for the matmuls and convolutions of `float32` inputs on supported GPUs, whenever possible.\n", "\n", "Quantax does **NOT** change these defaults. Most computations work well under this setting, but be careful in the following cases.\n", "\n", "- **Matrix inversion / linear solves.** Several solvers in `quantax/optimizer/solver.py` switch to FP64 (`with jax.enable_x64()`) or FP32 matmuls (`precision='highest'`) locally to keep the inverse accurate.\n", "- **Local updates.** Higher precision might be required in {doc}`local updates <../tutorials/local_updates>` to avoid error accumulation.\n", "\n", "To enable double precision globally, call `jax.config.update(\"jax_enable_x64\", True)`. Alternatively, call `quantax.set_default_dtype(jnp.float64)` to set `jnp.float64` as the default data type in Quantax, which also turns on `jax_enable_x64`. To disable TF32, call `jax.config.update(\"jax_default_matmul_precision\", \"float32\")`.\n", "\n", "The variational models in `quantax.model` deliberately ignore the Quantax default dtype and pick their own (usually `float32`) for efficiency, so {py:func}`quantax.set_default_dtype` does not change them." ] }, { "cell_type": "markdown", "id": "231933ed", "metadata": {}, "source": [ "## To JIT or not to JIT\n", "\n", "[`jax.jit`](https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html#just-in-time-compilation-with-jax-jit) greatly accelerates JAX functions, but it imposes constraints on the jitted function and can make the code harder to read.\n", "\n", "In Quantax, our strategy is to apply `jax.jit` only to functions that are likely to be the bottleneck of a simulation, and to keep the rest unjitted for flexibility. For instance, {py:meth}`quantax.operator.Operator.expectation` is left unjitted so that it can compute the expectation of arbitrary states, while the forward pass of {py:class}`quantax.state.Variational` — the actual bottleneck when computing expectation values — is jitted." ] }, { "cell_type": "markdown", "id": "ecfd1436", "metadata": {}, "source": [ "## Wavefunction overflow\n", "\n", "Many-body wavefunctions often span a wide range of magnitudes, even beyond the limit of FP64. In many packages, e.g. [NetKet](https://www.netket.org/) and [jVMC](https://jvmc.readthedocs.io/en/latest/), the wavefunction amplitudes are expressed by $\\log \\psi$ instead of $\\psi$ to avoid overflow (and underflow).\n", "\n", "In Quantax, we don't utilize this approach for the following reasons.\n", "\n", "1. Ground-state wavefunctions are often real-valued with signs. Expressing them using $\\log \\psi$ forces one to perform optimization in the complex space.\n", "2. $\\log\\psi$ is discontinuous at $\\psi=0$, which might cause stability issues.\n", "\n", "Instead, Quantax keeps $\\psi$ but factors out its magnitude into a separate, overflow-safe scale. Two custom array types in `quantax.utils` implement this, each a PyTree with two leaves:\n", "\n", "- {py:class}`~quantax.utils.LogArray` — `value = sign * exp(logabs)`, with `sign` being $\\pm 1$ or a complex phase and `logabs` the real log-magnitude.\n", "- {py:class}`~quantax.utils.ScaleArray` — `value = significand * exp(exponent)`, with `exponent` a real normalization factor pulled out to keep `significand` in range.\n", "\n", "The sign/phase stays attached to the amplitude, so a real sign-structured wavefunction remains real and continuous at $\\psi=0$. Arithmetic is overloaded to act in the scaled space (e.g. multiplying two `LogArray`s adds the `logabs`, never forming the overflowing product); convert back to `jax.Array` with `arr.value()` or `jnp.asarray(arr)`.\n", "\n", "For details of constructing neural wavefunctions with `LogArray` and `ScaleArray`, see tutorial {doc}`build your network <../tutorials/build_net>`." ] }, { "cell_type": "markdown", "id": "b15dd4a6", "metadata": {}, "source": [ "## Sharding\n", "\n", "Quantax uses JAX's auto sharding by default. For details, see the JAX guide [Distributed arrays and automatic parallelization](https://docs.jax.dev/en/latest/parallel.html).\n", "\n", "The mesh has two axes, `(\"process\", \"device\")`, with shape `(jax.process_count(), jax.local_device_count())`.\n", "\n", "Some arrays, such as neural-network weights, are replicated across all devices (see {py:func}`~quantax.utils.get_replicated_sharding`). Others, such as Monte Carlo samples, are distributed over the devices (see {py:func}`~quantax.utils.get_distributed_sharding`)." ] }, { "cell_type": "markdown", "id": "35a9dcdb", "metadata": {}, "source": [ "## HPC usage\n", "\n", "There are 3 ways to submit JAX jobs on HPC. Here we provide the caveats in the python script and the slurm scripts on the NERSC cluster as examples. In all cases, the users only need a few lines of code to control the distributed initialization of JAX, while the Quantax API works the same way as on a single device." ] }, { "cell_type": "markdown", "id": "4380d220", "metadata": {}, "source": [ "### Single-node, only one process\n", "\n", "In the single-node case, one doesn't have to do anything special in the python script. The slurm script is given below." ] }, { "cell_type": "code", "execution_count": null, "id": "9caebf86", "metadata": { "vscode": { "languageId": "shellscript" } }, "outputs": [], "source": [ "#!/bin/bash -l\n", "#SBATCH --nodes=1\n", "#SBATCH --ntasks-per-node=1\n", "#SBATCH --gpus-per-task=\n", "#SBATCH --output=gpu-out.%j\n", "#SBATCH --error=gpu-err.%j\n", "#SBATCH --time=\n", "#SBATCH --constraint=gpu\n", "#SBATCH --qos=regular\n", "#SBATCH --account=\n", "\n", "source /bin/activate\n", "\n", "srun --cpu-bind=none python -u " ] }, { "cell_type": "markdown", "id": "1eecfa4b", "metadata": {}, "source": [ "### Multi-node, one process per GPU" ] }, { "cell_type": "code", "execution_count": null, "id": "slurm-multinode-per-gpu", "metadata": { "vscode": { "languageId": "shellscript" } }, "outputs": [], "source": [ "#!/bin/bash -l\n", "#SBATCH --nodes=\n", "#SBATCH --ntasks-per-node=4 # 4 GPUs / node on NERSC\n", "#SBATCH --gpus-per-task=1\n", "#SBATCH --gpu-bind=none\n", "#SBATCH --mem=0\n", "#SBATCH --output=gpu-out.%j\n", "#SBATCH --error=gpu-err.%j\n", "#SBATCH --time=\n", "#SBATCH --constraint=gpu\n", "#SBATCH --qos=regular\n", "#SBATCH -c 32\n", "#SBATCH --account=\n", "\n", "source /bin/activate\n", "\n", "# Specific to multi-node jobs on NERSC\n", "module load nccl/2.24.3\n", "\n", "# Some JAX versions fail without turning off xla_gpu_shard_autotuning\n", "export XLA_FLAGS=\"--xla_gpu_shard_autotuning=false\"\n", "\n", "srun --cpu-bind=cores python -u " ] }, { "cell_type": "markdown", "id": "85f5d112", "metadata": {}, "source": [ "One should configure the distributed initialization of JAX in the beginning of the python script." ] }, { "cell_type": "code", "execution_count": null, "id": "e2509949", "metadata": {}, "outputs": [], "source": [ "import jax\n", "jax.distributed.initialize()\n", "\n", "..." ] }, { "cell_type": "markdown", "id": "7e3fcc93", "metadata": {}, "source": [ "### Multi-node, one process per node" ] }, { "cell_type": "code", "execution_count": null, "id": "4600efca", "metadata": { "vscode": { "languageId": "shellscript" } }, "outputs": [], "source": [ "#!/bin/bash -l\n", "#SBATCH --nodes=\n", "#SBATCH --ntasks-per-node=1\n", "#SBATCH --gpus-per-node=4 # 4 GPUs / node on NERSC\n", "#SBATCH --gpu-bind=none\n", "#SBATCH --mem=0\n", "#SBATCH --output=gpu-out.%j\n", "#SBATCH --error=gpu-err.%j\n", "#SBATCH --time=\n", "#SBATCH --constraint=gpu\n", "#SBATCH --qos=regular\n", "#SBATCH -c 128\n", "#SBATCH --account=\n", "\n", "source /bin/activate\n", "\n", "# Specific to multi-node jobs on NERSC\n", "module load nccl/2.24.3\n", "\n", "# Some JAX versions fail without turning off xla_gpu_shard_autotuning\n", "export XLA_FLAGS=\"--xla_gpu_shard_autotuning=false\"\n", "\n", "srun --cpu-bind=none python -u " ] }, { "cell_type": "markdown", "id": "f6fe965e", "metadata": {}, "source": [ "`jax.distributed.initialize` often fails to handle this case automatically. The following code resolves it." ] }, { "cell_type": "code", "execution_count": null, "id": "48ae6022", "metadata": {}, "outputs": [], "source": [ "import os\n", "import jax\n", "\n", "gpus_per_node = int(os.environ[\"SLURM_GPUS_ON_NODE\"])\n", "\n", "jax.distributed.initialize(\n", " num_processes=int(os.environ[\"SLURM_NTASKS\"]),\n", " process_id=int(os.environ[\"SLURM_PROCID\"]),\n", " local_device_ids=list(range(gpus_per_node)),\n", ")\n", "\n", "..." ] } ], "metadata": { "kernelspec": { "display_name": "quantax_env", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12.11" } }, "nbformat": 4, "nbformat_minor": 5 }