diff --git a/.github/workflows/mpi_tests.yml b/.github/workflows/mpi_tests.yml new file mode 100644 index 0000000000..66c2d2514a --- /dev/null +++ b/.github/workflows/mpi_tests.yml @@ -0,0 +1,141 @@ +name: MPI tests + +on: + push: + branches: + - master + - dev + pull_request: + branches: + - master + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + mpi_tests: + runs-on: ubuntu-latest + env: + GH_TOKEN: ${{ github.token }} + + strategy: + matrix: + python-version: ["3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Filter changes + id: changes + uses: dorny/paths-filter@v3 + with: + filters: | + has_changes: + - 'desc/**' + - 'tests/test_mpi*' + - 'requirements.txt' + - 'devtools/dev-requirements.txt' + - 'setup.cfg' + - '.github/workflows/mpi_tests.yml' + + - name: Check for relevant changes + id: check_changes + run: echo "has_changes=${{ !contains(github.event.pull_request.labels.*.name, 'only-docs-comments') && steps.changes.outputs.has_changes }}" >> $GITHUB_ENV + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Check full Python version + run: | + python --version + python_version=$(python --version 2>&1 | cut -d' ' -f2) + echo "Python version: $python_version" + echo "version=$python_version" >> $GITHUB_ENV + + - name: Install MPI (OpenMPI + compiler) + if: env.has_changes == 'true' + run: | + sudo apt-get update + sudo apt-get install -y libopenmpi-dev openmpi-bin + + - name: Set up virtual environment + run: | + python -m venv .venv-${{ env.version }} + source .venv-${{ env.version }}/bin/activate + python -m pip install --upgrade pip + pip install -r devtools/dev-requirements.txt + pip install matplotlib==3.9.2 + # install mpi4py after OpenMPI is available + pip install mpi4py + + - name: Action Details + if: env.has_changes == 'true' + run: | + source .venv-${{ env.version }}/bin/activate + which python + python --version + mpirun --version + pip list + + - name: Test with pytest (MPI setup) + if: env.has_changes == 'true' + run: | + source .venv-${{ env.version }}/bin/activate + python -m pytest -v -m mpi_setup\ + --durations=0 \ + --cov-report xml:cov.xml \ + --cov-config=setup.cfg \ + --cov=desc/ \ + --db ./prof.db + + - name: Test with pytest (MPI run) + if: env.has_changes == 'true' + run: | + source .venv-${{ env.version }}/bin/activate + # ensure each MPI rank writes to a different coverage file + mpirun -n 3 --oversubscribe \ + bash -c 'COVERAGE_FILE=.coverage.$OMPI_COMM_WORLD_RANK \ + python -m pytest -v -m mpi_run \ + --durations=0 \ + --cov=desc/ \ + --cov-config=setup.cfg \ + --cov-append \ + --cov-report=' + + - name: Run MPI tutorials + if: env.has_changes == 'true' + run: | + source .venv-${{ env.version }}/bin/activate + mpirun -n 2 --oversubscribe python docs/notebooks/tutorials/mpi-tutorials/mpi-eq-solve.py + mpirun -n 2 --oversubscribe python docs/notebooks/tutorials/mpi-tutorials/mpi-proximal.py + + - name: Combine coverage files + if: always() && env.has_changes == 'true' + run: | + source .venv-${{ env.version }}/bin/activate + coverage combine + coverage xml -o cov.xml + + - name: save coverage file + if: always() && env.has_changes == 'true' + uses: actions/upload-artifact@v4 + with: + name: mpi_test_artifact-${{ matrix.python-version }} + path: | + ./cov.xml + ./mpl_results.html + ./prof.db + + - name: Upload coverage + if: env.has_changes == 'true' + id : codecov + uses: codecov/codecov-action@v5 + with: + name: codecov-umbrella + files: ./cov.xml + fail_ci_if_error: true + verbose: true diff --git a/.github/workflows/notebook_tests.yml b/.github/workflows/notebook_tests.yml index 016243ac06..8031eb0a4c 100644 --- a/.github/workflows/notebook_tests.yml +++ b/.github/workflows/notebook_tests.yml @@ -16,13 +16,12 @@ concurrency: jobs: notebook_tests: - runs-on: ubuntu-latest env: GH_TOKEN: ${{ github.token }} strategy: matrix: - python-version: ['3.10'] + python-version: ["3.10"] group: [1, 2, 3] steps: @@ -45,7 +44,6 @@ jobs: id: check_changes run: echo "has_changes=${{ !contains(github.event.pull_request.labels.*.name, 'only-docs-comments') && steps.changes.outputs.has_changes }}" >> $GITHUB_ENV - - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v7 with: @@ -94,6 +92,7 @@ jobs: pytest -v --nbmake "./docs/notebooks" \ --nbmake-timeout=2000 \ --ignore=./docs/notebooks/zernike_eval.ipynb \ + --ignore=./docs/notebooks/tutorials/multi_device.ipynb \ --splits 3 \ --group ${{ matrix.group }} \ --splitting-algorithm least_duration diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a46a132cc..3a9ad82315 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,11 @@ New Features - Adds ``desc.objectives.DeflationOperator``, a new objective class which can be used to apply deflation techniques to equilibrium and optimization problems to find multiple local minima or multiple solutions from a single initial point, either by wrapping an existing ``desc.objectives._Objective`` object or by including as an additional penalty or constraint. Also adds a tutorial showing this functionality. - Sub-objectives of an `ObjectiveFunction` can now have different `use_jit` values than the `ObjectiveFunction`. These objectives have to be built before building the `ObjectiveFunction`. - Adds ``num_neighbors`` parameter to ``CoilSetMinDistance`` that limits the pairwise distance computation to the nearest neighbors per coil, reducing memory useage for large coilsets. +- Adds initial support for multi-device optimization with MPI. This allows to compute derivatives and costs on multiple devices (GPUs/CPUs), and to split memory usage during these operations across devices. See the [documentation](https://desc-docs.readthedocs.io/en/stable/notebooks/tutorials/multi_device.html) for details. Couple important notes: + - MPI is not a default dependency of DESC, so, to use MPI functionality, the users should verify their MPI installation themselves. + - Using MPI is recommended only for the cases where you get out-of-memory error. If your problem fits to single GPU memory, it's unlikely that MPI will give speed improvement. + - MPI is not implemented for matrix decompositions (i.e. QR/SVD/Cholesky) which default optimizer ``lsq-exact`` uses. For the cases where Jacobian doesn't fit to GPU memory, matrix decompositions will be performed on CPU and will be slow. Feel free to open a PR, if you have knowledge on parallel QR/SVD or Cholesky. + - CUDA-aware MPI is not supported yet. - Method to plot frequency spectrum of inverse stream map in field line coordinates ``Bounce2D.plot_angle_spectrum``. - Method to compute bounce integrals in batches is now added to the public API ``Bounce2D.batch``. - Initiated deprecation of ``Bounce2D.compute_fieldline_length`` in favor of ``eq.compute("V_psi")``. @@ -154,7 +159,6 @@ New Features - `chunk_size` argument is now used for chunking the number of field lines. For the chunking of Biot-Savart integration for the magnetic field, users can use `bs_chunk_size` instead. - Bug Fixes - Fixes straight field line equilibrium conversion, see #1880 diff --git a/desc/__init__.py b/desc/__init__.py index 30842a4055..c9e275f75a 100644 --- a/desc/__init__.py +++ b/desc/__init__.py @@ -2,10 +2,13 @@ import importlib import os +import platform import re +import subprocess import warnings import colorama +import psutil from termcolor import colored from ._version import get_versions @@ -58,37 +61,133 @@ def __getattr__(name): BANNER = colored(_BANNER, "magenta") +# mpi-cuda = True is not supported yet +config = { + "devices": None, + "avail_mems": None, + "kind": None, + "num_device": None, + # Set to True if CUDA-aware MPI is installed (not tested) + "mpi-cuda": False, + # Suppress the warning in `desc.backend.safe_transfer_to_device` + "SUPPRESS_GPU_MEMORY_WARNING": False, +} -config = {"device": None, "avail_mem": None, "kind": None} +def _get_processor_name(): + """Get the processor name of the current system.""" + if platform.system() == "Windows": + return platform.processor() + elif platform.system() == "Darwin": + os.environ["PATH"] = os.environ["PATH"] + os.pathsep + "/usr/sbin" + command = "sysctl -n machdep.cpu.brand_string" + return subprocess.check_output(command).strip() + elif platform.system() == "Linux": + command = "cat /proc/cpuinfo" + all_info = subprocess.check_output(command, shell=True).decode().strip() + for line in all_info.split("\n"): + if "model name" in line: + return re.sub(pattern=".*model name.*:", repl="", string=line, count=1) + return "CPU" -def set_device(kind="cpu", gpuid=None): + +def _set_cpu_count(n): + """Divide 1 physical CPU into multiple virtual CPUs. + + By default, JAX sees the whole CPU as a single device, regardless of the number of + cores or threads. It then uses multiple cores and threads for lower level + parallelism within individual operations. + + Alternatively, you can force JAX to expose a given number of "virtual" CPUs that + can then be used manually for higher level parallelism (as in at the level of + multiple objective functions.) + + This function is mainly for testing on CI purposes of the parallelism in DESC. + It won't use multiple CPUs even if there are multiple CPUs available on the + machine. It will just divide the first CPU into multiple virtual CPUs. + + Parameters + ---------- + n : int + Number of virtual CPUs for high level parallelism. + + Notes + ----- + This function must be called before importing anything else from DESC or JAX, + and before calling ``desc.set_device``, otherwise it will have no effect. + """ + xla_flags = os.getenv("XLA_FLAGS", "") + xla_flags = re.sub( + r"--xla_force_host_platform_device_count=\S+", "", xla_flags + ).split() + os.environ["XLA_FLAGS"] = " ".join( + [f"--xla_force_host_platform_device_count={n}"] + xla_flags + ) + + +def set_device(kind="cpu", gpuid=None, num_device=1, mpi=None): # noqa: C901 """Sets the device to use for computation. If kind==``'gpu'`` and a gpuid is specified, uses the specified GPU. If - gpuid==``None`` or a wrong GPU id is given, checks available GPUs and selects the - one with the most available memory. - Respects environment variable CUDA_VISIBLE_DEVICES for selecting from multiple - available GPUs + gpuid==``None`` or a wrong GPU id is given, checks available GPUs and selects + the one with the most available memory. Respects environment variable + `CUDA_VISIBLE_DEVICES` for selecting from multiple available GPUs. + + Notes + ----- + This function must be called before importing anything else from DESC or JAX, + otherwise it will have no effect. Parameters ---------- kind : {``'cpu'``, ``'gpu'``} - whether to use CPU or GPU. + Whether to use CPU or GPU. + gpuid : int, optional + GPU id to use. Default is None. Supported only when num_device is 1. + num_device : int, optional + Number of devices to use. For `cpu`, this is the number of nodes. + For `gpu`, this is equal to the number of GPUs connected to a single node. + Default is 1. + mpi : MPI object, optional + MPI communicator. Used to get distinct CPU information for multi-node + jobs where each rank runs on different node. Communicator is not used + if the backend is ``'gpu'``. Supplying communicator doesn't + change the computations, it can only change the output of + ``desc.backend.print_backend_info()``. """ config["kind"] = kind + config["num_device"] = num_device + + cpu_mem = psutil.virtual_memory().available / 1024**3 # RAM in GB + cpu_info = _get_processor_name() + config["cpu_info"] = f"{cpu_info} CPU" + config["cpu_mem"] = cpu_mem + if kind == "cpu": os.environ["JAX_PLATFORMS"] = "cpu" os.environ["CUDA_VISIBLE_DEVICES"] = "" - import psutil - - cpu_mem = psutil.virtual_memory().available / 1024**3 # RAM in GB - config["device"] = "CPU" - config["avail_mem"] = cpu_mem + if num_device == 1: + config["devices"] = [f"{cpu_info} CPU"] + config["avail_mems"] = [cpu_mem] + else: + if mpi is None: + warnings.warn( + "To get the full list of CPUs, provide the MPI communicator.", + UserWarning, + ) + # return the same device multiple times + cpu_names = [f"{i} {cpu_info}" for i in range(num_device)] + else: + comm = mpi.COMM_WORLD + rank = comm.Get_rank() + cpu_name = f"{rank} {cpu_info}" + cpu_names = comm.allgather(cpu_name) + config["devices"] = cpu_names + # This memory is not individual but the total memory + config["avail_mems"] = [cpu_mem] * num_device - if kind == "gpu": - # Set CUDA_DEVICE_ORDER so the IDs assigned by CUDA match those from nvidia-smi + elif kind == "gpu": os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # pynvml namespace is exposed through nvidia-ml-py from pynvml import ( @@ -137,55 +236,56 @@ def _gpu_info(): set_device(kind="cpu") return - maxmem = 0 - selected_gpu = None gpu_ids = [dev["index"] for dev in devices] if "CUDA_VISIBLE_DEVICES" in os.environ: cuda_ids = [ s for s in re.findall(r"\b\d+\b", os.environ["CUDA_VISIBLE_DEVICES"]) ] - # check that the visible devices actually exist and are gpus gpu_ids = [i for i in cuda_ids if i in gpu_ids] if len(gpu_ids) == 0: - # cuda visible devices = '' -> don't use any gpu warnings.warn( colored( - ( - "CUDA_VISIBLE_DEVICES={} ".format( - os.environ["CUDA_VISIBLE_DEVICES"] - ) - + "did not match any physical GPU " - + "(id={}), falling back to CPU".format( - [dev["index"] for dev in devices] - ) - ), + f"CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']} did " + "not match any physical GPU " + f"(id={[dev['index'] for dev in devices]}), falling back to CPU", "yellow", ) ) set_device(kind="cpu") return + devices = [dev for dev in devices if dev["index"] in gpu_ids] + memories = {dev["index"]: dev["mem_total"] - dev["mem_used"] for dev in devices} + + if num_device == 1: + selected_gpu = max( + devices, key=lambda dev: dev["mem_total"] - dev["mem_used"] + ) + if gpuid is not None: + if str(gpuid) in gpu_ids: + selected_gpu = next( + dev for dev in devices if dev["index"] == str(gpuid) + ) + else: + warnings.warn( + colored( + f"Specified gpuid {gpuid} not found, selecting GPU with " + "most memory", + "yellow", + ) + ) + devices = [selected_gpu] - if gpuid is not None and (str(gpuid) in gpu_ids): - selected_gpu = [dev for dev in devices if dev["index"] == str(gpuid)][0] else: - for dev in devices: - mem = dev["mem_total"] - dev["mem_used"] - if mem > maxmem: - maxmem = mem - selected_gpu = dev - config["device"] = selected_gpu["type"] + " (id={})".format( - selected_gpu["index"] - ) - if gpuid is not None and not (str(gpuid) in gpu_ids): - warnings.warn( - colored( - "Specified gpuid {} not found, falling back to ".format(str(gpuid)) - + config["device"], - "yellow", + if num_device > len(devices): + raise ValueError( + f"Requested {num_device} GPUs, but only {len(devices)} available" ) - ) - config["avail_mem"] = ( - selected_gpu["mem_total"] - selected_gpu["mem_used"] - ) / 1024 # in GB - os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_gpu["index"]) + if gpuid is not None: + # TODO: implement multiple GPU selection + raise ValueError("Cannot specify `gpuid` when requesting multiple GPUs") + + devs = devices[:num_device] + config["avail_mems"] = [memories[dev["index"]] / 1024 for dev in devs] # in GB + config["devices"] = [f"{dev['type']} (id={dev['index']})" for dev in devs] + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(dev["index"]) for dev in devs) diff --git a/desc/backend.py b/desc/backend.py index 362ad901a8..4932606c7f 100644 --- a/desc/backend.py +++ b/desc/backend.py @@ -19,7 +19,7 @@ use_jax = False set_device(kind="cpu") else: - if desc_config.get("device") is None: + if desc_config.get("devices") is None: set_device("cpu") try: with warnings.catch_warnings(): @@ -54,15 +54,48 @@ def print_backend_info(): if use_jax: print( f"Using JAX backend: jax version={jax.__version__}, " - + f"jaxlib version={jaxlib.__version__}, dtype={y.dtype}." + f"jaxlib version={jaxlib.__version__}, dtype={y.dtype}." ) else: print(f"Using NumPy backend: version={np.__version__}, dtype={y.dtype}.") - print( - "Using device: {}, with {:.2f} GB available memory.".format( - desc_config.get("device"), desc_config.get("avail_mem") + + if desc_config["kind"] == "cpu": + if desc_config["num_device"] == 1: + print( + f"CPU Info: {desc_config['cpu_info']} with " + f"{desc_config['cpu_mem']:.2f} GB available memory" + ) + else: + print( + f"Using {desc_config['num_device']} CPUs with " + + f"{desc_config['avail_mems'][0]:.2f} GB total available memory:" + ) + for dev in desc_config["devices"]: + print(f"\t CPU : {dev}") + + print( + "\nNote: The backend information assumes that the user has 1 " + "process per CPU (node). Using multiple processes per CPU (node) is " + "not the most efficient way to use MPI with purely CPUs." + ) + elif desc_config["kind"] == "gpu": + print( + f"CPU Info: {desc_config['cpu_info']} with " + f"{desc_config['cpu_mem']:.2f} GB available memory" ) - ) + print(f"Using {desc_config['num_device']} device:") + for i, dev in enumerate(desc_config["devices"]): + print( + f"\t Device : {dev} with {desc_config['avail_mems'][i]:.2f} " + "GB available memory" + ) + + if desc_config["num_device"] != 1: + print( + "\nNote: The backend information only reflects the devices for " + "the current process. The full set of devices used by other processes " + "may be different." + ) def _diag_to_full(d, e): @@ -638,6 +671,69 @@ def bodyfun(state): x = jax.lax.custom_root(res, x0, solve, _tangent_solve, has_aux=False) return x + def safe_mpi_Bcast(arr, comm, root=0): + """Safe Bcast function for Jax arrays. + + JAX arrays cannot be directly broadcasted using MPI's Bcast, but numpy + arrays can. If CUDA-aware MPI is available, JAX arrays on GPU can be + broadcasted directly. This function checks the type of the array and + perform the broadcast safely. + + Parameters + ---------- + arr : jnp.ndarray or np.ndarray + Array to broadcast. + comm : MPI.Comm + MPI communicator. + root : int + Rank of root process. Default is 0. + + Returns + ------- + arr : jnp.ndarray or np.ndarray + Broadcasted array. + """ + if not desc_config["mpi-cuda"]: + arr = np.array(arr) + comm.Bcast(arr, root=root) + # don't use this returned value for root == rank, as it will replace the jax + # array with a numpy array, which is not ideal + return arr + + def safe_transfer_to_device(arr): + """Safely transfer array to device. + + Handles the final array device if the array is too big for GPU, + or the backend is CPU. + + Parameters + ---------- + arr : jnp.ndarray or np.ndarray + Array to transfer. + + Returns + ------- + arr : jnp.ndarray + Array on the target device. + """ + size_gb = arr.nbytes / 1024**3 + + # this can still fail if arr is big even for CPU + if desc_config["kind"] == "cpu" or size_gb > desc_config["avail_mems"][0] * 0.9: + if ( + not desc_config["SUPPRESS_GPU_MEMORY_WARNING"] + and desc_config["kind"] == "gpu" + ): + warnings.warn( + "The total size of the arrays exceeds the available memory of the " + "GPU[id=0]. Moving the array to CPU. This may cause performance " + "degredation. To suppress this warning, use \n" + "`from desc import config as desc_config` \n" + "`desc_config['SUPPRESS_GPU_MEMORY_WARNING'] = True`" + ) + return jnp.asarray(arr, device=jax.devices("cpu")[0]) + return jnp.asarray(arr) + # we can't really test the numpy backend stuff in automated testing, so we ignore it # for coverage purposes @@ -1113,3 +1209,11 @@ def take( else: out = np.take(a, indices, axis, out, mode) return out + + def safe_mpi_Bcast(arr, comm, root=0): + """Numpy implementation of desc.backend.safe_mpi_Bcast.""" + return comm.Bcast(arr, root=root) + + def safe_transfer_to_device(arr): + """Numpy implementation of desc.backend.safe_transfer_to_device.""" + return arr diff --git a/desc/objectives/_bootstrap.py b/desc/objectives/_bootstrap.py index d149a852cc..b952308e6e 100644 --- a/desc/objectives/_bootstrap.py +++ b/desc/objectives/_bootstrap.py @@ -72,6 +72,8 @@ def __init__( degree=None, name="Bootstrap current self-consistency (Redl)", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -91,6 +93,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index eaafe93a65..83d9deafa5 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -79,6 +79,8 @@ def __init__( grid=None, name=None, jac_chunk_size=None, + device_id=0, + rank=None, ): self._grid = grid self._data_keys = data_keys @@ -95,6 +97,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): # noqa:C901 @@ -397,6 +401,8 @@ def __init__( grid=None, name="coil length", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 2 * np.pi @@ -414,6 +420,8 @@ def __init__( grid=grid, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -501,6 +509,8 @@ def __init__( grid=None, name="coil curvature", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (0, 1) @@ -518,6 +528,8 @@ def __init__( grid=grid, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -601,6 +613,8 @@ def __init__( grid=None, name="coil torsion", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -618,6 +632,8 @@ def __init__( grid=grid, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -701,6 +717,8 @@ def __init__( grid=None, name="coil current length", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -717,6 +735,8 @@ def __init__( grid=grid, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -809,6 +829,8 @@ def __init__( grid=None, name="coil integrated curvature", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 2 * np.pi @@ -825,6 +847,8 @@ def __init__( grid=grid, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -944,6 +968,8 @@ def __init__( softmin_alpha=1.0, dist_chunk_size=None, num_neighbors=None, + device_id=0, + rank=None, ): from desc.coils import CoilSet @@ -970,6 +996,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1154,6 +1182,8 @@ def __init__( use_softmin=False, softmin_alpha=1.0, dist_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (0, 1) @@ -1189,6 +1219,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1439,6 +1471,8 @@ def __init__( use_softmin=False, softmin_alpha=1.0, dist_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (1, np.inf) @@ -1462,6 +1496,8 @@ def __init__( use_softmin=use_softmin, softmin_alpha=softmin_alpha, dist_chunk_size=dist_chunk_size, + device_id=device_id, + rank=rank, ) @@ -1510,6 +1546,8 @@ def __init__( deriv_mode="auto", grid=None, name="coil arclength variance", + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -1526,6 +1564,8 @@ def __init__( deriv_mode=deriv_mode, grid=grid, name=name, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1669,6 +1709,8 @@ def __init__( vacuum=False, name="Quadratic flux", jac_chunk_size=None, + device_id=0, + rank=None, *, bs_chunk_size=None, B_plasma_chunk_size=None, @@ -1702,6 +1744,8 @@ def __init__( normalize_target=normalize_target, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1890,6 +1934,8 @@ def __init__( name="Surface Quadratic Flux", field_fixed=False, jac_chunk_size=None, + device_id=0, + rank=None, *, bs_chunk_size=None, **kwargs, @@ -1915,6 +1961,8 @@ def __init__( normalize_target=normalize_target, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -2113,6 +2161,8 @@ def __init__( field_fixed=False, eq_fixed=False, jac_chunk_size=None, + device_id=0, + rank=None, *, bs_chunk_size=None, **kwargs, @@ -2147,6 +2197,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -2380,6 +2432,8 @@ def __init__( deriv_mode="auto", jac_chunk_size=None, name="linking current", + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -2400,6 +2454,8 @@ def __init__( deriv_mode=deriv_mode, jac_chunk_size=jac_chunk_size, name=name, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -2567,6 +2623,8 @@ def __init__( deriv_mode="auto", jac_chunk_size=None, name="coil-coil linking number", + device_id=0, + rank=None, ): from desc.coils import CoilSet @@ -2589,6 +2647,8 @@ def __init__( deriv_mode=deriv_mode, jac_chunk_size=jac_chunk_size, name=name, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -2726,6 +2786,8 @@ def __init__( regularization="K", source_grid=None, name="surface-current-regularization", + device_id=0, + rank=None, ): from desc.magnetic_fields import ( CurrentPotentialField, @@ -2766,6 +2828,8 @@ def __init__( deriv_mode=deriv_mode, jac_chunk_size=jac_chunk_size, name=name, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index 2ee7f8deb9..e8988889ba 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -61,6 +61,8 @@ def __init__( grid=None, name="force", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -76,6 +78,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -215,6 +219,8 @@ def __init__( grid=None, name="force-anisotropic", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -230,6 +236,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -351,6 +359,8 @@ def __init__( grid=None, name="radial force", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -366,6 +376,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -487,6 +499,8 @@ def __init__( grid=None, name="helical force", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -502,6 +516,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -624,6 +640,8 @@ def __init__( gamma=0, name="energy", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -640,6 +658,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -766,6 +786,8 @@ def __init__( grid=None, name="current density", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -781,6 +803,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_fast_ion.py b/desc/objectives/_fast_ion.py index a1b169c110..029b44b1bc 100644 --- a/desc/objectives/_fast_ion.py +++ b/desc/objectives/_fast_ion.py @@ -119,6 +119,8 @@ def __init__( spline=True, use_bounce1d=False, Nemov=True, + device_id=0, + rank=None, **kwargs, ): try: @@ -171,6 +173,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index 2025623054..b92a280a96 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -91,6 +91,8 @@ def __init__( field_fixed=False, name="Vacuum boundary error", jac_chunk_size=None, + device_id=0, + rank=None, *, bs_chunk_size=None, **kwargs, @@ -118,6 +120,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -484,6 +488,8 @@ def __init__( eq_fixed=False, name="Boundary error", jac_chunk_size=None, + device_id=0, + rank=None, *, bs_chunk_size=None, B_plasma_chunk_size=None, @@ -529,6 +535,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1055,6 +1063,8 @@ def __init__( deriv_mode="auto", name="NESTOR Boundary", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -1075,6 +1085,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 49706903be..6a6051ac3d 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -271,6 +271,8 @@ def __init__( name="Generic", jac_chunk_size=None, compute_kwargs=None, + device_id=0, + rank=None, **kwargs, ): errorif( @@ -295,6 +297,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) self._print_value_fmt = f"{name} objective value: " self._p = _parse_parameterization(thing) @@ -566,6 +570,8 @@ def __init__( name="Custom", jac_chunk_size=None, compute_kwargs=None, + device_id=0, + rank=None, **kwargs, ): errorif( @@ -592,6 +598,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) self._p = _parse_parameterization(thing) @@ -800,6 +808,8 @@ def __init__( deflation_type="power", multiple_deflation_type="prod", single_shift=False, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -870,6 +880,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 164bf64fb9..81d5f8ec60 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -61,6 +61,8 @@ def __init__( grid=None, name="aspect ratio", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 2 @@ -76,6 +78,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -204,6 +208,8 @@ def __init__( grid=None, name="elongation", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 1 @@ -219,6 +225,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -345,6 +353,8 @@ def __init__( grid=None, name="volume", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 1 @@ -360,6 +370,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -540,6 +552,8 @@ def __init__( name="plasma-vessel distance", use_signed_distance=False, jac_chunk_size=None, + device_id=0, + rank=None, **kwargs, ): if target is None and bounds is None: @@ -580,6 +594,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -866,6 +882,8 @@ def __init__( grid=None, name="mean curvature", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (-np.inf, 0) @@ -881,6 +899,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1005,6 +1025,8 @@ def __init__( grid=None, name="principal-curvature", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 1 @@ -1020,6 +1042,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1139,6 +1163,8 @@ def __init__( grid=None, name="B-scale-length", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (1, np.inf) @@ -1154,6 +1180,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1270,6 +1298,8 @@ def __init__( grid=None, name="coordinate goodness", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -1286,6 +1316,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -1408,6 +1440,8 @@ def __init__( deriv_mode="auto", name="mirror ratio", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0.2 @@ -1423,6 +1457,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_neoclassical.py b/desc/objectives/_neoclassical.py index d2d926b3a5..6b2c072dcb 100644 --- a/desc/objectives/_neoclassical.py +++ b/desc/objectives/_neoclassical.py @@ -88,6 +88,8 @@ def __init__( nufft_eps=1e-6, spline=True, use_bounce1d=False, + device_id=0, + rank=None, **kwargs, ): try: @@ -138,6 +140,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index 4510acbdb5..e50f2b210c 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -64,6 +64,8 @@ def __init__( name="QS Boozer", jac_chunk_size=None, surf_batch_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -83,6 +85,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) self._print_value_fmt = "Quasi-symmetry ({},{}) Boozer error: ".format( @@ -255,6 +259,8 @@ def __init__( helicity=(1, 0), name="QS two-term", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -271,6 +277,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) self._print_value_fmt = "Quasi-symmetry ({},{}) two-term error: ".format( @@ -416,6 +424,8 @@ def __init__( grid=None, name="QS triple product", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -431,6 +441,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -592,6 +604,8 @@ def __init__( name="omnigenity", jac_chunk_size=None, surf_batch_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -625,6 +639,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -915,6 +931,8 @@ def __init__( grid=None, name="Isodynamicity", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -930,6 +948,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_power_balance.py b/desc/objectives/_power_balance.py index 87ceb12999..f954eb8cee 100644 --- a/desc/objectives/_power_balance.py +++ b/desc/objectives/_power_balance.py @@ -58,6 +58,8 @@ def __init__( grid=None, name="fusion power", jac_chunk_size=None, + device_id=0, + rank=None, ): errorif( fuel not in ["DT"], ValueError, f"fuel must be one of ['DT'], got {fuel}." @@ -77,6 +79,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -225,6 +229,8 @@ def __init__( grid=None, name="heating power", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -242,6 +248,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_profiles.py b/desc/objectives/_profiles.py index 51b3b6b255..0bf4da1122 100644 --- a/desc/objectives/_profiles.py +++ b/desc/objectives/_profiles.py @@ -62,6 +62,8 @@ def __init__( grid=None, name="pressure", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -77,6 +79,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -196,6 +200,8 @@ def __init__( grid=None, name="rotational transform", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -211,6 +217,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -343,6 +351,8 @@ def __init__( grid=None, name="shear", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (-np.inf, 0) @@ -358,6 +368,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -482,6 +494,8 @@ def __init__( grid=None, name="toroidal current", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -497,6 +511,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index b0b04d53a8..4496b4c7b1 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -73,6 +73,8 @@ def __init__( grid=None, name="Mercier Stability", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (0, np.inf) @@ -88,6 +90,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -237,6 +241,8 @@ def __init__( grid=None, name="Magnetic Well", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: bounds = (0, np.inf) @@ -252,6 +258,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): @@ -433,6 +441,8 @@ def __init__( w1=10.0, name="ideal ballooning lambda", jac_chunk_size=None, + device_id=0, + rank=None, ): if target is None and bounds is None: target = 0 @@ -466,6 +476,8 @@ def __init__( deriv_mode=deriv_mode, name=name, jac_chunk_size=jac_chunk_size, + device_id=device_id, + rank=rank, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 9ec99d6808..5e67f31c3d 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -8,8 +8,11 @@ from desc.backend import ( desc_config, execute_on_cpu, + jax, jit, jnp, + safe_mpi_Bcast, + safe_transfer_to_device, tree_flatten, tree_leaves, tree_map, @@ -119,6 +122,17 @@ is called on the raw compute value, before any shifting, scaling, or normalization. Operates over all coils, not each individual coil. """ +doc_device_id = """ + device_id : int, optional + Device ID to run the objective on. Defaults to 0. If different objectives + are on different devices, the ObjectiveFunction will run each sub-objective + on the device specified in the sub-objective. +""" +doc_rank = """ + rank : int, optional + MPI rank to run the objective on. Defaults to 0. Objectives on the same + rank should have the same `device_id`. +""" docs = { "target": doc_target, "bounds": doc_bounds, @@ -129,6 +143,8 @@ "deriv_mode": doc_deriv_mode, "name": doc_name, "jac_chunk_size": doc_jac_chunk_size, + "device_id": doc_device_id, + "rank": doc_rank, } doc_bounce = """ @@ -348,6 +364,8 @@ class ObjectiveFunction(IOAble): accurately estimate the available device memory, so the "auto" chunk_size option will yield a larger chunk size than may be needed. It is recommended to manually choose a chunk_size if an OOM error is experienced in this case. + mpi : MPI object, optional + MPI communicator. Required when using multiple devices. """ @@ -358,6 +376,7 @@ class ObjectiveFunction(IOAble): "_objectives", "_use_jit", ] + # these will be updated for MPI later _static_attrs = [ "_built", "_compile_mode", @@ -367,6 +386,7 @@ class ObjectiveFunction(IOAble): "_name", "_things_per_objective_idx", "_use_jit", + "_is_mpi", "_static_attrs", ] @@ -377,6 +397,7 @@ def __init__( deriv_mode="auto", name="ObjectiveFunction", jac_chunk_size="auto", + mpi=None, ): if not isinstance(objectives, (tuple, list)): objectives = (objectives,) @@ -404,6 +425,271 @@ def __init__( self._built = False self._compiled = False self._name = name + device_ids = [obj._device_id for obj in objectives] + self._is_mpi = len(set(device_ids)) > 1 + if mpi is not None: + # for multiple node cases, each process sees 1 CPU + # for those cases we cannot put objectives on different devices + # instead we will run each objective on the given rank + self._is_mpi = True + ranks = [obj._rank for obj in objectives] + # give a reasonable default if all None + if ranks == [None] * len(ranks): + ranks = np.arange(len(objectives)) + errorif( + any(rank is None for rank in ranks), + ValueError, + "If a rank is given to any of the sub-objective, it has to be " + f"given to all of them. Given ranks: {ranks}", + ) + self._rank_per_objective = ranks + self._rank_per_objective = np.asarray(self._rank_per_objective) + # here, we guess the number of devices per node by max(device_ids) + 1 + # device id can be same for different devices on different nodes, these will + # have different ranks, for the check, we take the mod for mapping + errorif( + ( + np.mod(self._rank_per_objective, max(device_ids) + 1) != device_ids + ).any(), + ValueError, + "Some objective's rank and device id are inconsistent. The device id " + "of an objective must equal its rank modulo the number of devices per " + f"node ({max(device_ids) + 1}). Got rank_per_objective=" + f"{self._rank_per_objective} and device_ids={device_ids}.", + ) + self.mpi = mpi + self.comm = self.mpi.COMM_WORLD + self.rank = self.comm.Get_rank() + self.size = self.comm.Get_size() + self.running = True + # rank_per_objective is 0-indexed, so the number of ranks it expects + # is its maximum value + 1. This must match the number of MPI ranks. + n_ranks_needed = max(self._rank_per_objective) + 1 + msg = ( + f"rank_per_objective uses {n_ranks_needed} rank(s) (highest rank " + f"index is {max(self._rank_per_objective)}), but {self.size} MPI " + f"rank(s) are running. These must match. " + ) + errorif( + n_ranks_needed > self.size, + ValueError, + f"{msg}Some objectives are assigned to a rank index that is too large " + "for the number of MPI ranks running. Either run more MPI ranks or " + "lower the rank indices in rank_per_objective.", + ) + self._obj_per_rank = [ + np.where(self._rank_per_objective == i)[0] for i in range(self.size) + ] + # if the constaints should also use MPI, we will store the constraints here + # such that we don't need to pass the constraint objects between workers. + # This can be thought of making them globally accessible to the single + # worker loop. This is set externally (on every rank, before the worker + # loop is started) to the constraints that will be passed to the optimizer, + # and turned into a parallel ObjectiveFunction by _build_constraints. + self._constraints = None + # we will use this string to check if the computations should be done on the + # objective or the constraint when we receive the message in the worker loop + # The possible values are "obj" and "con". It is set to "con" in + # _build_constraints for the constraint ObjectiveFunction. + self._obj_type = "obj" + self._static_attrs += [ + "mpi", + "comm", + "rank", + "size", + "running", + "_obj_per_rank", + "_rank_per_objective", + "_f_sizes", + "_f_displs", + "_constraints", + "_obj_type", + ] + + if self._is_mpi and mpi is None: + raise ValueError( + "MPI communicator must be passed when objectives are on different " + "devices." + ) + + def __enter__(self): + errorif( + not self._built, + RuntimeError, + "In parallel mode, ObjectiveFunction must be built before entering " + "context manager.", + ) + errorif( + not self._is_mpi, + RuntimeError, + "ObjectiveFunction must be parallel to be used as a context manager.", + ) + # when entering the context manager, we start the worker loop + # this allows the root rank to send messages to the workers to compute and stop + self._worker_loop() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + # this will be called when the context manager exits + # we send a stop message to the workers + if self.rank == 0: + # only the root rank can send the stop message + # in general the message contains 3 parts + # but for the stop message we only need the first part + message = ("STOP", None, None) + self.comm.bcast(message, root=0) + self.running = False + + def _build_constraints(self, verbose=0): + """Combine and build the nonlinear constraints of a parallel optimization. + + ``self._constraints`` is set to the nonlinear constraints of the problem by + ``desc.optimize.build_for_mpi``, which is called by every rank before + the context manager is entered. Here they are combined into their own parallel + ObjectiveFunction, which uses the worker loop of this ObjectiveFunction instead + of having one of its own. Since every rank does this, the constraint objects + never have to be communicated between the ranks. + + The constraints only use MPI if the user distributed them over ranks or + devices. Otherwise this is reset to None and they are computed on the root + rank, like they would be without MPI. + + Parameters + ---------- + verbose : int, optional + Level of output. + + """ + if isinstance(self._constraints, ObjectiveFunction): + return # already combined and built + cons = self._constraints + cons = [cons] if isinstance(cons, _Objective) else list(cons) + if not ( + any(con._rank is not None for con in cons) + or len({con._device_id for con in cons}) > 1 + ): + # the ranks would default to np.arange(len(cons)) which has nothing to do + # with what the user wants, so don't parallelize the constraints at all + self._constraints = None + return + self._constraints = ObjectiveFunction(cons, deriv_mode="blocked", mpi=self.mpi) + self._constraints._obj_type = "con" + self._constraints.build(verbose=verbose) + # the constraints must take the same state vector as the objective + self._constraints._set_things(self.things) + + def _worker_loop(self): + """Worker loop for MPI parallelization. + + This function is called when the ObjectiveFunction is used as a context manager. + + with obj: + if rank == 0: + eq.optimize(objective=obj) + + Worker processes will be in a loop waiting for messages from the root rank + during the context manager. The root rank will send messages to the workers + to compute the objective function and its derivatives, and to stop. The workers + will then broadcast the results back to the root rank. Once the context manager + exits, the loop will be terminated by the root rank. + + The same loop serves both the objective and the nonlinear constraints, the + first part of the message tells which one an operation belongs to. + + Therefore we can use MPI parallelization with the ObjectiveFunction while + preventing execution of redundant calculations on different ranks. + This is very similar to the strategy used in Simsopt. + + """ + if self.rank == 0: + # Root rank won't enter worker loop + return + + def alloc_array(shape, device=None): + if not desc_config["mpi-cuda"]: + return np.empty(shape, dtype=np.float64) + return jnp.empty(shape, dtype=jnp.float64, device=device) + + while self.running: + # The message contains 3 parts: + # message[0] is ":", the first part + # tells if the operation belongs to the objective or to the constraints + # message[1] is the size of state vector (for compute and jvp's) + # message[2] is the shape of tangents (for only jvp's) + message = (None, None, None) + message = self.comm.bcast(message, root=0) + + if message[0] == "STOP": + print(f"Rank {self.rank} STOPPING") + break + + kind, op = message[0].split(":") + objfun = self if kind == "obj" else self._constraints + obj_idx_rank = objfun._obj_per_rank[self.rank] + objs = [objfun.objectives[i] for i in obj_idx_rank] + # a rank may have no constraint assigned to it, it still has to take part + # in the collective calls, but with empty buffers + device = objs[0]._device if len(objs) else None + + # get arrays by Bcast which uses buffers and faster than bcast + x = alloc_array(message[1], device=device) + x = safe_mpi_Bcast(x, self.comm, root=0) + + if "compute" in op: + if len(objs): + params = objfun.unpack_state(x) + params = jax.device_put(params, device) + params = [params[i] for i in obj_idx_rank] + out = compute_per_process(params, objs, op=op) + else: + out = jnp.empty(0) + if not desc_config["mpi-cuda"]: + out = np.array(out) + self.comm.Gatherv( + out, + (None, objfun._f_sizes, objfun._f_displs, self.mpi.DOUBLE), + root=0, + ) + elif "jvp" in op: + splits = np.cumsum([t.dim_x for t in objfun.things]) + x = jnp.split(x, splits) + vs = alloc_array(message[2], device=device) + vs = safe_mpi_Bcast(vs, self.comm, root=0) + vs = jnp.split(vs, splits, axis=-1) + + # put xi and vi on the same device as the objective + xs = jax.device_put(x, device) + vs = jax.device_put(vs, device) + # only pass the relevant parts of x and v to each objective + xs = [ + [xs[i] for i in objfun._things_per_objective_idx[idx]] + for idx in obj_idx_rank + ] + vs = [ + [vs[i] for i in objfun._things_per_objective_idx[idx]] + for idx in obj_idx_rank + ] + if not len(objs): + out = jnp.empty((0, message[2][0])) + elif "proximal" not in op: + out = jvp_per_process(xs, vs, objs, op=op).T + elif "proximal_jvp" in op: + out = jvp_proximal_per_process( + xs, vs, objs, op=op.replace("proximal_jvp_", "") + ) + + if not desc_config["mpi-cuda"]: + out = np.array(out) + self.comm.Gatherv( + out, + ( + None, + objfun._f_sizes * message[2][0], + objfun._f_displs * message[2][0], + self.mpi.DOUBLE, + ), + root=0, + ) def _unjit(self): """Remove jit compiled methods.""" @@ -435,7 +721,7 @@ def _unjit(self): pass @execute_on_cpu - def build(self, use_jit=None, verbose=1): + def build(self, use_jit=None, verbose=1): # noqa: C901 """Build the objective. Parameters @@ -461,17 +747,35 @@ def build(self, use_jit=None, verbose=1): ) self._use_jit = False + # cannot use different devices under jit, and MPI calls cannot be traced, + # unjit to allow for both. The actual computations are still jitted by the + # module level *_per_process functions. + if self._is_mpi: + self._use_jit = False + timer = Timer() timer.start("Objective build") # build objectives self._dim_f = 0 for objective in self.objectives: + obj_things = objective._things if not objective.built: if verbose > 0: print("Building objective: " + objective.name) objective.build(use_jit=self.use_jit, verbose=verbose) self._dim_f += objective.dim_f + if objective._device_id != 0: + if verbose > 0 and self.rank == 0: + print( + f"Putting objective {objective.name} on device " + f"{objective._device_id}" + ) + objective = jax.device_put(objective, objective._device) + # same object on different device will have different id + # need to overwrite by original to keep it the same and make + # _set_things work + objective._things = obj_things if self._dim_f == 1: self._scalar = True else: @@ -511,6 +815,16 @@ def build(self, use_jit=None, verbose=1): else: self._deriv_mode = "blocked" + warnif( + self._is_mpi and self._deriv_mode != "blocked", + UserWarning, + "When using multiple devices, the ObjectiveFunction will run each " + "sub-objective on the device specified in the sub-objective. " + "Setting the deriv_mode to 'blocked' to ensure that each sub-objective " + "runs on the correct device.", + ) + if self._is_mpi: + self._deriv_mode = "blocked" errorif( isposint(self._jac_chunk_size) and self._deriv_mode in ["blocked"], ValueError, @@ -522,11 +836,12 @@ def build(self, use_jit=None, verbose=1): # Heuristic estimates of fwd mode Jacobian memory usage, # slightly conservative, based on using ForceBalance as the objective estimated_memory_usage = 2.4e-7 * self.dim_f * self.dim_x + 1 # in GB - avail_mem = desc_config.get("avail_mem") + avail_mem = desc_config["avail_mems"][0] # in GB max_chunk_size = round( (avail_mem / estimated_memory_usage - 0.22) / 0.85 * self.dim_x ) self._jac_chunk_size = max([1, max_chunk_size]) + if self._deriv_mode == "blocked": chunk_sizes = [obj._jac_chunk_size for obj in self.objectives] if len(set(chunk_sizes)) > 1: @@ -539,6 +854,38 @@ def build(self, use_jit=None, verbose=1): # use the chunk size of the first objective self._jac_chunk_size = self.objectives[0]._jac_chunk_size + if self._is_mpi: + # sizes and displacements for Gatherv + self._f_sizes = np.array( + [ + sum([self.objectives[i].dim_f for i in ids]) + for ids in self._obj_per_rank + ] + ) + self._f_displs = np.array( + [sum(self._f_sizes[:i]) for i in range(self.size)] + ) + + if self._is_mpi and verbose > 0: + if self.rank == 0: + objective_names_per_rank = [ + [self._objectives[i].__class__.__name__ for i in objective_ids] + for objective_ids in self._obj_per_rank + ] + kind = "objective" if self._obj_type == "obj" else "constraint" + print("-" * 60) + for rank in range(self.size): + print( + f"Rank {rank} will run {kind}(s): " + f"{objective_names_per_rank[rank]}" + ) + print("-" * 60) + + # nonlinear constraints can run in parallel as well, they use the worker loop + # of this objective, see _build_constraints + if self._is_mpi and self._constraints is not None: + self._build_constraints(verbose=verbose) + if not self._use_jit: self._unjit() @@ -592,16 +939,49 @@ def _set_things(self, things=None): def _compute_op(self, x, constants=None, op="compute_unscaled"): """Helper function to compute various operations.""" constants = self._get_deprecated_constants(constants) - params = self.unpack_state(x) - assert len(params) == len(constants) == len(self.objectives) - f = jnp.concatenate( - [ - getattr(obj, op)(*par, constants=const) - for par, obj, const in zip(params, self.objectives, constants) - ] - ) + if not self._is_mpi: + params = self.unpack_state(x) + assert len(params) == len(constants) == len(self.objectives) + f = jnp.concatenate( + [ + getattr(obj, op)(*par, constants=const) + for par, obj, const in zip(params, self.objectives, constants) + ] + ) + else: + f = self._parallel_compute(x, op) return f + def _parallel_compute(self, x, op): + """Compute the objective function in parallel using MPI.""" + if self.rank == 0: + message = (self._obj_type + ":" + op, x.shape, None) + self.comm.bcast(message, root=0) + safe_mpi_Bcast(x, self.comm, root=0) + + obj_idx_rank = self._obj_per_rank[self.rank] + if len(obj_idx_rank): + params = self.unpack_state(x) + f_rank = compute_per_process( + [params[i] for i in obj_idx_rank], + [self.objectives[i] for i in obj_idx_rank], + op=op, + ) + else: + f_rank = jnp.empty(0) + if not desc_config["mpi-cuda"]: + f_rank = np.array(f_rank) + recvbuf = np.empty(self.dim_f, dtype=np.float64) + else: + recvbuf = jnp.empty(self.dim_f, dtype=jnp.float64) + self.comm.Gatherv( + f_rank, + (recvbuf, self._f_sizes, self._f_displs, self.mpi.DOUBLE), + root=0, + ) + recvbuf = safe_transfer_to_device(recvbuf) + return recvbuf + @jit def compute_unscaled(self, x, constants=None): """Compute the raw value of the objective function. @@ -757,6 +1137,9 @@ def print_value(self, x, x0=None, constants=None, fse=None, f0se=None): else: f0i = None offset += dim + if self._is_mpi: + par = jax.device_put(par, obj._device) + par0 = jax.device_put(par0, obj._device) outi = obj.print_value( args=par, args0=par0, fse=fi, f0se=f0i, constants=const ) @@ -836,17 +1219,27 @@ def x(self, *things): def grad(self, x, constants=None): """Compute gradient vector of self.compute_scalar wrt x.""" constants = self._get_deprecated_constants(constants) - return jnp.atleast_1d( - Derivative(self.compute_scalar, mode="grad")(x, constants).squeeze() - ) + if not self._is_mpi: + return jnp.atleast_1d( + Derivative(self.compute_scalar, mode="grad")(x, constants).squeeze() + ) + else: + raise NotImplementedError( + "Gradient computation is not implemented for MPI ObjectiveFunction." + ) @jit def hess(self, x, constants=None): """Compute Hessian matrix of self.compute_scalar wrt x.""" constants = self._get_deprecated_constants(constants) - return jnp.atleast_2d( - Derivative(self.compute_scalar, mode="hess")(x, constants).squeeze() - ) + if not self._is_mpi: + return jnp.atleast_2d( + Derivative(self.compute_scalar, mode="hess")(x, constants).squeeze() + ) + else: + raise NotImplementedError( + "Gradient computation is not implemented for MPI ObjectiveFunction." + ) @jit def jac_scaled(self, x, constants=None): @@ -874,25 +1267,75 @@ def _jvp_blocked(self, v, x, constants=None, op="scaled"): # is needed for perturbations. Just pass that to jvp_batched for now return self._jvp_batched(v, x, constants, op) - xs_splits = np.cumsum([t.dim_x for t in self.things]) - xs = jnp.split(x, xs_splits) - vs = jnp.split(v[0], xs_splits, axis=-1) - J = [] - assert len(self.objectives) == len(constants) - # basic idea is we compute the jacobian of each objective wrt each thing - # one by one, and assemble into big block matrix - # if objective doesn't depend on a given thing, that part is set to 0. - for k, (obj, const) in enumerate(zip(self.objectives, constants)): - # get the xs that go to that objective - thing_idx = self._things_per_objective_idx[k] - xi = [xs[i] for i in thing_idx] - vi = [vs[i] for i in thing_idx] - Ji_ = getattr(obj, "jvp_" + op)(vi, xi, constants=const) - J += [Ji_] - # this is the transpose of the jvp when v is a matrix, for consistency with - # jvp_batched - J = jnp.hstack(J) - return J + if not self._is_mpi: + xs_splits = np.cumsum([t.dim_x for t in self.things]) + xs = jnp.split(x, xs_splits) + vs = jnp.split(v[0], xs_splits, axis=-1) + J = [] + assert len(self.objectives) == len(constants) + # basic idea is we compute the jacobian of each objective wrt each thing + # one by one, and assemble into big block matrix + # if objective doesn't depend on a given thing, that part is set to 0. + for k, (obj, const) in enumerate(zip(self.objectives, constants)): + # get the xs that go to that objective + thing_idx = self._things_per_objective_idx[k] + xi = [xs[i] for i in thing_idx] + vi = [vs[i] for i in thing_idx] + Ji_ = getattr(obj, "jvp_" + op)(vi, xi, constants=const) + J += [Ji_] + # this is the transpose of the jvp when v is a matrix, for consistency with + # jvp_batched + return jnp.hstack(J) + else: + if self.rank == 0: + # broadcasting x and v as single array is faster than + # boradcasting the list + message = (self._obj_type + ":jvp_" + op, x.shape, v[0].shape) + self.comm.bcast(message, root=0) + safe_mpi_Bcast(x, self.comm, root=0) + safe_mpi_Bcast(v[0], self.comm, root=0) + + xs = jnp.split(x, np.cumsum([t.dim_x for t in self.things])) + vs = jnp.split(v[0], np.cumsum([t.dim_x for t in self.things]), axis=-1) + + obj_idx_rank = self._obj_per_rank[self.rank] + # jvp_per_process returns the Jacobian in a transposed way which is + # hard to stack vertically by MPI (colums get scrambled), that's why + # we will do multiple transpose operations. The first one is to be able + # to stack the Jacobian parts vertically, the second one is to return + # the Jacobian in the expected way by other functions. + if len(obj_idx_rank): + J_rank = jvp_per_process( + [ + [xs[i] for i in self._things_per_objective_idx[idx]] + for idx in obj_idx_rank + ], + [ + [vs[i] for i in self._things_per_objective_idx[idx]] + for idx in obj_idx_rank + ], + [self.objectives[i] for i in obj_idx_rank], + op="jvp_" + op, + ).T + else: + J_rank = jnp.empty((0, message[2][0])) + if not desc_config["mpi-cuda"]: + J_rank = np.array(J_rank) + recvbuf = np.empty((self.dim_f, message[2][0]), dtype=np.float64) + else: + recvbuf = jnp.empty((self.dim_f, message[2][0]), dtype=jnp.float64) + self.comm.Gatherv( + J_rank, + ( + recvbuf, + self._f_sizes * message[2][0], + self._f_displs * message[2][0], + self.mpi.DOUBLE, + ), + root=0, + ) + recvbuf = safe_transfer_to_device(recvbuf) + return recvbuf.T def _jvp_batched(self, v, x, constants=None, op="scaled"): v = ensure_tuple(v) @@ -984,6 +1427,11 @@ def jvp_unscaled(self, v, x, constants=None): return J def _vjp(self, v, x, constants=None, op="scaled"): + errorif( + self._is_mpi, + NotImplementedError, + "Vector-Jacobian product is not implemented for MPI ObjectiveFunction.", + ) fun = lambda x: getattr(self, "compute_" + op)(x, constants) return Derivative.compute_vjp(fun, 0, v, x) @@ -1264,6 +1712,9 @@ class _Objective(IOAble, ABC): "_print_value_fmt", "_scalar", "_units", + "_device", + "_device_id", + "_rank", "_static_attrs", ] @@ -1279,6 +1730,8 @@ def __init__( deriv_mode="auto", name=None, jac_chunk_size=None, + device_id=0, + rank=None, ): if self._scalar: assert self._coordinates == "" @@ -1292,6 +1745,21 @@ def __init__( assert jac_chunk_size is None or isposint(jac_chunk_size) self._jac_chunk_size = jac_chunk_size + self._device_id = device_id + self._rank = rank + # This will help the data placement if we have multiple GPU devices. + # Linear objectives have a separate ObjectiveFunction, hence they cannot use the + # "with" context manager of the main objective. + # They should run on the default device to avoid staling code. + if ( + desc_config["num_device"] != 1 + and desc_config["kind"] == "gpu" + and not self._linear + ): + self._device = jax.devices("gpu")[device_id] + else: + # multiple CPUs should already have the data on their rank + self._device = None self._target = target self._bounds = bounds @@ -2002,3 +2470,48 @@ def __call__(self, things): assert len(flat) == self.length unique, _, _ = unique_list(flat) return unique + + +# These will run on workers, and we want to safely jit them +@functools.partial(jit, static_argnames="op") +def compute_per_process(params, objectives, op): + """Compute the objective function on each process.""" + return jnp.concatenate( + [ + getattr(obj, op)(*param, constants=None) + for (obj, param) in zip(objectives, params) + ] + ) + + +@functools.partial(jit, static_argnames="op") +def jvp_per_process(x, v, objectives, op): + """Compute the Jacobian-vector product on each process.""" + return jnp.hstack( + [ + getattr(obj, op)(v[idx], x[idx], constants=None) + for idx, obj in enumerate(objectives) + ] + ) + + +@functools.partial(jit, static_argnames="op") +def jvp_proximal_per_process(x, v, objectives, op): + """Compute the Jacobian-vector product on each process, for proximal.""" + J_rank = [] + for idx, obj in enumerate(objectives): + if obj._deriv_mode == "rev": + # obj might not allow fwd mode, so compute full rev mode Jacobian and do + # matmul manually. This is slightly inefficient, but usually with rev mode + # dim_f << dim_x so it is not too bad. + Ji = getattr(obj, "jac_" + op)(*x[idx]) + J_rank.append( + jnp.array([Jii @ vii.T for Jii, vii in zip(Ji, v[idx])]).sum(axis=0) + ) + else: + J_rank.append( + getattr(obj, "jvp_" + op)( + [_vi for _vi in v[idx]], x[idx], constants=None + ).T + ) + return jnp.vstack(J_rank) diff --git a/desc/optimize/__init__.py b/desc/optimize/__init__.py index 7093f03c4a..41875bb341 100644 --- a/desc/optimize/__init__.py +++ b/desc/optimize/__init__.py @@ -6,5 +6,11 @@ from .aug_lagrangian_ls import lsq_auglag from .fmin_scalar import fmintr from .least_squares import lsqtr -from .optimizer import Optimizer, optimizers, register_optimizer +from .optimizer import ( + Optimizer, + build_for_mpi, + optimizers, + register_optimizer, + run_with_mpi, +) from .stochastic import sgd diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py index 1041b2a1fb..1f755dd183 100644 --- a/desc/optimize/_constraint_wrappers.py +++ b/desc/optimize/_constraint_wrappers.py @@ -4,7 +4,14 @@ import numpy as np -from desc.backend import jit, jnp, put +from desc.backend import ( + desc_config, + jit, + jnp, + put, + safe_mpi_Bcast, + safe_transfer_to_device, +) from desc.batching import batched_vectorize from desc.objectives import ( BoundaryRSelfConsistency, @@ -13,6 +20,7 @@ get_fixed_boundary_constraints, maybe_add_self_consistency, ) +from desc.objectives.objective_funs import jvp_proximal_per_process from desc.objectives.utils import ( _Project, _Recover, @@ -1067,15 +1075,27 @@ def grad(self, x, constants=None): v = jnp.eye(x.shape[0]) constants = setdefault(constants, [None, None]) xg, xf = self._update_equilibrium(x, store=True) - jvpfun = lambda u: self._get_tangent(u, xf, constants, op="scaled_error") - tangents = batched_vectorize( - jvpfun, - signature="(n)->(k)", - chunk_size=self._constraint._jac_chunk_size, - )(v) - g = self._objective.compute_scaled_error(xg, constants[0]) - g_vjp = self._objective.vjp_scaled_error(g, xg, constants[0]) - return tangents @ g_vjp + if not (self._constraint._is_mpi or self._objective._is_mpi): + jvpfun = lambda u: self._get_tangent(u, xf, constants, op="scaled_error") + tangents = batched_vectorize( + jvpfun, + signature="(n)->(k)", + chunk_size=self._constraint._jac_chunk_size, + )(v) + g = self._objective.compute_scaled_error(xg, constants[0]) + g_vjp = self._objective.vjp_scaled_error(g, xg, constants[0]) + return tangents @ g_vjp + elif self._constraint._is_mpi: + # TODO: implement parallel constraint for ProximalProjection + raise NotImplementedError( + "Parallel constraint for ProximalProjection not implemented yet. " + "Please use only one Equilibrium constraint." + ) + else: + # TODO: apply vjp for multidevice similar to #2030 + f = jnp.atleast_1d(self.compute_scaled_error(x, constants)) + J = self.jac_scaled_error(x, constants) + return f.T @ J def hess(self, x, constants=None): """Compute Hessian of self.compute_scalar. @@ -1228,23 +1248,36 @@ def _jvp(self, v, x, constants=None, op="scaled_error"): # we don't need to divide this part into blocked and batched because # self._constraint._deriv_mode will handle it - jvpfun = lambda u: self._get_tangent(u, xf, constants, op=op) - tangents = batched_vectorize( - jvpfun, - signature="(n)->(k)", - chunk_size=self._constraint._jac_chunk_size, - )(v) + if not self._constraint._is_mpi: + jvpfun = lambda u: self._get_tangent(u, xf, constants, op=op) + tangents = batched_vectorize( + jvpfun, + signature="(n)->(k)", + chunk_size=self._constraint._jac_chunk_size, + )(v) + else: + # TODO: implement parallel constraint for ProximalProjection + # Note: the workers no longer need a second loop for this, the constraints + # share the loop of the objective. What is left is that _get_tangent is + # vectorized over v by batched_vectorize, and MPI calls cannot be traced, + # so the tangents have to be computed for all directions at once instead. + raise NotImplementedError( + "Parallel constraint for ProximalProjection not implemented yet. " + "Please use only one Equilibrium constraint." + ) if self._objective._deriv_mode == "batched": # objective's method already know about its jac_chunk_size return getattr(self._objective, "jvp_" + op)(tangents, xg, constants[0]) else: - return _proximal_jvp_blocked_pure( - self._objective, - jnp.split(tangents, np.cumsum(self._dimx_per_thing), axis=-1), - jnp.split(xg, np.cumsum(self._dimx_per_thing)), - op, - ) + if not self._objective._is_mpi: + vgs = jnp.split(tangents, np.cumsum(self._dimx_per_thing), axis=-1) + xgs = jnp.split(xg, np.cumsum(self._dimx_per_thing)) + return _proximal_jvp_blocked_pure(self._objective, vgs, xgs, op) + else: + return _proximal_jvp_blocked_parallel( + self._objective, tangents, xg, np.cumsum(self._dimx_per_thing), op + ) def _get_tangent(self, v, xf, constants, op): # Note: This function is vectorized over v. So, v is expected to be 1D array @@ -1388,4 +1421,50 @@ def _proximal_jvp_blocked_pure(objective, vgs, xgs, op): else: outi = getattr(obj, "jvp_" + op)([_vi for _vi in vi], xi).T out.append(outi) + return jnp.concatenate(out).T + + +def _proximal_jvp_blocked_parallel(objective, vgs, xgs, splits, op): + if objective.rank == 0: + message = (objective._obj_type + ":proximal_jvp_" + op, xgs.shape, vgs.shape) + objective.comm.bcast(message, root=0) + safe_mpi_Bcast(xgs, comm=objective.comm, root=0) + safe_mpi_Bcast(vgs, comm=objective.comm, root=0) + + xgs = jnp.split(xgs, splits) + vgs = jnp.split(vgs, splits, axis=-1) + + obj_idx_rank = objective._obj_per_rank[objective.rank] + xs = [ + [xgs[i] for i in objective._things_per_objective_idx[idx]] + for idx in obj_idx_rank + ] + vs = [ + [vgs[i] for i in objective._things_per_objective_idx[idx]] + for idx in obj_idx_rank + ] + objs = [objective.objectives[i] for i in obj_idx_rank] + J_rank = jvp_proximal_per_process(xs, vs, objs, op=op) + if not desc_config["mpi-cuda"]: + J_rank = np.array(J_rank) + recvbuf = np.empty((objective.dim_f, J_rank.shape[1]), dtype=np.float64) + else: + recvbuf = jnp.empty((objective.dim_f, J_rank.shape[1]), dtype=jnp.float64) + objective.comm.Gatherv( + J_rank, + ( + recvbuf, + objective._f_sizes * J_rank.shape[1], + objective._f_displs * J_rank.shape[1], + objective.mpi.DOUBLE, + ), + root=0, + ) + recvbuf = safe_transfer_to_device(recvbuf) + + # we collected the Jacobian in the proper way above, but as a convention + # the _jvp methods return the transpose of the Jacobian. For example, + # _jac methods always take the transpose of the returned quantity by _jvp. + # To be consistent with that, we return the transpose here. + return recvbuf.T diff --git a/desc/optimize/optimizer.py b/desc/optimize/optimizer.py index 2d4b114a55..56af0a198b 100644 --- a/desc/optimize/optimizer.py +++ b/desc/optimize/optimizer.py @@ -576,8 +576,123 @@ def _parse_constraints(constraints): return linear_constraints, nonlinear_constraints +def build_for_mpi(objective, constraints=(), verbose=1): + """Build a parallel objective and its nonlinear constraints on every rank. + + When using MPI, only the root rank runs the optimization, the other ranks wait in + the worker loop of the ObjectiveFunction. Everything the workers need must therefore + be known to them before the context manager is entered, which is what this function + does: it parses the constraints the same way the optimizer does and hands the + nonlinear ones to the ObjectiveFunction, so that they can use the same worker loop, + then builds everything. In a parallel script it replaces ``objective.build()``, + + objective = build_for_mpi(objective, constraints) + with objective: + if rank == 0: + eq.optimize(objective=objective, constraints=constraints, ...) + + Parameters + ---------- + objective : ObjectiveFunction + Objective function to optimize, created with an ``mpi`` communicator. + constraints : tuple of Objective + The same constraints that will be passed to the optimizer. The nonlinear ones + are computed in parallel only if they are given a ``rank`` or a ``device_id``, + otherwise they are computed on the root rank as usual. + verbose : int, optional + Level of output. Only the root rank prints. + + Returns + ------- + objective : ObjectiveFunction + Built objective function, that also knows about the nonlinear constraints. + + """ + errorif( + not isinstance(objective, ObjectiveFunction), + TypeError, + "objective should be of type ObjectiveFunction.", + ) + if not isinstance(constraints, (tuple, list)): + constraints = (constraints,) + _, nonlinear_constraints = _parse_constraints(constraints) + # wrappers like LinearConstraintProjection don't have this attribute + is_mpi = getattr(objective, "_is_mpi", False) + errorif( + not is_mpi and any(con._rank is not None for con in nonlinear_constraints), + ValueError, + "Some nonlinear constraints are assigned to a rank, but the ObjectiveFunction " + "is not parallel. The constraints use the worker loop of the ObjectiveFunction," + " so the MPI communicator must be given to it, even if the objectives " + "themselves all run on a single device, ie ObjectiveFunction(..., mpi=MPI).", + ) + verbose = verbose if (not is_mpi or objective.rank == 0) else 0 + + if is_mpi: + # ObjectiveFunction.build combines and builds these, see _build_constraints + objective._constraints = ( + nonlinear_constraints if nonlinear_constraints else None + ) + if not objective.built: + objective.build(verbose=verbose) + elif is_mpi and objective._constraints is not None: + objective._build_constraints(verbose=verbose) + + # make the objective and the constraints take the same state vector, this is what + # combine_args does on the root rank once the optimization starts + things = unique_list( + flatten_list([objective.things] + [con.things for con in constraints]) + )[0] + objective._set_things(things) + if is_mpi and objective._constraints is not None: + objective._constraints._set_things(things) + return objective + + +@contextlib.contextmanager +def run_with_mpi(objective, constraints=(), verbose=1): + """Build a problem for MPI and keep the worker ranks listening to the root rank. + + Combines ``build_for_mpi`` with the context manager of the ObjectiveFunction, so + that a parallel script is + + with run_with_mpi(objective, constraints) as is_root: + if is_root: + eq.optimize(objective=objective, constraints=constraints, ...) + + Every rank builds the objective and the nonlinear constraints, then the worker + ranks wait for the root rank to send them work, until the context is exited. Only + the root rank gets ``is_root=True``, since the optimization itself is done by it + alone. If the objective is not parallel, this only builds it and every rank is the + root rank, so that the same script can be run with and without MPI. + + Parameters + ---------- + objective : ObjectiveFunction + Objective function to optimize, created with an ``mpi`` communicator. + constraints : tuple of Objective + The same constraints that will be passed to the optimizer. The nonlinear ones + are computed in parallel only if they are given a ``rank`` or a ``device_id``, + otherwise they are computed on the root rank as usual. + verbose : int, optional + Level of output. Only the root rank prints. + + Yields + ------ + is_root : bool + Whether this rank is the one that should run the optimization. + + """ + objective = build_for_mpi(objective, constraints, verbose) + if not objective._is_mpi: + yield True + else: + with objective: + yield objective.rank == 0 + + def _maybe_wrap_nonlinear_constraints( - eq, objective, nonlinear_constraints, method, options + eq, objective, nonlinear_constraints, method, options, nonlinear_constraint=None ): """Use ProximalProjection to handle nonlinear constraints.""" if eq is None: # not deal with an equilibrium problem -> no ProximalProjection @@ -603,7 +718,11 @@ def _maybe_wrap_nonlinear_constraints( solve_options = options.pop("solve_options", {}) objective = ProximalProjection( objective, - constraint=_combine_constraints(nonlinear_constraints), + constraint=( + nonlinear_constraint + if nonlinear_constraint is not None + else _combine_constraints(nonlinear_constraints) + ), perturb_options=perturb_options, solve_options=solve_options, eq=eq, @@ -631,8 +750,23 @@ def get_combined_constraint_objectives( # noqa: C901 # parse and combine constraints into linear & nonlinear objective functions linear_constraints, nonlinear_constraints = _parse_constraints(constraints) + # for a parallel objective, the nonlinear constraints are already combined and + # built by every rank in build_for_mpi, reuse that ObjectiveFunction so + # that the root rank and the workers use the same one + mpi_constraint = getattr(objective, "_constraints", None) + mpi_constraint = ( + mpi_constraint if isinstance(mpi_constraint, ObjectiveFunction) else None + ) + errorif( + mpi_constraint is not None + and {id(con) for con in mpi_constraint.objectives} + != {id(con) for con in nonlinear_constraints}, + ValueError, + "The nonlinear constraints given to the optimizer are not the same as the " + "ones given to build_for_mpi.", + ) objective, nonlinear_constraints = _maybe_wrap_nonlinear_constraints( - eq, objective, nonlinear_constraints, opt_method, options + eq, objective, nonlinear_constraints, opt_method, options, mpi_constraint ) is_prox = isinstance(objective, ProximalProjection) for t in things: @@ -642,7 +776,11 @@ def get_combined_constraint_objectives( # noqa: C901 continue linear_constraints = maybe_add_self_consistency(t, linear_constraints) linear_constraint = _combine_constraints(linear_constraints) - nonlinear_constraint = _combine_constraints(nonlinear_constraints) + nonlinear_constraint = ( + mpi_constraint + if (mpi_constraint is not None and len(nonlinear_constraints)) + else _combine_constraints(nonlinear_constraints) + ) # make sure everything is built if objective is not None and not objective.built: diff --git a/devtools/check_unmarked_tests.py b/devtools/check_unmarked_tests.py index 150a67d2cb..019e6618fe 100644 --- a/devtools/check_unmarked_tests.py +++ b/devtools/check_unmarked_tests.py @@ -8,7 +8,7 @@ import ast import sys -REQUIRED_MARKS = {"unit", "regression", "benchmark", "memory"} +REQUIRED_MARKS = {"unit", "regression", "benchmark", "memory", "mpi_run", "mpi_setup"} def _pytest_marks(decorators): diff --git a/docs/index.rst b/docs/index.rst index 2096b112f6..e6aee7dc9e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -24,6 +24,7 @@ installation notebooks/tutorials/use_outputs.ipynb performance_tips + notebooks/tutorials/multi_device.ipynb .. toctree:: diff --git a/docs/notebooks/tutorials/mpi-tutorials/mpi-eq-solve.py b/docs/notebooks/tutorials/mpi-tutorials/mpi-eq-solve.py new file mode 100644 index 0000000000..4130b01873 --- /dev/null +++ b/docs/notebooks/tutorials/mpi-tutorials/mpi-eq-solve.py @@ -0,0 +1,118 @@ +import os +import sys + +# Add the path to the parent directory to augment search for module +sys.path.insert(0, os.path.abspath(".")) +sys.path.append(os.path.abspath("../../../")) +sys.path.append(os.path.abspath("../../../../")) + +import numpy as np +from mpi4py import MPI + +from desc import _set_cpu_count, set_device + +kind = "cpu" # or "gpu" +num_device = 2 +# ====== Using CPUs ====== +# These will be used for dividing the single CPU into multiple virtual CPUs +# such that JAX and XLA thinks there are multiple devices +if kind == "cpu": + # !!! If you have multiple CPUs, you shouldn't call `_set_cpu_count` !!! + _set_cpu_count(num_device) + set_device("cpu", num_device=num_device, mpi=MPI) + +# ====== Using GPUs ====== +# When we have multiple processes using the same devices (for example, 3 processes +# using 3 GPUs), each process will try to pre-allocate 75% of the GPU memory which will +# cause the memory allocation to fail. To avoid this, we can set the allocator to `platform` +# such that there is no pre-allocation. This is a bit conservative (and probably there is room +# for improvement), but if a process needs more memory, it can use more memory on the fly. +elif kind == "gpu": + os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" + set_device("gpu", num_device=num_device) + +from desc import config as desc_config +from desc.backend import jax, print_backend_info +from desc.examples import get +from desc.grid import LinearGrid +from desc.objectives import ForceBalance, ObjectiveFunction +from desc.objectives.getters import get_fixed_boundary_constraints +from desc.optimize import run_with_mpi + +if __name__ == "__main__": + rank = MPI.COMM_WORLD.Get_rank() + size = MPI.COMM_WORLD.Get_size() + if rank == 0: + print(f"====== TOTAL OF {size} RANKS ======") + + # see which rank is running on which device + # Note: JAX has 2 functions for this: `jax.devices()` and `jax.local_devices()` + # `jax.devices()` will return all devices available to JAX, while `jax.local_devices()` + # will return only the devices that are available to the current process. This is + # useful when you have multiple processes running on multiple nodes and you want + # to see which devices are available to each process. + if desc_config["kind"] == "gpu": + print( + f"Rank {rank} can see {jax.local_devices(backend='gpu')} " + f"and {jax.local_devices(backend='cpu')}\n" + ) + else: + print(f"Rank {rank} can see {jax.local_devices(backend='cpu')}\n") + + if rank == 0: + print("====== BACKEND INFO ======") + print_backend_info() + print("\n") + + eq = get("HELIOTRON") + if desc_config["kind"] == "cpu": + # for local testing use lower resolution + eq.change_resolution(M=3, N=2, M_grid=6, N_grid=4) + + # setup 2 grids for 2 objectives covering different flux surfaces + rhos = np.linspace(0.1, 1.0, eq.L_grid) + grid1 = LinearGrid( + rho=rhos[: rhos.size // 2], + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + ) + grid2 = LinearGrid( + rho=rhos[rhos.size // 2 :], + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + ) + # ranks will be automatically assigned as 0 and 1, respectively + obj = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1, device_id=0), + ForceBalance(eq, grid=grid2, device_id=1), + ], + mpi=MPI, + deriv_mode="blocked", + ) + cons = get_fixed_boundary_constraints(eq) + + # Until this line, the code is performed on all ranks, so it might print some + # information multiple times. The following part will only be performed on the + # master rank + + # this context manager builds the problem on every rank, then puts the workers in + # a loop to listen to the master to compute the objective function and its + # derivatives. Only the master rank gets is_root=True. + with run_with_mpi(obj, cons) as is_root: + # apart from cost evaluation and derivatives, everything else will be only + # performed on the master rank + if is_root: + eq.solve( + objective=obj, + constraints=cons, + maxiter=10, + ftol=0, + gtol=0, + xtol=0, + verbose=3, + ) + + # if you put a code here, it will be performed on all ranks diff --git a/docs/notebooks/tutorials/mpi-tutorials/mpi-proximal.py b/docs/notebooks/tutorials/mpi-tutorials/mpi-proximal.py new file mode 100644 index 0000000000..d4bbaa680c --- /dev/null +++ b/docs/notebooks/tutorials/mpi-tutorials/mpi-proximal.py @@ -0,0 +1,160 @@ +import os +import sys + +# Add the path to the parent directory to augment search for module +sys.path.insert(0, os.path.abspath(".")) +sys.path.append(os.path.abspath("../../../")) +sys.path.append(os.path.abspath("../../../../")) + +from mpi4py import MPI + +from desc import _set_cpu_count, set_device + +kind = "cpu" # or "gpu" +num_device = 2 +# ====== Using CPUs ====== +# These will be used for dividing the single CPU into multiple virtual CPUs +# such that JAX and XLA thinks there are multiple devices +if kind == "cpu": + # !!! If you have multiple CPUs, you shouldn't call `_set_cpu_count` !!! + _set_cpu_count(num_device) + set_device("cpu", num_device=num_device, mpi=MPI) + +# ====== Using GPUs ====== +# When we have multiple processes using the same devices (for example, 3 processes +# using 3 GPUs), each process will try to pre-allocate 75% of the GPU memory which will +# cause the memory allocation to fail. To avoid this, we can set the allocator to `platform` +# such that there is no pre-allocation. This is a bit conservative (and probably there is room +# for improvement), but if a process needs more memory, it can use more memory on the fly. +elif kind == "gpu": + os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" + set_device("gpu", num_device=num_device) + + +import numpy as np + +from desc import config as desc_config +from desc.backend import jax, jnp, print_backend_info +from desc.examples import get +from desc.grid import LinearGrid +from desc.objectives import ( + AspectRatio, + FixBoundaryR, + FixBoundaryZ, + FixCurrent, + FixPressure, + FixPsi, + ForceBalance, + ObjectiveFunction, + QuasisymmetryTwoTerm, +) +from desc.optimize import Optimizer, run_with_mpi + +if __name__ == "__main__": + rank = MPI.COMM_WORLD.Get_rank() + size = MPI.COMM_WORLD.Get_size() + if rank == 0: + print(f"====== TOTAL OF {size} RANKS ======") + + # see which rank is running on which device + # Note: JAX has 2 functions for this: `jax.devices()` and `jax.local_devices()` + # `jax.devices()` will return all devices available to JAX, while `jax.local_devices()` + # will return only the devices that are available to the current process. This is + # useful when you have multiple processes running on multiple nodes and you want + # to see which devices are available to each process. + if desc_config["kind"] == "gpu": + print( + f"Rank {rank} is running on {jax.local_devices(backend='gpu')} " + f"and {jax.local_devices(backend='cpu')}\n" + ) + else: + print(f"Rank {rank} is running on {jax.local_devices(backend='cpu')}\n") + + if rank == 0: + print("====== BACKEND INFO ======") + print_backend_info() + print("\n") + + eq = get("precise_QA") + if desc_config["kind"] == "cpu": + eq.change_resolution(M=3, N=2, M_grid=6, N_grid=4) + + # create two grids with different rho values, this will effectively separate + # the quasisymmetry objective into two parts + grid1 = LinearGrid( + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + rho=jnp.linspace(0.2, 0.5, 4), + sym=True, + ) + grid2 = LinearGrid( + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + rho=jnp.linspace(0.6, 1.0, 6), + sym=True, + ) + + # when using parallel objectives, the user needs to supply the device_id and rank + obj1 = QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid1, device_id=0, rank=0 + ) + obj2 = QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid2, device_id=1, rank=1 + ) + obj3 = AspectRatio(eq=eq, target=8, weight=100, device_id=0, rank=0) + objs = [obj1, obj2, obj3] + + # Parallel objective function needs the MPI communicator + # If you don't specify `deriv_mode=blocked`, you will get a warning and DESC will + # automatically switch to `blocked`. + # this is not built here, `run_with_mpi` below builds it on every rank + objective = ObjectiveFunction(objs, deriv_mode="blocked", mpi=MPI) + + # we will fix some modes as usual + R_modes = np.vstack( + ( + [0, 0, 0], + eq.surface.R_basis.modes[ + np.max(np.abs(eq.surface.R_basis.modes), 1) > 1, : + ], + ) + ) + Z_modes = eq.surface.Z_basis.modes[ + np.max(np.abs(eq.surface.Z_basis.modes), 1) > 1, : + ] + # nonlinear constraints can be given a device_id and rank to run them on different + # devices as well, but that is not supported by the proximal wrapper yet, so here + # ForceBalance is computed on the master rank + constraints = ( + ForceBalance(eq=eq), + FixBoundaryR(eq=eq, modes=R_modes), + FixBoundaryZ(eq=eq, modes=Z_modes), + FixPressure(eq=eq), + FixPsi(eq=eq), + FixCurrent(eq=eq), + ) + optimizer = Optimizer("proximal-lsq-exact") + + # Until this line, the code is performed on all ranks, so it might print some + # information multiple times. The following part will only be performed on the + # master rank + + # this context manager builds the problem on every rank, then puts the workers in + # a loop to listen to the master to compute the objective function and its + # derivatives. Only the master rank gets is_root=True, and prints. + with run_with_mpi(objective, constraints, verbose=3) as is_root: + # apart from cost evaluation and derivatives, everything else will be only + # performed on the master rank + if is_root: + eq.optimize( + objective=objective, + constraints=constraints, + optimizer=optimizer, + maxiter=3, + verbose=3, + options={"initial_trust_ratio": 1.0}, + ) + + # if you put a code here, it will be performed on all ranks diff --git a/docs/notebooks/tutorials/multi_device.ipynb b/docs/notebooks/tutorials/multi_device.ipynb new file mode 100644 index 0000000000..c9c9628012 --- /dev/null +++ b/docs/notebooks/tutorials/multi_device.ipynb @@ -0,0 +1,802 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "source": [ + "# How to use Multiple Devices\n", + "\n", + "In this tutorial, we will see how to use multiple devices to run DESC. This will make the optimization problem scalable to computing clusters. In short, for existing scripts, all you need to do is to provide `device_id` and `rank` for objective and constraints, pass the MPI communicator to `ObjectiveFunction` and call the optimization in the following form,\n", + "\n", + "```python\n", + "from desc.optimize import run_with_mpi\n", + "\n", + "objs = (\n", + " Objective(..., device_id=0, rank=0), \n", + " Objective(..., device_id=1, rank=1),\n", + " Objective(..., device_id=2, rank=2),\n", + ")\n", + "obj = ObjectiveFunction(objs, mpi=MPI)\n", + "cons = (\n", + " Objective(..., device_id=0, rank=0), \n", + " Objective(..., device_id=1, rank=1),\n", + ")\n", + "with run_with_mpi(obj, cons, verbose=3) as is_root:\n", + " if is_root:\n", + " eq.optimize(obj, cons, ...)\n", + "```\n", + "\n", + "This tutorials will not be able to run on a Jupyter Notebook, so we will provide the content of the script here but run an underlying python script to show the results.\n", + "\n", + "## Solving Equilibrium" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "import os\n", + "\n", + "sys.path.insert(0, os.path.abspath(\".\"))\n", + "sys.path.append(os.path.abspath(\"../../../\"))\n", + "\n", + "from IPython.display import Markdown" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/markdown": [ + "```python\n", + "import os\n", + "import sys\n", + "\n", + "# Add the path to the parent directory to augment search for module\n", + "sys.path.insert(0, os.path.abspath(\".\"))\n", + "sys.path.append(os.path.abspath(\"../../../\"))\n", + "sys.path.append(os.path.abspath(\"../../../../\"))\n", + "\n", + "import numpy as np\n", + "from mpi4py import MPI\n", + "\n", + "from desc import _set_cpu_count, set_device\n", + "\n", + "kind = \"cpu\" # or \"gpu\"\n", + "num_device = 2\n", + "# ====== Using CPUs ======\n", + "# These will be used for dividing the single CPU into multiple virtual CPUs\n", + "# such that JAX and XLA thinks there are multiple devices\n", + "if kind == \"cpu\":\n", + " # !!! If you have multiple CPUs, you shouldn't call `_set_cpu_count` !!!\n", + " _set_cpu_count(num_device)\n", + " set_device(\"cpu\", num_device=num_device, mpi=MPI)\n", + "\n", + "# ====== Using GPUs ======\n", + "# When we have multiple processes using the same devices (for example, 3 processes\n", + "# using 3 GPUs), each process will try to pre-allocate 75% of the GPU memory which will\n", + "# cause the memory allocation to fail. To avoid this, we can set the allocator to `platform`\n", + "# such that there is no pre-allocation. This is a bit conservative (and probably there is room\n", + "# for improvement), but if a process needs more memory, it can use more memory on the fly.\n", + "elif kind == \"gpu\":\n", + " os.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n", + " set_device(\"gpu\", num_device=num_device)\n", + "\n", + "from desc import config as desc_config\n", + "from desc.backend import jax, print_backend_info\n", + "from desc.examples import get\n", + "from desc.grid import LinearGrid\n", + "from desc.objectives import ForceBalance, ObjectiveFunction\n", + "from desc.objectives.getters import get_fixed_boundary_constraints\n", + "from desc.optimize import run_with_mpi\n", + "\n", + "if __name__ == \"__main__\":\n", + " rank = MPI.COMM_WORLD.Get_rank()\n", + " size = MPI.COMM_WORLD.Get_size()\n", + " if rank == 0:\n", + " print(f\"====== TOTAL OF {size} RANKS ======\")\n", + "\n", + " # see which rank is running on which device\n", + " # Note: JAX has 2 functions for this: `jax.devices()` and `jax.local_devices()`\n", + " # `jax.devices()` will return all devices available to JAX, while `jax.local_devices()`\n", + " # will return only the devices that are available to the current process. This is\n", + " # useful when you have multiple processes running on multiple nodes and you want\n", + " # to see which devices are available to each process.\n", + " if desc_config[\"kind\"] == \"gpu\":\n", + " print(\n", + " f\"Rank {rank} can see {jax.local_devices(backend='gpu')} \"\n", + " f\"and {jax.local_devices(backend='cpu')}\\n\"\n", + " )\n", + " else:\n", + " print(f\"Rank {rank} can see {jax.local_devices(backend='cpu')}\\n\")\n", + "\n", + " if rank == 0:\n", + " print(\"====== BACKEND INFO ======\")\n", + " print_backend_info()\n", + " print(\"\\n\")\n", + "\n", + " eq = get(\"HELIOTRON\")\n", + " if desc_config[\"kind\"] == \"cpu\":\n", + " # for local testing use lower resolution\n", + " eq.change_resolution(M=3, N=2, M_grid=6, N_grid=4)\n", + "\n", + " # setup 2 grids for 2 objectives covering different flux surfaces\n", + " rhos = np.linspace(0.1, 1.0, eq.L_grid)\n", + " grid1 = LinearGrid(\n", + " rho=rhos[: rhos.size // 2],\n", + " M=eq.M_grid,\n", + " N=eq.N_grid,\n", + " NFP=eq.NFP,\n", + " )\n", + " grid2 = LinearGrid(\n", + " rho=rhos[rhos.size // 2 :],\n", + " M=eq.M_grid,\n", + " N=eq.N_grid,\n", + " NFP=eq.NFP,\n", + " )\n", + " # ranks will be automatically assigned as 0 and 1, respectively\n", + " obj = ObjectiveFunction(\n", + " [\n", + " ForceBalance(eq, grid=grid1, device_id=0),\n", + " ForceBalance(eq, grid=grid2, device_id=1),\n", + " ],\n", + " mpi=MPI,\n", + " deriv_mode=\"blocked\",\n", + " )\n", + " cons = get_fixed_boundary_constraints(eq)\n", + "\n", + " # Until this line, the code is performed on all ranks, so it might print some\n", + " # information multiple times. The following part will only be performed on the\n", + " # master rank\n", + "\n", + " # this context manager builds the problem on every rank, then puts the workers in\n", + " # a loop to listen to the master to compute the objective function and its\n", + " # derivatives. Only the master rank gets is_root=True.\n", + " with run_with_mpi(obj, cons) as is_root:\n", + " # apart from cost evaluation and derivatives, everything else will be only\n", + " # performed on the master rank\n", + " if is_root:\n", + " eq.solve(\n", + " objective=obj,\n", + " constraints=cons,\n", + " maxiter=10,\n", + " ftol=0,\n", + " gtol=0,\n", + " xtol=0,\n", + " verbose=3,\n", + " )\n", + "\n", + " # if you put a code here, it will be performed on all ranks\n", + "\n", + "```" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Display the content of mpi-eq-solve.py\n", + "with open(\"mpi-tutorials/mpi-eq-solve.py\", \"r\") as f:\n", + " code = f.read()\n", + "\n", + "Markdown(f\"```python\\n{code}\\n```\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "====== TOTAL OF 2 RANKS ======\n", + "Rank 0 can see [CpuDevice(id=0), CpuDevice(id=1)]\n", + "\n", + "====== BACKEND INFO ======\n", + "DESC version=0.17.3+271.ge76ccd55e.dirty.\n", + "Rank 1 can see [CpuDevice(id=0), CpuDevice(id=1)]\n", + "\n", + "Using JAX backend: jax version=0.9.2, jaxlib version=0.9.2, dtype=float64.\n", + "Using 2 CPUs with 18.87 GB total available memory:\n", + "\t CPU : 0 13th Gen Intel(R) Core(TM) i9-13900HX\n", + "\t CPU : 1 13th Gen Intel(R) Core(TM) i9-13900HX\n", + "\n", + "Note: The backend information assumes that the user has 1 process per CPU (node). Using multiple processes per CPU (node) is not the most efficient way to use MPI with purely CPUs.\n", + "\n", + "\n", + "Building objective: force\n", + "Precomputing transforms\n", + "Building objective: force\n", + "Precomputing transforms\n", + "Putting objective force on device 1\n", + "------------------------------------------------------------\n", + "Rank 0 will run objective(s): ['ForceBalance']\n", + "Rank 1 will run objective(s): ['ForceBalance']\n", + "------------------------------------------------------------\n", + "Building objective: lcfs R\n", + "Building objective: lcfs Z\n", + "Building objective: fixed Psi\n", + "Building objective: fixed pressure\n", + "Building objective: fixed iota\n", + "Building objective: fixed sheet current\n", + "Building objective: self_consistency R\n", + "Building objective: self_consistency Z\n", + "Building objective: lambda gauge\n", + "Building objective: axis R self consistency\n", + "Building objective: axis Z self consistency\n", + "\u001b[32mTimer: Objective build = 802 ms\u001b[0m\n", + "\u001b[32mTimer: LinearConstraintProjection build = 2.40 sec\u001b[0m\n", + "Number of parameters: 551\n", + "Number of objectives: 8424\n", + "\u001b[32mTimer: Initializing the optimization = 3.23 sec\u001b[0m\n", + "\n", + "Starting optimization\n", + "Using method: lsq-exact\n", + "Solver options:\n", + "------------------------------------------------------------\n", + "Maximum Function Evaluations : 51\n", + "Maximum Allowed Total Δx Norm : inf\n", + "Scaled Termination : True\n", + "Trust Region Method : qr\n", + "Initial Trust Radius : 5.583e+03\n", + "Maximum Trust Radius : inf\n", + "Minimum Trust Radius : 2.220e-16\n", + "Trust Radius Increase Ratio : 2.000e+00\n", + "Trust Radius Decrease Ratio : 2.500e-01\n", + "Trust Radius Increase Threshold : 7.500e-01\n", + "Trust Radius Decrease Threshold : 2.500e-01\n", + "------------------------------------------------------------ \n", + "\n", + " Iteration Total nfev Cost Cost reduction Step norm Optimality \n", + " 0 1 1.540e+00 1.132e+00 \n", + " 1 2 4.303e-01 1.110e+00 4.473e-01 3.546e-01 \n", + " 2 3 9.374e-02 3.365e-01 2.949e-01 1.240e-01 \n", + " 3 5 2.716e-02 6.658e-02 1.387e-01 1.557e-01 \n", + " 4 7 7.589e-03 1.957e-02 6.497e-02 7.389e-02 \n", + " 5 8 1.072e-03 6.517e-03 3.394e-02 2.471e-02 \n", + " 6 11 6.304e-04 4.417e-04 6.563e-03 5.988e-03 \n", + " 7 12 6.040e-04 2.638e-05 6.114e-03 6.742e-04 \n", + " 8 14 6.034e-04 6.960e-07 3.717e-03 1.056e-03 \n", + " 9 15 6.024e-04 9.979e-07 3.879e-03 1.004e-03 \n", + " 10 17 6.010e-04 1.326e-06 1.098e-03 2.785e-04 \n", + "Warning: Maximum number of iterations has been exceeded.\n", + " Current function value: 6.010e-04\n", + " Total delta_x: 3.276e-01\n", + " Iterations: 10\n", + " Function evaluations: 17\n", + " Jacobian evaluations: 11\n", + "\u001b[32mTimer: Solution time = 15.5 sec\u001b[0m\n", + "\u001b[32mTimer: Avg time per step = 1.41 sec\u001b[0m\n", + "==============================================================================================================\n", + " Start --> End\n", + "Total (sum of squares): 1.540e+00 --> 6.010e-04, \n", + "Maximum absolute Force error: 2.530e+05 --> 9.751e+03 (N)\n", + "Minimum absolute Force error: 1.089e-10 --> 1.311e-10 (N)\n", + "Average absolute Force error: 5.001e+04 --> 1.075e+03 (N)\n", + "Maximum absolute Force error: 2.035e-02 --> 7.842e-04 (normalized)\n", + "Minimum absolute Force error: 8.759e-18 --> 1.054e-17 (normalized)\n", + "Average absolute Force error: 4.022e-03 --> 8.648e-05 (normalized)\n", + "Maximum absolute Force error: 1.231e+07 --> 2.136e+05 (N)\n", + "Minimum absolute Force error: 2.182e-12 --> 3.323e-14 (N)\n", + "Average absolute Force error: 1.467e+05 --> 3.861e+03 (N)\n", + "Maximum absolute Force error: 9.903e-01 --> 1.718e-02 (normalized)\n", + "Minimum absolute Force error: 1.755e-19 --> 2.672e-21 (normalized)\n", + "Average absolute Force error: 1.180e-02 --> 3.105e-04 (normalized)\n", + "R boundary error: 0.000e+00 --> 0.000e+00 (m)\n", + "Z boundary error: 0.000e+00 --> 0.000e+00 (m)\n", + "Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)\n", + "Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)\n", + "Fixed iota profile error: 0.000e+00 --> 0.000e+00 (dimensionless)\n", + "Fixed sheet current error: 0.000e+00 --> 0.000e+00 (~)\n", + "==============================================================================================================\n", + "\n", + "Rank 1 STOPPING\n", + "\u001b[0m\u001b[0m" + ] + } + ], + "source": [ + "!mpirun -n 2 python mpi-tutorials/mpi-eq-solve.py" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using other Objectives\n", + "Above we used MPI for force balance objective, but we can also use it for general optimization.\n", + "\n", + "**Note:** Currently, if the optimizer solves the equilibrium at each step, this equilibrium solve cannot use MPI." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/markdown": [ + "```python\n", + "import os\n", + "import sys\n", + "\n", + "# Add the path to the parent directory to augment search for module\n", + "sys.path.insert(0, os.path.abspath(\".\"))\n", + "sys.path.append(os.path.abspath(\"../../../\"))\n", + "sys.path.append(os.path.abspath(\"../../../../\"))\n", + "\n", + "from mpi4py import MPI\n", + "\n", + "from desc import _set_cpu_count, set_device\n", + "\n", + "kind = \"cpu\" # or \"gpu\"\n", + "num_device = 2\n", + "# ====== Using CPUs ======\n", + "# These will be used for dividing the single CPU into multiple virtual CPUs\n", + "# such that JAX and XLA thinks there are multiple devices\n", + "if kind == \"cpu\":\n", + " # !!! If you have multiple CPUs, you shouldn't call `_set_cpu_count` !!!\n", + " _set_cpu_count(num_device)\n", + " set_device(\"cpu\", num_device=num_device, mpi=MPI)\n", + "\n", + "# ====== Using GPUs ======\n", + "# When we have multiple processes using the same devices (for example, 3 processes\n", + "# using 3 GPUs), each process will try to pre-allocate 75% of the GPU memory which will\n", + "# cause the memory allocation to fail. To avoid this, we can set the allocator to `platform`\n", + "# such that there is no pre-allocation. This is a bit conservative (and probably there is room\n", + "# for improvement), but if a process needs more memory, it can use more memory on the fly.\n", + "elif kind == \"gpu\":\n", + " os.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n", + " set_device(\"gpu\", num_device=num_device)\n", + "\n", + "\n", + "import numpy as np\n", + "\n", + "from desc import config as desc_config\n", + "from desc.backend import jax, jnp, print_backend_info\n", + "from desc.examples import get\n", + "from desc.grid import LinearGrid\n", + "from desc.objectives import (\n", + " AspectRatio,\n", + " FixBoundaryR,\n", + " FixBoundaryZ,\n", + " FixCurrent,\n", + " FixPressure,\n", + " FixPsi,\n", + " ForceBalance,\n", + " ObjectiveFunction,\n", + " QuasisymmetryTwoTerm,\n", + ")\n", + "from desc.optimize import Optimizer, run_with_mpi\n", + "\n", + "if __name__ == \"__main__\":\n", + " rank = MPI.COMM_WORLD.Get_rank()\n", + " size = MPI.COMM_WORLD.Get_size()\n", + " if rank == 0:\n", + " print(f\"====== TOTAL OF {size} RANKS ======\")\n", + "\n", + " # see which rank is running on which device\n", + " # Note: JAX has 2 functions for this: `jax.devices()` and `jax.local_devices()`\n", + " # `jax.devices()` will return all devices available to JAX, while `jax.local_devices()`\n", + " # will return only the devices that are available to the current process. This is\n", + " # useful when you have multiple processes running on multiple nodes and you want\n", + " # to see which devices are available to each process.\n", + " if desc_config[\"kind\"] == \"gpu\":\n", + " print(\n", + " f\"Rank {rank} is running on {jax.local_devices(backend='gpu')} \"\n", + " f\"and {jax.local_devices(backend='cpu')}\\n\"\n", + " )\n", + " else:\n", + " print(f\"Rank {rank} is running on {jax.local_devices(backend='cpu')}\\n\")\n", + "\n", + " if rank == 0:\n", + " print(\"====== BACKEND INFO ======\")\n", + " print_backend_info()\n", + " print(\"\\n\")\n", + "\n", + " eq = get(\"precise_QA\")\n", + " if desc_config[\"kind\"] == \"cpu\":\n", + " eq.change_resolution(M=3, N=2, M_grid=6, N_grid=4)\n", + "\n", + " # create two grids with different rho values, this will effectively separate\n", + " # the quasisymmetry objective into two parts\n", + " grid1 = LinearGrid(\n", + " M=eq.M_grid,\n", + " N=eq.N_grid,\n", + " NFP=eq.NFP,\n", + " rho=jnp.linspace(0.2, 0.5, 4),\n", + " sym=True,\n", + " )\n", + " grid2 = LinearGrid(\n", + " M=eq.M_grid,\n", + " N=eq.N_grid,\n", + " NFP=eq.NFP,\n", + " rho=jnp.linspace(0.6, 1.0, 6),\n", + " sym=True,\n", + " )\n", + "\n", + " # when using parallel objectives, the user needs to supply the device_id and rank\n", + " obj1 = QuasisymmetryTwoTerm(\n", + " eq=eq, helicity=(1, eq.NFP), grid=grid1, device_id=0, rank=0\n", + " )\n", + " obj2 = QuasisymmetryTwoTerm(\n", + " eq=eq, helicity=(1, eq.NFP), grid=grid2, device_id=1, rank=1\n", + " )\n", + " obj3 = AspectRatio(eq=eq, target=8, weight=100, device_id=0, rank=0)\n", + " objs = [obj1, obj2, obj3]\n", + "\n", + " # Parallel objective function needs the MPI communicator\n", + " # If you don't specify `deriv_mode=blocked`, you will get a warning and DESC will\n", + " # automatically switch to `blocked`.\n", + " # this is not built here, `run_with_mpi` below builds it on every rank\n", + " objective = ObjectiveFunction(objs, deriv_mode=\"blocked\", mpi=MPI)\n", + "\n", + " # we will fix some modes as usual\n", + " R_modes = np.vstack(\n", + " (\n", + " [0, 0, 0],\n", + " eq.surface.R_basis.modes[\n", + " np.max(np.abs(eq.surface.R_basis.modes), 1) > 1, :\n", + " ],\n", + " )\n", + " )\n", + " Z_modes = eq.surface.Z_basis.modes[\n", + " np.max(np.abs(eq.surface.Z_basis.modes), 1) > 1, :\n", + " ]\n", + " # nonlinear constraints can be given a device_id and rank to run them on different\n", + " # devices as well, but that is not supported by the proximal wrapper yet, so here\n", + " # ForceBalance is computed on the master rank\n", + " constraints = (\n", + " ForceBalance(eq=eq),\n", + " FixBoundaryR(eq=eq, modes=R_modes),\n", + " FixBoundaryZ(eq=eq, modes=Z_modes),\n", + " FixPressure(eq=eq),\n", + " FixPsi(eq=eq),\n", + " FixCurrent(eq=eq),\n", + " )\n", + " optimizer = Optimizer(\"proximal-lsq-exact\")\n", + "\n", + " # Until this line, the code is performed on all ranks, so it might print some\n", + " # information multiple times. The following part will only be performed on the\n", + " # master rank\n", + "\n", + " # this context manager builds the problem on every rank, then puts the workers in\n", + " # a loop to listen to the master to compute the objective function and its\n", + " # derivatives. Only the master rank gets is_root=True, and prints.\n", + " with run_with_mpi(objective, constraints, verbose=3) as is_root:\n", + " # apart from cost evaluation and derivatives, everything else will be only\n", + " # performed on the master rank\n", + " if is_root:\n", + " eq.optimize(\n", + " objective=objective,\n", + " constraints=constraints,\n", + " optimizer=optimizer,\n", + " maxiter=3,\n", + " verbose=3,\n", + " options={\"initial_trust_ratio\": 1.0},\n", + " )\n", + "\n", + " # if you put a code here, it will be performed on all ranks\n", + "\n", + "```" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Display the content of mpi-proximal.py\n", + "with open(\"mpi-tutorials/mpi-proximal.py\", \"r\") as f:\n", + " code = f.read()\n", + "\n", + "Markdown(f\"```python\\n{code}\\n```\")" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Rank 1 is running on [CpuDevice(id=0), CpuDevice(id=1)]\n", + "\n", + "====== TOTAL OF 2 RANKS ======\n", + "Rank 0 is running on [CpuDevice(id=0), CpuDevice(id=1)]\n", + "\n", + "====== BACKEND INFO ======\n", + "DESC version=0.17.3+271.ge76ccd55e.dirty.\n", + "Using JAX backend: jax version=0.9.2, jaxlib version=0.9.2, dtype=float64.\n", + "Using 2 CPUs with 18.76 GB total available memory:\n", + "\t CPU : 0 13th Gen Intel(R) Core(TM) i9-13900HX\n", + "\t CPU : 1 13th Gen Intel(R) Core(TM) i9-13900HX\n", + "\n", + "Note: The backend information assumes that the user has 1 process per CPU (node). Using multiple processes per CPU (node) is not the most efficient way to use MPI with purely CPUs.\n", + "\n", + "\n", + "Building objective: QS two-term\n", + "Precomputing transforms\n", + "\u001b[32mTimer: Precomputing transforms = 1.01 sec\u001b[0m\n", + "Building objective: QS two-term\n", + "Precomputing transforms\n", + "\u001b[32mTimer: Precomputing transforms = 1.01 sec\u001b[0m\n", + "Putting objective QS two-term on device 1\n", + "Building objective: aspect ratio\n", + "Precomputing transforms\n", + "\u001b[32mTimer: Precomputing transforms = 939 ms\u001b[0m\n", + "------------------------------------------------------------\n", + "Rank 0 will run objective(s): ['QuasisymmetryTwoTerm', 'AspectRatio']\n", + "Rank 1 will run objective(s): ['QuasisymmetryTwoTerm']\n", + "------------------------------------------------------------\n", + "\u001b[32mTimer: Objective build = 3.79 sec\u001b[0m\n", + "Building objective: force\n", + "Precomputing transforms\n", + "\u001b[32mTimer: Precomputing transforms = 1.32 sec\u001b[0m\n", + "\u001b[32mTimer: Objective build = 1.38 sec\u001b[0m\n", + "\u001b[32mTimer: Objective build = 1.08 ms\u001b[0m\n", + "\u001b[32mTimer: Eq Update LinearConstraintProjection build = 2.52 sec\u001b[0m\n", + "\u001b[32mTimer: Proximal projection build = 12.7 sec\u001b[0m\n", + "Building objective: lcfs R\n", + "Building objective: lcfs Z\n", + "Building objective: fixed pressure\n", + "Building objective: fixed Psi\n", + "Building objective: fixed current\n", + "\u001b[32mTimer: Objective build = 611 ms\u001b[0m\n", + "\u001b[32mTimer: LinearConstraintProjection build = 1.10 sec\u001b[0m\n", + "Number of parameters: 8\n", + "Number of objectives: 631\n", + "\u001b[32mTimer: Initializing the optimization = 14.5 sec\u001b[0m\n", + "\n", + "Starting optimization\n", + "Using method: proximal-lsq-exact\n", + "Solver options:\n", + "------------------------------------------------------------\n", + "Maximum Function Evaluations : 16\n", + "Maximum Allowed Total Δx Norm : inf\n", + "Scaled Termination : True\n", + "Trust Region Method : qr\n", + "Initial Trust Radius : 6.219e+02\n", + "Maximum Trust Radius : inf\n", + "Minimum Trust Radius : 2.220e-16\n", + "Trust Radius Increase Ratio : 2.000e+00\n", + "Trust Radius Decrease Ratio : 2.500e-01\n", + "Trust Radius Increase Threshold : 7.500e-01\n", + "Trust Radius Decrease Threshold : 2.500e-01\n", + "------------------------------------------------------------ \n", + "\n", + " Iteration Total nfev Cost Cost reduction Step norm Optimality \n", + " 0 1 2.001e+04 1.870e+02 \n", + " 1 4 8.742e+03 1.126e+04 3.689e-02 8.655e+01 \n", + " 2 5 3.984e+03 4.758e+03 6.537e-02 6.133e+01 \n", + " 3 6 2.205e+03 1.779e+03 8.708e-02 1.535e+01 \n", + "Warning: Maximum number of iterations has been exceeded.\n", + " Current function value: 2.205e+03\n", + " Total delta_x: 1.233e-01\n", + " Iterations: 3\n", + " Function evaluations: 6\n", + " Jacobian evaluations: 4\n", + "\u001b[32mTimer: Solution time = 32.6 sec\u001b[0m\n", + "\u001b[32mTimer: Avg time per step = 8.17 sec\u001b[0m\n", + "==============================================================================================================\n", + " Start --> End\n", + "Total (sum of squares): 2.001e+04 --> 2.205e+03, \n", + "Maximum absolute Quasi-symmetry (1,2) two-term error: 1.910e-01 --> 9.784e-01 (T^3)\n", + "Minimum absolute Quasi-symmetry (1,2) two-term error: 2.766e-07 --> 1.949e-03 (T^3)\n", + "Average absolute Quasi-symmetry (1,2) two-term error: 6.521e-02 --> 2.375e-01 (T^3)\n", + "Maximum absolute Quasi-symmetry (1,2) two-term error: 3.429e-01 --> 1.757e+00 (normalized)\n", + "Minimum absolute Quasi-symmetry (1,2) two-term error: 4.966e-07 --> 3.500e-03 (normalized)\n", + "Average absolute Quasi-symmetry (1,2) two-term error: 1.171e-01 --> 4.265e-01 (normalized)\n", + "Maximum absolute Quasi-symmetry (1,2) two-term error: 1.425e+01 --> 7.592e+00 (T^3)\n", + "Minimum absolute Quasi-symmetry (1,2) two-term error: 1.678e-03 --> 8.196e-04 (T^3)\n", + "Average absolute Quasi-symmetry (1,2) two-term error: 2.460e-01 --> 3.879e-01 (T^3)\n", + "Maximum absolute Quasi-symmetry (1,2) two-term error: 2.558e+01 --> 1.363e+01 (normalized)\n", + "Minimum absolute Quasi-symmetry (1,2) two-term error: 3.012e-03 --> 1.472e-03 (normalized)\n", + "Average absolute Quasi-symmetry (1,2) two-term error: 4.417e-01 --> 6.965e-01 (normalized)\n", + "Aspect ratio: 8.000e+00 --> 8.000e+00 (dimensionless)\n", + "Maximum absolute Force error: 5.997e+05 --> 2.364e+04 (N)\n", + "Minimum absolute Force error: 3.370e+00 --> 5.342e+00 (N)\n", + "Average absolute Force error: 1.335e+04 --> 2.503e+03 (N)\n", + "Maximum absolute Force error: 5.499e-01 --> 2.168e-02 (normalized)\n", + "Minimum absolute Force error: 3.090e-06 --> 4.898e-06 (normalized)\n", + "Average absolute Force error: 1.224e-02 --> 2.295e-03 (normalized)\n", + "R boundary error: 5.081e-18 --> 3.981e-18 (m)\n", + "Z boundary error: 2.877e-18 --> 2.453e-18 (m)\n", + "Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)\n", + "Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)\n", + "Fixed current profile error: 0.000e+00 --> 0.000e+00 (A)\n", + "==============================================================================================================\n", + "\n", + "Rank 1 STOPPING\n", + "\u001b[0m\u001b[0m" + ] + } + ], + "source": [ + "!mpirun -n 2 python mpi-tutorials/mpi-proximal.py" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Slurm for Multi-Node and Multi-Process Scripts\n", + "\n", + "**Note :** These instructions may differ for the cluster you are trying to use. The reason we give this example is to set some terminology for users that are not familiar with multi-node and multi-processing.\n", + "\n", + "**Note :** For more details, one can check Princeton University Research Computing page [here](https://researchcomputing.princeton.edu/support/knowledge-base/slurm#Multinode--Multithreaded-Jobs).\n", + "\n", + "One needs to use proper slurm script to run parallel code on a cluster. Here, we will give an example in which we use 2 nodes, 8 processes per node and 4 CPU cores per process. *Node* means the actual CPU chip, so we will have 2 CPUs (you can think of it as having 2 computers that are connected to each other). We will have 16 processes and 64 CPU cores in total. Additionally, you can specify the number of GPUs per node." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```bash\n", + "\n", + "#!/bin/bash\n", + "#SBATCH --job-name=mpi-example # create a short name for your job\n", + "#SBATCH --nodes=2 # node count\n", + "#SBATCH --ntasks-per-node=8 # total number of tasks per node\n", + "#SBATCH --cpus-per-task=4 # cpu-cores per task (>1 if multi-threaded tasks)\n", + "#SBATCH --mem-per-cpu=4G # memory per cpu-core (4G is default)\n", + "#SBATCH --time=00:10:00 # total run time limit (HH:MM:SS)\n", + "#SBATCH --gres=gpu:4 # number of GPUs per node (in this case 8 GPUs in total)\n", + "\n", + "export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\n", + "export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK\n", + "module purge\n", + "\n", + "# module names and version might be different for clusters\n", + "module load anaconda3/2024.6\n", + "module load openmpi/gcc/4.1.6\n", + "\n", + "# activate the environment that has DESC requirements\n", + "# as well as proper mpi4py installation\n", + "conda activate mpi-env\n", + "\n", + "srun python your-script.py\n", + "\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When using MPI with multiple nodes, each process will see 1 CPU (with multiple cores), and if you requested GPUs, only the GPUs connected to that CPU will be visible to your program. For example, if you want to use 2 nodes with 3 GPUs and 3 processes per node, you can have 6 objectives each on an independent device.\n", + "\n", + "```python\n", + "\n", + "# each node will see 3 GPUs\n", + "num_device = 3\n", + "os.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n", + "set_device(\"gpu\", num_device=num_device)\n", + "\n", + "\n", + "...\n", + "\n", + "\n", + "# this will run on node 1, GPU 0 (rank=0)\n", + "obj1 = QuasisymmetryTwoTerm(eq=eq, helicity=(1, eq.NFP), grid=grid1, device_id=0)\n", + "# this will run on node 1, GPU 1 (rank=1)\n", + "obj2 = QuasisymmetryTwoTerm(eq=eq, helicity=(1, eq.NFP), grid=grid2, device_id=1)\n", + "# this will run on node 1, GPU 2 (rank=2)\n", + "obj3 = QuasisymmetryTwoTerm(eq=eq, helicity=(1, eq.NFP), grid=grid3, device_id=2)\n", + "# this will run on node 2, GPU 0 (rank=3)\n", + "obj4 = AspectRatio(eq=eq, target=8, weight=100, device_id=0)\n", + "# this will run on node 2, GPU 1 (rank=4)\n", + "obj5 = Objective(..., device_id=1)\n", + "# this will run on node 2, GPU 2 (rank=5)\n", + "obj6 = Objective(..., device_id=2)\n", + "objs = [obj1, obj2, obj3, obj4, obj5, obj6]\n", + "\n", + "# Parallel objective function needs the MPI communicator\n", + "objective = ObjectiveFunction(objs, deriv_mode=\"blocked\", mpi=MPI)\n", + "\n", + "```\n", + "\n", + "When you write your script for multiple nodes, the number of devices and the device IDs must be selected as if there is only 1 node and only the local GPUs are visible. Other nodes will be used through `rank` of MPI communicator.\n", + "\n", + "**Note: Most clusters have multiple GPUs connected to each node, so before using multiple nodes, use all the GPUs available to that node. Multi-node communication is slower and your script will be easier to write properly.**\n", + "\n", + "**Note: If you want to run multiple objectives on the same device, you can specify the ``rank`` keywords for each sub-objective. By default, the initializer will assign different ranks for each sub-objective, in above example it defaults to `np.arange(len(objs))`. Keep in mind that if a rank is given to any sub-objective, it must be given to all of them. Moreover, given ranks and device ids must be consistent, i.e. different devices should have different ranks.** \n", + "\n", + "An example optimization setup where each rank has 1 or more objectives,\n", + "\n", + "```python\n", + "\n", + "# each node will see 3 GPUs\n", + "num_device = 3\n", + "os.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n", + "set_device(\"gpu\", num_device=num_device)\n", + "\n", + "\n", + "...\n", + "\n", + "\n", + "# this will run on node 1, GPU 0 (rank=0)\n", + "obj1 = SomeObjective(..., device_id=0, rank=0)\n", + "# this will run on node 1, GPU 1 (rank=1)\n", + "obj2 = SomeObjective(..., device_id=1, rank=1)\n", + "# this will run on node 1, GPU 1 (rank=1)\n", + "obj3 = SomeObjective(..., device_id=1, rank=1)\n", + "# this will run on node 1, GPU 2 (rank=2)\n", + "obj4 = SomeObjective(..., device_id=2, rank=2)\n", + "# this will run on node 2, GPU 0 (rank=3)\n", + "obj5 = SomeObjective(..., device_id=0, rank=3)\n", + "# this will run on node 2, GPU 0 (rank=3)\n", + "obj6 = SomeObjective(..., device_id=0, rank=3)\n", + "# this will run on node 2, GPU 1 (rank=4)\n", + "obj7 = SomeObjective(..., device_id=1, rank=4)\n", + "# this will run on node 2, GPU 2 (rank=5)\n", + "obj8 = SomeObjective(..., device_id=2, rank=5)\n", + "# this will run on node 2, GPU 2 (rank=5)\n", + "obj9 = SomeObjective(..., device_id=2, rank=5)\n", + "objs = [obj1, obj2, obj3, obj4, obj5, obj6, obj7, obj8, obj9]\n", + "\n", + "objective = ObjectiveFunction(\n", + " objs, \n", + " deriv_mode=\"blocked\", \n", + " mpi=MPI, \n", + ")\n", + "\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Tip\n", + "MPI parallelism is only used to compute the value and the derivatives of the objectives. The heavy linear algebra operations are not parallelized and will always run on the `rank=0`, so if you really want to maximize the memory limit, give the least memory hungry sub-objectives to `rank=0` or don't give any at all. This way, you will be able to fit bigger Jacobians to the memory of the first device. " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "mpi", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/setup.cfg b/setup.cfg index 64d6be0c06..3c3ec7e3d9 100644 --- a/setup.cfg +++ b/setup.cfg @@ -19,8 +19,10 @@ source = desc/ # _version.py is generated code, no need to count it +# __init__.py deals with device selection that CI cannot test omit = desc/_version.py + desc/__init__.py desc/examples/precise_QH.py desc/examples/precise_QA.py desc/examples/reactor_QA.py @@ -48,6 +50,8 @@ markers= slow: marks tests as slow (deselect with 'pytest -m "not slow"'). fast: mark tests as fast. memory: marks tests that check memory usage + mpi_setup: marks tests that require MPI but not MPI processes + mpi_run: marks tests that require MPI and need MPI processes filterwarnings= error ignore::pytest.PytestUnraisableExceptionWarning @@ -81,10 +85,12 @@ per-file-ignores = desc/compute/data_index.py: E501 # need imports in weird order for selecting device before benchmarks tests/benchmarks/*.py: E402 + tests/test_multidevice.py: E402 # stop complaining about setting gpu before import other desc stuff desc/examples/precise_QA.py: E402 desc/examples/precise_QH.py: E402 desc/examples/reactor_QA.py: E402 + docs/notebooks/tutorials/mpi-tutorials/*.py: E402 max-line-length = 88 exclude = docs/* diff --git a/tests/test_multidevice.py b/tests/test_multidevice.py new file mode 100644 index 0000000000..a3e6a24a2e --- /dev/null +++ b/tests/test_multidevice.py @@ -0,0 +1,623 @@ +"""Tests for the multidevice capabilities.""" + +import warnings + +# This file has to run on a separate process because it changes the number of CPUs +from desc import _set_cpu_count, set_device + +num_device = 3 +with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _set_cpu_count(num_device) + set_device(kind="cpu", num_device=num_device) + +import numpy as np +import pytest + +try: + from mpi4py import MPI +except ModuleNotFoundError: + print("mpi4py is not installed, skipping MPI tests.") + pytest.skip("mpi4py is not installed, skipping MPI tests.", allow_module_level=True) + +from desc import config as desc_config +from desc.examples import get +from desc.grid import LinearGrid +from desc.objectives import ( + AspectRatio, + FixBoundaryR, + FixBoundaryZ, + FixCurrent, + FixPressure, + FixPsi, + ForceBalance, + ObjectiveFunction, + QuasisymmetryTwoTerm, + get_fixed_boundary_constraints, +) +from desc.optimize import ( + LinearConstraintProjection, + Optimizer, + ProximalProjection, + build_for_mpi, + run_with_mpi, +) + + +@pytest.mark.mpi_setup +def test_set_cpu_count(): + """Test that _set_cpu_count works.""" + # we already called the function, just check the desc_config + assert desc_config["kind"] == "cpu" + assert desc_config["num_device"] == num_device + assert len(desc_config["devices"]) == num_device + assert len(desc_config["avail_mems"]) == num_device + + +@pytest.mark.mpi_run +def test_multidevice_objective_attributes(): + """Test that objective attributes are same.""" + eq = get("precise_QH") + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.8], sym=True) + + obj1 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1), + ForceBalance(eq, grid=grid2), + ForceBalance(eq, grid=grid3), + ], + deriv_mode="blocked", + ) + obj1.build() + + # deriv_mode will be set to "blocked" automatically + with pytest.warns(UserWarning, match="When using multiple devices"): + obj2 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1, device_id=0, rank=0), + ForceBalance(eq, grid=grid2, device_id=1, rank=1), + ForceBalance(eq, grid=grid3, device_id=2, rank=2), + ], + mpi=MPI, + ) + obj2.build() + + for obj1i, obj2i in zip(obj1.objectives, obj2.objectives): + assert obj1i._loss_function == obj2i._loss_function + np.testing.assert_allclose(obj1i._weight, obj2i._weight) + np.testing.assert_allclose(obj1i._target, obj2i._target) + np.testing.assert_allclose(obj1i._normalization, obj2i._normalization) + np.testing.assert_allclose(obj1i._dim_f, obj2i._dim_f) + key = "quad_weights" + np.testing.assert_allclose( + obj1i._constants[key], obj2i._constants[key], err_msg=key + ) + + +@pytest.mark.mpi_run +def test_multidevice_compute(): + """Test that objective compute gives same results.""" + eq = get("precise_QH") + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.8], sym=True) + + obj1 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1), + ForceBalance(eq, grid=grid2), + ForceBalance(eq, grid=grid3), + ], + deriv_mode="blocked", + ) + obj1.build() + + # deriv_mode will be set to "blocked" automatically + with pytest.warns(UserWarning, match="When using multiple devices"): + obj2 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1, device_id=0, rank=0), + ForceBalance(eq, grid=grid2, device_id=1, rank=1), + ForceBalance(eq, grid=grid3, device_id=2, rank=2), + ], + mpi=MPI, + ) + obj2.build() + + with run_with_mpi(obj2) as is_root: + if is_root: + f1 = obj1.compute_scalar(obj1.x(eq)) + f2 = obj2.compute_scalar(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.compute_unscaled(obj1.x(eq)) + f2 = obj2.compute_unscaled(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.compute_scaled(obj1.x(eq)) + f2 = obj2.compute_scaled(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.compute_scaled_error(obj1.x(eq)) + f2 = obj2.compute_scaled_error(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + +@pytest.mark.mpi_run +def test_multidevice_derivatives(): + """Test that objective derivatives gives same results.""" + eq = get("precise_QH") + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.8], sym=True) + + obj1 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1), + ForceBalance(eq, grid=grid2), + ForceBalance(eq, grid=grid3), + ], + deriv_mode="blocked", + ) + obj1.build() + + # deriv_mode will be set to "blocked" automatically + with pytest.warns(UserWarning, match="When using multiple devices"): + obj2 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1, device_id=0, rank=0), + ForceBalance(eq, grid=grid2, device_id=1, rank=1), + ForceBalance(eq, grid=grid3, device_id=2, rank=2), + ], + mpi=MPI, + ) + obj2.build() + + with run_with_mpi(obj2) as is_root: + if is_root: + with pytest.raises(NotImplementedError): + _ = obj2.grad(obj2.x(eq)) + + f1 = obj1.jac_unscaled(obj1.x(eq)) + f2 = obj2.jac_unscaled(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.jac_scaled(obj1.x(eq)) + f2 = obj2.jac_scaled(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.jac_scaled_error(obj1.x(eq)) + f2 = obj2.jac_scaled_error(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + +@pytest.mark.mpi_run +def test_multidevice_linear_proj_derivatives(): + """Test that linear projection derivatives gives same results.""" + eq = get("precise_QH") + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.8], sym=True) + + objf1 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1), + ForceBalance(eq, grid=grid2), + ForceBalance(eq, grid=grid3), + ], + deriv_mode="blocked", + ) + objf1.build() + + # deriv_mode will be set to "blocked" automatically + with pytest.warns(UserWarning, match="When using multiple devices"): + objf2 = ObjectiveFunction( + [ + ForceBalance(eq, grid=grid1, device_id=0, rank=0), + ForceBalance(eq, grid=grid2, device_id=1, rank=1), + ForceBalance(eq, grid=grid3, device_id=2, rank=2), + ], + mpi=MPI, + ) + objf2.build() + + cons = get_fixed_boundary_constraints(eq) + cons = ObjectiveFunction(cons) + obj1 = LinearConstraintProjection(objective=objf1, constraint=cons) + obj2 = LinearConstraintProjection(objective=objf2, constraint=cons) + obj1.build() + obj2.build() + + with run_with_mpi(objf2) as is_root: + if is_root: + with pytest.raises(NotImplementedError): + _ = obj2.grad(obj2.x(eq)) + + f1 = obj1.jac_unscaled(obj1.x(eq)) + f2 = obj2.jac_unscaled(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.jac_scaled(obj1.x(eq)) + f2 = obj2.jac_scaled(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = obj1.jac_scaled_error(obj1.x(eq)) + f2 = obj2.jac_scaled_error(obj2.x(eq)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + +@pytest.mark.mpi_run +def test_multidevice_nonlinear_constraint_derivatives(): + """Test that parallel nonlinear constraints give same results.""" + eq = get("precise_QH") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(1, 1, 1, 2, 2, 2) + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2, 0.6], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.9], sym=True) + + objective = ObjectiveFunction( + [ + QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid1, device_id=0, rank=0 + ), + QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid2, device_id=1, rank=1 + ), + AspectRatio(eq=eq, target=8, device_id=2, rank=2), + ], + deriv_mode="blocked", + mpi=MPI, + ) + # the nonlinear constraints only use rank 0 and 1, rank 2 gets none of them + constraints = ( + ForceBalance(eq=eq, grid=grid1, device_id=0, rank=0), + ForceBalance(eq=eq, grid=grid2, device_id=1, rank=1), + ) + get_fixed_boundary_constraints(eq) + + # same constraints on a single device, to compare against + con1 = ObjectiveFunction( + [ForceBalance(eq=eq, grid=grid1), ForceBalance(eq=eq, grid=grid2)], + deriv_mode="blocked", + ) + con1.build(verbose=0) + + with run_with_mpi(objective, constraints, verbose=0) as is_root: + if is_root: + con2 = objective._constraints + assert con2._is_mpi + assert [len(ids) for ids in con2._obj_per_rank] == [1, 1, 0] + + x = objective.x(eq) + f1 = con1.compute_scaled_error(x) + f2 = con2.compute_scaled_error(x) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = con1.jac_scaled_error(x) + f2 = con2.jac_scaled_error(x) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + # constraints that are not given a rank are computed on the root rank as usual + objective = build_for_mpi( + objective, + (ForceBalance(eq=eq, grid=grid1),) + get_fixed_boundary_constraints(eq), + verbose=0, + ) + assert objective._constraints is None + + +@pytest.mark.mpi_run +def test_multidevice_proximal_derivatives(): + """Test that proximal derivatives gives same results.""" + eq = get("precise_QH") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(1, 1, 1, 2, 2, 2) + + eq1 = eq.copy() + eq2 = eq.copy() + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6, 0.8], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.9], sym=True) + + obj1 = QuasisymmetryTwoTerm(eq=eq1, helicity=(1, eq.NFP), grid=grid1) + obj2 = QuasisymmetryTwoTerm(eq=eq1, helicity=(1, eq.NFP), grid=grid2) + obj3 = QuasisymmetryTwoTerm(eq=eq1, helicity=(1, eq.NFP), grid=grid3) + objs = [obj1, obj2, obj3] + + objective1 = ObjectiveFunction(objs, deriv_mode="blocked") + objective1.build(verbose=0) + + con1 = ObjectiveFunction(ForceBalance(eq1)) + con1.build(verbose=0) + + obj1 = QuasisymmetryTwoTerm( + eq=eq2, helicity=(1, eq.NFP), grid=grid1, device_id=0, rank=0 + ) + obj2 = QuasisymmetryTwoTerm( + eq=eq2, helicity=(1, eq.NFP), grid=grid2, device_id=1, rank=1 + ) + obj3 = QuasisymmetryTwoTerm( + eq=eq2, helicity=(1, eq.NFP), grid=grid3, device_id=2, rank=2 + ) + objs = [obj1, obj2, obj3] + + objective2 = ObjectiveFunction(objs, deriv_mode="blocked", mpi=MPI) + objective2.build(verbose=0) + con2 = ObjectiveFunction(ForceBalance(eq2)) + con2.build(verbose=0) + + perturb_options = {"order": 1} + solve_options = {"maxiter": 1} + prox1 = ProximalProjection( + objective=objective1, + constraint=con1, + eq=eq1, + solve_options=solve_options, + perturb_options=perturb_options, + ) + prox2 = ProximalProjection( + objective=objective2, + constraint=con2, + eq=eq2, + solve_options=solve_options, + perturb_options=perturb_options, + ) + prox1.build() + prox2.build() + + with run_with_mpi(objective2) as is_root: + if is_root: + f1 = prox1.grad(prox1.x(eq1)) + f2 = prox2.grad(prox2.x(eq2)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = prox1.jac_unscaled(prox1.x(eq1)) + f2 = prox2.jac_unscaled(prox2.x(eq2)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = prox1.jac_scaled(prox1.x(eq1)) + f2 = prox2.jac_scaled(prox2.x(eq2)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + f1 = prox1.jac_scaled_error(prox1.x(eq1)) + f2 = prox2.jac_scaled_error(prox2.x(eq2)) + np.testing.assert_allclose(f2, f1, atol=1e-8) + + +@pytest.mark.mpi_run +def test_multidevice_objective_build(): + """Test that objective function build works fine.""" + eq = get("HELIOTRON") + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6, 0.8], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2, 0.6], sym=True) + grid4 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.4, 0.8, 0.9], sym=True) + + # default rank will be 0, 1, 2, 3, respectively + obj1 = ForceBalance(eq, grid=grid1, device_id=0) + obj2 = ForceBalance(eq, grid=grid2, device_id=1) + obj3 = ForceBalance(eq, grid=grid3, device_id=2) + obj4 = ForceBalance(eq, grid=grid4, device_id=0) + + # need to pass MPI communicator to the ObjectiveFunction + with pytest.raises(ValueError, match="MPI communicator"): + # this one is multi-device + obj = ObjectiveFunction([obj1, obj2, obj3]) + + # if any rank is given, all should be given + obj1._rank = 0 + with pytest.raises(ValueError, match="If a rank is given to any"): + # this one is multi-device + obj = ObjectiveFunction([obj1, obj2, obj3], mpi=MPI) + + # make it inconsistent with device_id + obj1._rank = 0 + obj2._rank = 2 + obj3._rank = 1 + # need to use multiple ranks if using multiple devices + with pytest.raises(ValueError, match="rank and device id are inconsistent"): + # this one is multi-device + obj = ObjectiveFunction([obj1, obj2, obj3], mpi=MPI) + + obj1._rank = 0 + obj2._rank = 1 + obj3._rank = 2 + obj = ObjectiveFunction([obj1, obj2, obj3], mpi=MPI) + # deriv_mode will be set to "blocked" automatically + with pytest.warns(UserWarning, match="When using multiple devices"): + obj.build() + + # reset objectives that are built + for o in [obj1, obj2, obj3]: + o._built = False + o._use_jit = True + + # this one is single device, and grids have different sizes + obj1._rank = 0 + obj4._rank = 0 + objj = ObjectiveFunction([obj1, obj4]) + objj.build() + + assert obj._is_mpi + assert not objj._is_mpi + + np.testing.assert_allclose(obj.x(eq), objj.x(eq)) + + # multi-device objective must be blocked + assert obj._deriv_mode == "blocked" + assert objj._deriv_mode == "batched" + + +@pytest.mark.mpi_run +def test_multidevice_eq_solve(): + """Test that eq.solve still reduces force error.""" + rank = MPI.COMM_WORLD.Get_rank() + size = MPI.COMM_WORLD.Get_size() + assert size == num_device + assert rank < num_device + + eq = get("HELIOTRON") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(6, 6, 3, 12, 12, 6) + + gM = eq.M_grid + gN = eq.N_grid + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6, 0.8], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.9], sym=True) + + obj1 = ForceBalance(eq, grid=grid1, device_id=0, rank=0) + obj2 = ForceBalance(eq, grid=grid2, device_id=1, rank=1) + obj3 = ForceBalance(eq, grid=grid3, device_id=2, rank=2) + + # deriv_mode will be set to "blocked" automatically + with pytest.warns(UserWarning, match="When using multiple devices"): + obj = ObjectiveFunction([obj1, obj2, obj3], mpi=MPI) + # there is no constraint here, this only builds the objective on every rank, + # run_with_mpi below then only has to start the worker loop + obj = build_for_mpi(obj) + + # creating grids like grid3 = [grid1, grid2] doesn't give the same + # node, spacing and weight ordering, so we can't compare the Jacobians + # or the objective values directly. Instead, we compare the objective + # values before and after a single iteration of the solver. This should + # always decrease the objective value. + with run_with_mpi(obj) as is_root: + if is_root: + f0 = obj.compute_scalar(obj.x(eq)).block_until_ready() + eq.solve(objective=obj, maxiter=2, verbose=3) + f1 = obj.compute_scalar(obj.x(eq)) + + assert f1 < f0 + + +@pytest.mark.mpi_run +def test_multidevice_eq_optimize(): + """Test that eq.optimize still reduces error.""" + rank = MPI.COMM_WORLD.Get_rank() + size = MPI.COMM_WORLD.Get_size() + assert size == num_device + assert rank < num_device + + eq = get("precise_QA") + eq.change_resolution(M=3, N=2, M_grid=6, N_grid=4) + eq_no_mpi = eq.copy() + + # create two grids with different rho values, this will effectively separate + # the quasisymmetry objective into two parts + gM = 2 + gN = 2 + grid1 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.2], sym=True) + grid2 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.6, 0.8], sym=True) + grid3 = LinearGrid(M=gM, N=gN, NFP=eq.NFP, rho=[0.9], sym=True) + + # we will fix some modes as usual + k = 1 + sRm = eq.surface.R_basis.modes + sZm = eq.surface.Z_basis.modes + R_modes = np.vstack(([0, 0, 0], sRm[np.max(np.abs(sRm), 1) > k, :])) + Z_modes = sZm[np.max(np.abs(sZm), 1) > k, :] + + verbose = 3 if rank == 0 else 0 + maxiter = 3 + + ### Single device optimization + + obj1 = QuasisymmetryTwoTerm(eq=eq_no_mpi, helicity=(1, eq.NFP), grid=grid1) + obj2 = QuasisymmetryTwoTerm(eq=eq_no_mpi, helicity=(1, eq.NFP), grid=grid2) + obj3 = QuasisymmetryTwoTerm(eq=eq_no_mpi, helicity=(1, eq.NFP), grid=grid3) + obj4 = AspectRatio(eq=eq_no_mpi, target=8, weight=100) + objs = [obj1, obj2, obj3, obj4] + + objective = ObjectiveFunction(objs, deriv_mode="blocked") + objective.build(verbose=verbose) + + constraints = ( + ForceBalance(eq=eq_no_mpi), + FixBoundaryR(eq=eq_no_mpi, modes=R_modes), + FixBoundaryZ(eq=eq_no_mpi, modes=Z_modes), + FixPressure(eq=eq_no_mpi), + FixPsi(eq=eq_no_mpi), + FixCurrent(eq=eq_no_mpi), + ) + optimizer = Optimizer("proximal-lsq-exact") + eq_no_mpi.optimize( + objective=objective, + constraints=constraints, + optimizer=optimizer, + maxiter=maxiter, + verbose=verbose, + ) + x1_no_mpi = objective.x(eq_no_mpi) + f1_no_mpi = objective.compute_scalar(x1_no_mpi) + + ### Multidevice optimization + + # Wait for everyone to finish their work before proceeding + MPI.COMM_WORLD.Barrier() + + # when using parallel objectives, the user needs to supply the device_id and rank + obj1 = QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid1, device_id=0, rank=0 + ) + obj2 = QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid2, device_id=1, rank=1 + ) + obj3 = QuasisymmetryTwoTerm( + eq=eq, helicity=(1, eq.NFP), grid=grid3, device_id=2, rank=2 + ) + obj4 = AspectRatio(eq=eq, target=8, weight=100, device_id=0, rank=0) + objs = [obj1, obj2, obj3, obj4] + + objective = ObjectiveFunction(objs, deriv_mode="blocked", mpi=MPI) + + constraints = ( + ForceBalance(eq=eq), + FixBoundaryR(eq=eq, modes=R_modes), + FixBoundaryZ(eq=eq, modes=Z_modes), + FixPressure(eq=eq), + FixPsi(eq=eq), + FixCurrent(eq=eq), + ) + optimizer = Optimizer("proximal-lsq-exact") + + # every rank builds the objective, the constraints are not given a rank so they + # are computed on the root rank as usual + with run_with_mpi(objective, constraints, verbose=verbose) as is_root: + if is_root: + f0 = objective.compute_scalar(objective.x(eq)) + eq.optimize( + objective=objective, + constraints=constraints, + optimizer=optimizer, + maxiter=maxiter, + verbose=3, + ) + x1 = objective.x(eq) + f1 = objective.compute_scalar(x1) + assert f1 < f0 + + np.testing.assert_allclose(x1_no_mpi, x1, atol=1e-8, rtol=1e-8) + np.testing.assert_allclose(f1_no_mpi, f1, atol=1e-8, rtol=1e-8)