{"id":25736886,"url":"https://github.com/differentiableuniverseinitiative/jaxdecomp","last_synced_at":"2025-05-08T01:15:32.202Z","repository":{"id":63589448,"uuid":"564730551","full_name":"DifferentiableUniverseInitiative/jaxDecomp","owner":"DifferentiableUniverseInitiative","description":"JAX bindings for the NVIDIA cuDecomp library","archived":false,"fork":false,"pushed_at":"2025-05-05T16:11:02.000Z","size":34946,"stargazers_count":35,"open_issues_count":4,"forks_count":1,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-05-08T01:15:14.996Z","etag":null,"topics":["cuda","hpc","jax","xla"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/DifferentiableUniverseInitiative.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2022-11-11T11:12:24.000Z","updated_at":"2025-04-30T14:33:41.000Z","dependencies_parsed_at":"2024-03-30T00:23:07.904Z","dependency_job_id":"aa4ea9cf-fd94-4b3d-b589-aab53a67a36e","html_url":"https://github.com/DifferentiableUniverseInitiative/jaxDecomp","commit_stats":null,"previous_names":[],"tags_count":11,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DifferentiableUniverseInitiative%2FjaxDecomp","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DifferentiableUniverseInitiative%2FjaxDecomp/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DifferentiableUniverseInitiative%2FjaxDecomp/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DifferentiableUniverseInitiative%2FjaxDecomp/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/DifferentiableUniverseInitiative","download_url":"https://codeload.github.com/DifferentiableUniverseInitiative/jaxDecomp/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252978805,"owners_count":21834920,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cuda","hpc","jax","xla"],"created_at":"2025-02-26T06:22:51.242Z","updated_at":"2025-05-08T01:15:32.174Z","avatar_url":"https://github.com/DifferentiableUniverseInitiative.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\n# jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs\n\n[![Build](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/actions/workflows/github-deploy.yml/badge.svg)](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/actions/workflows/github-deploy.yml)\n[![Code Formatting](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/actions/workflows/formatting.yml/badge.svg)](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/actions/workflows/formatting.yml)\n[![Tests](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/actions/workflows/tests.yml/badge.svg)](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/actions/workflows/tests.yml)\n[![MIT License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)\n\n\u003e **Important**\n\u003e Version `0.2.0` includes a **pure JAX backend** that **no longer requires MPI**. For multi-node runs, MPI and NCCL backends are still available through **cuDecomp**.\n\nJAX reimplementation and bindings for NVIDIA's [cuDecomp](https://nvidia.github.io/cuDecomp/index.html) library [(Romero et al. 2022)](https://dl.acm.org/doi/abs/10.1145/3539781.3539797), enabling **multi-node parallel FFTs and halo exchanges** directly in low-level NCCL/CUDA-Aware MPI from your JAX code.\n\n---\n\n## Usage\n\nBelow is a simple code snippet illustrating how to perform a **3D FFT** on a distributed 3D array, followed by a halo exchange. For demonstration purposes, we force 8 CPU devices via environment variables:\n\n```python\nimport os\nos.environ[\"XLA_FLAGS\"] = \"--xla_force_host_platform_device_count=8\"\nos.environ[\"JAX_PLATFORM_NAME\"] = \"cpu\"\n\nimport jax\nfrom jax.sharding import Mesh, PartitionSpec as P, NamedSharding\nimport jaxdecomp\n\n# Create a 2x4 mesh of devices on CPU\npdims = (2, 4)\nmesh = jax.make_mesh(pdims, axis_names=('x', 'y'))\nsharding = NamedSharding(mesh, P('x', 'y'))\n\n# Create a random 3D array and enforce sharding\na = jax.random.normal(jax.random.PRNGKey(0), (1024, 1024, 1024))\na = jax.lax.with_sharding_constraint(a, sharding)\n\n# Parallel FFTs\nk_array = jaxdecomp.fft.pfft3d(a)\nrec_array = jaxdecomp.fft.pifft3d(a)\n\n# Parallel halo exchange\nexchanged = jaxdecomp.halo_exchange(a, halo_extents=(16, 16), halo_periods=(True, True))\n```\n\nAll these functions are **JIT**-compatible and support **automatic differentiation** (with [some caveats](docs/02-caveats.md)).\n\nSee also:\n- [Basic Usage](docs/01-basic_usage.md)\n- [Distributed LPT Example](examples/lpt_nbody_demo.py)\n\n\u003e **Important**\n\u003e Multi-node FFTs work with both JAX and cuDecomp backends\\\n\u003e For CPU with JAX, Multi-node is supported starting JAX v0.5.1 (with `gloo` backend)\n\n---\n\n## Running on an HPC Cluster\n\nOn HPC clusters (e.g., Jean Zay, Perlmutter), you typically launch your script with:\n```bash\nsrun python demo.py\n```\nor\n```bash\nmpirun -n 8 python demo.py\n```\n\nSee the Slurm [README](slurms/README.md) and [template script](slurms/template.slurm) for more details.\n\n\n---\n\n## Using cuDecomp (MPI and NCCL)\n\nFor other features, compile and install with cuDecomp enabled as described in [install](#2-jax--cudecomp-backend-advanced):\n\n```python\nimport jaxdecomp\n\n# Optionally select communication backends (defaults to NCCL)\njaxdecomp.config.update('halo_comm_backend', jaxdecomp.HALO_COMM_MPI)\njaxdecomp.config.update('transpose_comm_backend', jaxdecomp.TRANSPOSE_COMM_MPI_A2A)\n\n# Then specify 'backend=\"cudecomp\"' in your FFT or halo calls:\nkarray = jaxdecomp.fft.pfft3d(global_array, backend='cudecomp')\nrecarray = jaxdecomp.fft.pifft3d(karray, backend='cudecomp')\nexchanged_array = jaxdecomp.halo_exchange(\n    padded_array, halo_extents=(16, 16), halo_periods=(True, True), backend='cudecomp'\n)\n```\n\n## Install\n\n### 1. Pure JAX Version (Easy / Recommended)\n\n`jaxDecomp` is on PyPI:\n\n1. **Install the appropriate JAX wheel**:\n   - **GPU**:\n     ```bash\n     pip install --upgrade \"jax[cuda]\"\n     ```\n   - **CPU**:\n     ```bash\n     pip install --upgrade \"jax[cpu]\"\n     ```\n2. **Install `jaxdecomp`**:\n   ```bash\n   pip install jaxdecomp\n   ```\n\nThis setup uses the pure-JAX backend—**no** MPI required.\n\n### 2. JAX + cuDecomp Backend (Advanced)\n\nIf you need to use `MPI` instead of `NCCL` for `GPU` or gloo for CPU, you can build from GitHub with cuDecomp enabled. This requires the [NVIDIA HPC SDK](https://developer.nvidia.com/hpc-sdk) or a similar environment providing a CUDA-aware MPI toolchain.\n\n```bash\npip install -U pip\npip install git+https://github.com/DifferentiableUniverseInitiative/jaxDecomp -Ccmake.define.JD_CUDECOMP_BACKEND=ON\n```\n\n- If CMake cannot find NVHPC, set:\n  ```bash\n  export CMAKE_PREFIX_PATH=$CMAKE_PREFIX_PATH:$NVCOMPILERS/$NVARCH/22.9/cmake\n  ```\n  and then install again.\n\n---\n\n## Machine-Specific Notes\n\n### IDRIS [Jean Zay](http://www.idris.fr/eng/jean-zay/cpu/jean-zay-cpu-hw-eng.html) HPE SGI 8600 supercomputer\n\n\nAs of February 2025, loading modules **in this exact order** works:\n\n```bash\nmodule load nvidia-compilers/23.9 cuda/12.2.0 cudnn/8.9.7.29-cuda openmpi/4.1.5-cuda nccl/2.18.5-1-cuda cmake\n\n# Install JAX\npip install --upgrade \"jax[cuda]\"\n\n# Install jaxDecomp with cuDecomp\nexport CMAKE_PREFIX_PATH=$NVHPC_ROOT/cmake # sometimes needed\npip install git+https://github.com/DifferentiableUniverseInitiative/jaxDecomp -Ccmake.define.JD_CUDECOMP_BACKEND=ON\n```\n\n**Note**: If using only the pure-JAX backend, you do not need NVHPC.\n\n#### NERSC [Perlmutter](https://docs.nersc.gov/systems/perlmutter/architecture/) HPE Cray EX supercomputer\n\nAs of November 2022:\n\n```bash\nmodule load PrgEnv-nvhpc python\nexport CRAY_ACCEL_TARGET=nvidia80\n\n# Install JAX\npip install --upgrade \"jax[cuda]\"\n\n# Install jaxDecomp w/ cuDecomp\nexport CMAKE_PREFIX_PATH=/opt/nvidia/hpc_sdk/Linux_x86_64/22.5/cmake\npip install git+https://github.com/DifferentiableUniverseInitiative/jaxDecomp -CCmake.define.JD_CUDECOMP_BACKEND=ON\n```\n\n---\n\n## Backend Configuration (cuDecomp Only)\n\nBy default, cuDecomp uses NCCL for inter-device communication. You can customize this at runtime:\n\n```python\nimport jaxdecomp\n\n# Choose MPI or NVSHMEM for halo and transpose ops\njaxdecomp.config.update('transpose_comm_backend', jaxdecomp.TRANSPOSE_COMM_MPI_A2A)\njaxdecomp.config.update('halo_comm_backend', jaxdecomp.HALO_COMM_MPI)\n```\n\nThis can also be managed via environment variables, as described in the [docs](https://github.com/DifferentiableUniverseInitiative/jaxDecomp/tree/main/docs).\n\n---\n\n## Autotune Computational Mesh (cuDecomp Only)\n\nThe cuDecomp library can **autotune** the partition layout to maximize performance:\n\n```python\nautomesh = jaxdecomp.autotune(shape=[512,512,512])\n# 'automesh' is an optimized partition layout.\n# You can then create a JAX Sharding spec from this:\nfrom jax.sharding import PositionalSharding\nsharding = PositionalSharding(automesh)\n```\n\n---\n\n**License**: This project is licensed under the [MIT License](https://opensource.org/licenses/MIT).\n\nFor more details, see the [examples](examples/) directory and the [documentation](docs/). Contributions and issues are welcome!\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdifferentiableuniverseinitiative%2Fjaxdecomp","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdifferentiableuniverseinitiative%2Fjaxdecomp","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdifferentiableuniverseinitiative%2Fjaxdecomp/lists"}