pytc.vmc.sharding¶

Multi-GPU sharding utilities for VMC.

This module provides helpers for distributing VMC walkers across multiple devices (GPUs/TPUs) using JAX’s built-in sharding.

The strategy is data-parallel over walkers: each device holds N/D walkers and performs per-walker computations independently. Only lightweight reductions (mean energy, J^T @ δE, J^T @ J) trigger cross-device communication, which JAX inserts automatically.

Usage¶

from pytc.vmc.sharding import create_mesh, shard_walker, get_sharding

mesh = create_mesh()                    # auto-detect devices
walkers = shard_walker(walkers, mesh)   # shard along walker axis
params  = replicate(params, mesh)       # replicate on all devices

After sharding, the existing jax.vmap and jax.jit code works unchanged — JAX’s SPMD compiler infers the communication.

Notes

folx.batched_vmap does not preserve sharding (it gathers results to all devices). When multi_gpu=True the code should fall back to jax.vmap so that per-walker outputs stay distributed.
Walkers must have n_walkers divisible by the number of devices. pad_walkers can add padding walkers to satisfy this.

Functions

`_assemble_sharded_from_local_pytrees`	Assemble per-device local pytrees into a globally sharded pytree.
`_slice_along_first_axis`	Slice a pytree along the leading axis for all non-scalar leaves.
`create_mesh`	Create a 1-D device mesh for walker-parallel sharding.
`get_replicated_sharding`	NamedSharding that replicates data on every device.
`get_vmap_fn`	Return the appropriate vmap implementation.
`get_walker_sharding`	NamedSharding that partitions the leading (walker) dimension.
`initialize_walkers_sharded`	Initialize walkers independently per device and return globally sharded walkers.
`is_multi_gpu`	True if more than one local device is available.
`n_devices`	Number of local devices available to this process.
`pad_n_walkers`	Return the smallest multiple of n_devices >= n_walkers.
`pad_walker`	Pad a Walker pytree along axis 0 to target_n_walkers.
`replicate`	Replicate a pytree on every device.
`shard_map_wrap`	Apply `shard_map` with compatibility flags across JAX versions.
`shard_vmap`	A version of jax.vmap that works across multiple devices using shard_map.
`shard_walker`	Place a Walker (or any pytree) with sharding along axis 0.
`sharded_batched_vmap`	A version of folx.batched_vmap that works across multiple devices.