Add multi-device sharding for NSS and SwiG

docs/guides/cli.md

@@ -354,6 +354,7 @@ n_tempered_steps = 5
 | `n_delete_frac` | `0.5` | Fraction of live points replaced per iteration |
 | `num_inner_steps_per_dim` | `10` | Slice steps per dimension for the nested kernel |
 | `termination_dlogz` | `0.1` | Stop when the remaining log-evidence contribution falls below this |
+| `n_devices` | `1` | Number of local devices used to shard live points; `n_live` and `n_delete` must be divisible by it |
 | `checkpoint_dir` | `{output.dir}/` | Directory for `checkpoint.pkl`; set by the CLI automatically |
 | `checkpoint_interval` | `600.0` | Seconds between checkpoint writes; `0` disables checkpointing |
 
@@ -369,6 +370,7 @@ n_tempered_steps = 5
 | `max_steps` | `10` | Maximum stepping-out expansions per slice |
 | `max_shrinkage` | `100` | Maximum shrinkage evaluations per slice |
 | `termination_dlogz` | `exp(-3)` (~`0.0498`) | Stop when the remaining evidence contribution falls below this |
+| `n_devices` | `1` | Number of local devices used to shard live points; `n_live` and `n_delete` must be divisible by it |
 | `checkpoint_dir` | `{output.dir}/` | Directory for `checkpoint.pkl`; set by the CLI automatically |
 | `checkpoint_interval` | `600.0` | Seconds between checkpoint writes; `0` disables checkpointing |
 

docs/guides/samplers.md

@@ -149,6 +149,7 @@ jim = Jim(
         n_delete_frac=0.5,
         num_inner_steps_per_dim=20,
         termination_dlogz=0.1,
+        n_devices=1,
     ),
     sample_transforms=sample_transforms,
     likelihood_transforms=likelihood_transforms,
@@ -162,6 +163,9 @@ Key parameters:
 - `n_live` — number of live points.
 - `n_delete_frac` — fraction of live points replaced per iteration.
 - `num_inner_steps_per_dim` — slice-sampler steps per dimension per dead point; increase for strongly correlated posteriors.
+- `n_devices` — number of local devices used to shard the live points and
+  replacement chains. Both `n_live` and `int(n_live * n_delete_frac)` must be
+  divisible by this value. The default `1` uses the unsharded kernel.
 
 **Repository:** [blackjax-devs/blackjax](https://github.com/blackjax-devs/blackjax)
 
@@ -202,6 +206,7 @@ jim = Jim(
         num_inner_steps_per_dim=1,
         num_gibbs_sweeps=2,
         termination_dlogz=math.exp(-3.0),
+        n_devices=1,
     ),
     likelihood_transforms=likelihood_transforms,
 )
@@ -218,6 +223,9 @@ Key parameters:
 - `n_live` / `n_delete_frac` — live-set size and replacement fraction, matching NSS.
 - `num_inner_steps_per_dim` — random-direction slice steps per block dimension.
 - `num_gibbs_sweeps` — complete block sweeps per particle replacement.
+- `n_devices` — number of local devices used to shard the live points and
+  replacement chains. The waveform cache remains local to each replacement
+  chain; only live-point state participates in collectives.
 
 ---
 

src/jimgw/samplers/blackjax/nss.py

@@ -16,9 +16,22 @@
 import blackjax
 from blackjax.ns.adaptive import AdaptiveNSState, init as _ns_adaptive_init
 from blackjax.ns.base import NSInfo, init_state_strategy as _init_state_strategy
-from blackjax.ns.nss import live_covariance, sample_direction_from_covariance
+from blackjax.ns.nss import (
+    live_covariance,
+    sample_direction_from_covariance,
+    slice_constrained_step,
+)
 from blackjax.ns.utils import finalise
+from blackjax.mcmc.slice import build_kernel as build_slice_kernel
+from blackjax.mcmc.slice import stepping_out
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 from jimgw.samplers.base import Sampler
+from jimgw.samplers.blackjax.sharding import (
+    build_sharded_from_mcmc_kernel,
+    make_live_mesh,
+    place_key,
+    place_state,
+)
 from jimgw.samplers.config import BlackJAXNSSConfig
 from jimgw.samplers.periodic import to_prior_space_proposal
 
@@ -84,9 +97,33 @@ def _checkpoint_tag(self) -> str:
     def _update_inner_kernel_params_fn(self) -> Callable:
         return live_covariance
 
-    def _build_nested_sampler(self, n_delete: int):
+    def _build_nested_sampler(self, n_delete: int, mesh: Mesh | None = None):
         config = self._config
         num_inner_steps = config.num_inner_steps_per_dim * self.n_dims
+        if mesh is not None:
+            init_state_fn = partial(
+                _init_state_strategy,
+                logprior_fn=self._log_prior_fn,
+                loglikelihood_fn=self._log_likelihood_fn,
+            )
+            slice_kernel = build_slice_kernel(
+                interval=stepping_out,
+                max_expansions=10,
+                max_shrinkage=100,
+            )
+            constrained_step = slice_constrained_step(
+                init_state_fn, slice_kernel, self._proposal
+            )
+            kernel = build_sharded_from_mcmc_kernel(
+                constrained_step,
+                n_inner_steps=num_inner_steps,
+                update_inner_kernel_params_fn=self._update_inner_kernel_params_fn,
+                n_delete=n_delete,
+                mesh=mesh,
+            )
+            return blackjax.SamplingAlgorithm(
+                lambda position, rng_key=None: position, kernel
+            )
         return blackjax.nss(
             logprior_fn=self._log_prior_fn,
             loglikelihood_fn=self._log_likelihood_fn,
@@ -120,6 +157,7 @@ def _sample(
         config = self._config
         n_live = config.n_live
         n_delete = int(n_live * config.n_delete_frac)
+        mesh = make_live_mesh(config.n_devices, n_live, n_delete)
         ckpt_path = (
             config.checkpoint_dir / "checkpoint.pkl"
             if config.checkpoint_dir is not None
@@ -137,7 +175,7 @@ def _validated_initial_particles(pos):
                 )
             return arr
 
-        nested_sampler = self._build_nested_sampler(n_delete)
+        nested_sampler = self._build_nested_sampler(n_delete, mesh)
 
         # Bypass BlackJAX's jax.vmap(init_state_fn) to avoid peak-memory OOM.
         # A full vmap over all live particles materialises O(n_live) concurrent
@@ -151,14 +189,18 @@ def _validated_initial_particles(pos):
         )
 
         def _batched_nss_init(positions):
+            if mesh is not None:
+                positions = jax.device_put(positions, NamedSharding(mesh, P("live")))
+
             def _batched_fn(pos):
                 return jax.lax.map(_single_init_fn, pos, batch_size=n_delete)
 
-            return _ns_adaptive_init(
+            state = _ns_adaptive_init(
                 positions,
                 init_state_fn=_batched_fn,
                 update_inner_kernel_params_fn=self._update_inner_kernel_params_fn,
             )
+            return place_state(state, mesh) if mesh is not None else state
 
         # Resume from checkpoint if one exists.
         if (
@@ -172,6 +214,9 @@ def _batched_fn(pos):
                 state = _ckpt["state"]
                 dead = _ckpt["dead"]
                 rng_key = _ckpt["rng_key"]
+                if mesh is not None:
+                    state = place_state(state, mesh)
+                    rng_key = place_key(rng_key, mesh)
                 n_iter = _ckpt["n_iter"]
                 self._prev_elapsed = float(_ckpt["elapsed_time"])
                 logger.info(
@@ -204,6 +249,9 @@ def _batched_fn(pos):
             dead = []
             n_iter = 0
 
+        if mesh is not None:
+            rng_key = place_key(rng_key, mesh)
+
         def _terminate(state: AdaptiveNSState) -> bool:
             dlogz = jnp.logaddexp(0, state.integrator.logZ_live - state.integrator.logZ)
             return bool(jnp.isfinite(dlogz) and dlogz < config.termination_dlogz)
@@ -223,17 +271,18 @@ def _terminate(state: AdaptiveNSState) -> bool:
             ):
                 _last_ckpt_t = config.write_checkpoint(
                     {
-                        "state": state,
-                        "dead": dead,
-                        "rng_key": rng_key,
+                        "state": jax.device_get(state),
+                        "dead": jax.device_get(dead),
+                        "rng_key": jax.device_get(rng_key),
                         "n_iter": n_iter,
                         "elapsed_time": self._prev_elapsed
                         + (time.perf_counter() - _method_t0),
                     },
                     self._checkpoint_tag,
                 )
 
-        self._final_state = finalise(state, dead)  # type: ignore[arg-type]  # AdaptiveNSState structurally satisfies NSState (.particles field)
+        final_state = finalise(state, dead)  # type: ignore[arg-type]  # AdaptiveNSState structurally satisfies NSState (.particles field)
+        self._final_state = jax.device_get(final_state)
         self._n_iterations = n_iter
 
         # Build anesthetic NestedSamples for use in get_samples() and get_diagnostics().

src/jimgw/samplers/blackjax/sharding.py

@@ -0,0 +1,177 @@
+"""Multi-device strategies for BlackJAX nested sampling."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from blackjax.ns.adaptive import build_kernel as build_adaptive_kernel
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
+
+_LIVE_AXIS = "live"
+
+
+def make_live_mesh(n_devices: int, n_live: int, n_delete: int) -> Mesh | None:
+    """Build a single-host device mesh for the live-particle axis."""
+    if n_devices == 1:
+        return None
+
+    devices = jax.local_devices()
+    if n_devices > len(devices):
+        raise ValueError(
+            f"n_devices={n_devices} requested, but only {len(devices)} local "
+            "JAX devices are available."
+        )
+    if n_live % n_devices:
+        raise ValueError(f"n_live={n_live} must be divisible by n_devices={n_devices}.")
+    if n_delete % n_devices:
+        raise ValueError(
+            f"n_delete={n_delete} must be divisible by n_devices={n_devices}."
+        )
+
+    return Mesh(np.asarray(devices[:n_devices]), axis_names=(_LIVE_AXIS,))
+
+
+def place_state(state, mesh: Mesh):
+    """Shard live particles and replicate adaptive/integrator state."""
+    live = NamedSharding(mesh, P(_LIVE_AXIS))
+    replicated = NamedSharding(mesh, P())
+    return state._replace(
+        particles=jax.tree.map(lambda x: jax.device_put(x, live), state.particles),
+        inner_kernel_params=jax.tree.map(
+            lambda x: jax.device_put(x, replicated), state.inner_kernel_params
+        ),
+        integrator=jax.tree.map(
+            lambda x: jax.device_put(x, replicated), state.integrator
+        ),
+    )
+
+
+def place_key(rng_key, mesh: Mesh):
+    """Replicate a sampler PRNG key over the mesh."""
+    return jax.device_put(rng_key, NamedSharding(mesh, P()))
+
+
+def _all_gather(mesh: Mesh, value, in_spec):
+    def gather(x):
+        return jax.lax.all_gather(x, _LIVE_AXIS, tiled=True)
+
+    return jax.shard_map(
+        gather,
+        mesh=mesh,
+        in_specs=in_spec,
+        out_specs=P(),
+        check_vma=False,
+    )(value)
+
+
+def delete_fn_sharded(mesh: Mesh, n_delete: int) -> Callable:
+    """Select the globally lowest-likelihood live particles."""
+
+    def delete_fn(state):
+        global_log_likelihood = _all_gather(
+            mesh, state.particles.loglikelihood, P(_LIVE_AXIS)
+        )
+        _, dead_idx = jax.lax.top_k(-global_log_likelihood, n_delete)
+        return dead_idx, dead_idx
+
+    return delete_fn
+
+
+def update_with_mcmc_take_last_sharded(
+    constrained_step_fn: Callable,
+    n_inner_steps: int,
+    n_delete: int,
+    mesh: Mesh,
+) -> Callable:
+    """Run independent constrained replacement chains across the mesh."""
+    live = NamedSharding(mesh, P(_LIVE_AXIS))
+
+    def update_function(rng_key, state, loglikelihood_0, **step_parameters):
+        particles = state.particles
+        particle_specs = jax.tree.map(
+            lambda x: P(_LIVE_AXIS) if x.ndim > 0 else P(), particles
+        )
+        global_particles = jax.shard_map(
+            lambda xs: jax.tree.map(
+                lambda x: jax.lax.all_gather(x, _LIVE_AXIS, tiled=True), xs
+            ),
+            mesh=mesh,
+            in_specs=(particle_specs,),
+            out_specs=jax.tree.map(lambda _: P(), particles),
+            check_vma=False,
+        )(particles)
+        global_log_likelihood = global_particles.loglikelihood
+
+        survivor_weights = (global_log_likelihood > loglikelihood_0).astype(jnp.float32)
+        survivor_weights = jnp.where(
+            survivor_weights.sum() > 0,
+            survivor_weights,
+            jnp.ones_like(survivor_weights),
+        )
+
+        choice_key, sample_key = jax.random.split(rng_key)
+        start_idx = jax.random.choice(
+            choice_key,
+            global_log_likelihood.shape[0],
+            shape=(n_delete,),
+            p=survivor_weights / survivor_weights.sum(),
+            replace=True,
+        )
+        start_states = jax.tree.map(lambda x: x[start_idx], global_particles)
+        start_states = jax.tree.map(lambda x: jax.device_put(x, live), start_states)
+        sample_keys = jax.device_put(jax.random.split(sample_key, n_delete), live)
+
+        state_specs = jax.tree.map(
+            lambda x: P(_LIVE_AXIS) if x.ndim > 0 else P(), start_states
+        )
+        parameter_specs = jax.tree.map(lambda _: P(), step_parameters)
+
+        def run_local_chains(keys, states, threshold, parameters):
+            def run_chain(key, chain_state):
+                keys = jax.random.split(key, n_inner_steps)
+
+                def body_fn(current_state, step_key):
+                    return constrained_step_fn(
+                        step_key,
+                        current_state,
+                        threshold,
+                        **parameters,
+                    )
+
+                return jax.lax.scan(body_fn, chain_state, keys)
+
+            return jax.vmap(run_chain)(keys, states)
+
+        return jax.shard_map(
+            run_local_chains,
+            mesh=mesh,
+            in_specs=(P(_LIVE_AXIS), state_specs, P(), parameter_specs),
+            out_specs=(state_specs, P(_LIVE_AXIS)),
+            check_vma=False,
+        )(sample_keys, start_states, loglikelihood_0, step_parameters)
+
+    return update_function
+
+
+def build_sharded_from_mcmc_kernel(
+    constrained_step_fn: Callable,
+    n_inner_steps: int,
+    update_inner_kernel_params_fn: Callable,
+    n_delete: int,
+    mesh: Mesh,
+) -> Callable:
+    """Assemble BlackJAX's adaptive NS kernel with sharded strategies."""
+    inner_kernel = update_with_mcmc_take_last_sharded(
+        constrained_step_fn,
+        n_inner_steps=n_inner_steps,
+        n_delete=n_delete,
+        mesh=mesh,
+    )
+    return build_adaptive_kernel(
+        delete_fn_sharded(mesh, n_delete),
+        inner_kernel,
+        update_inner_kernel_params_fn=update_inner_kernel_params_fn,
+    )