diff --git a/src/fdtdx/core/jax/sharding.py b/src/fdtdx/core/jax/sharding.py index 6772835b..3c343cfc 100644 --- a/src/fdtdx/core/jax/sharding.py +++ b/src/fdtdx/core/jax/sharding.py @@ -41,6 +41,60 @@ def get_named_sharding_from_shape( counter = 0 +def _replicated_mesh_sharding(compute_devices: list[jax.Device]) -> jax.sharding.NamedSharding: + """Build a REPLICATED NamedSharding over ``compute_devices`` -- NOT SingleDeviceSharding, whose + ``{device 0}`` device set is incompatible with the ``[0,1,...]`` set of the surrounding x-sharded + computation (jit raises "Received incompatible devices for jitted computation"). A replicated + PartitionSpec() keeps data on the full mesh but UN-sharded, so a subsequent placement (see + ``pin_to_single_device``) gathers only the small operand, never the full x-sharded domain, while + staying device-set-compatible with the parent computation. + + Split out of ``pin_to_single_device`` so the (cheap, pure) mesh/sharding construction can be unit- + tested directly with a MOCKED device list (including a duplicated physical device, which the actual + data placement in ``pin_to_single_device`` cannot accept -- see ``jax.device_put``'s "distinct + devices" requirement, exercised only in the subprocess-based multi-device tests). + """ + devices = mesh_utils.create_device_mesh((len(compute_devices),), devices=compute_devices) + mesh = jax.sharding.Mesh(devices=devices, axis_names=(SHARD_STR,)) + return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + + +def pin_to_single_device(arr: jax.Array) -> jax.Array: + """Pin ``arr`` to a single-device replicated layout via ``jax.device_put``. + + On a multi-device run the field/material arrays are x-sharded (NamedSharding + ``P(None,'shard',None,None)``). When a thin transverse plane is sliced out of such an + array and fed to a host ``pure_callback`` (the Tidy3D mode solver, which runs at + ``{maximal device=0}``), XLA must materialise a single-device operand for the callback. + Without an explicit constraint XLA may conservatively all-gather the *parent* array to + satisfy the callback's single-device requirement. Pinning the already-sliced PLANE to a + single-device replicated sharding tells XLA to gather only the thin band (a few-hundred-KB + plane), never the full ~100M-cell domain. + + Also used on the pure_callback's OWN outputs (see ``compute_mode``): those carry a maximal/ + single-device GSPMDSharding that ``jax.lax.with_sharding_constraint`` cannot convert (raises + "Cannot convert GSPMDSharding {maximal device=0} into SdyArray" via XLA's Shardy partitioner). + ``jax.device_put`` performs the placement/reshard directly and accepts both this maximal-sharded + case and the ordinary already-NamedSharding case (the sliced input this function was originally + written for), so one implementation covers both call patterns. + + No-op on a single-device run (``len(jax.devices()) == 1``): returns ``arr`` unchanged so the + single-GPU path stays byte-identical (no sharding constraint is inserted at all). + """ + if len(jax.devices()) == 1: + return arr + replicated = _replicated_mesh_sharding(jax.devices()) + # device_put, NOT jax.lax.with_sharding_constraint: the latter internally converts through XLA's + # Shardy partitioner (`api.jit(_identity_fn, out_shardings=sharding)`), which raises "Cannot convert + # GSPMDSharding {maximal device=0} into SdyArray" when `arr` is the RAW OUTPUT of a jax.pure_callback + # (pure_callback results carry a maximal/single-device GSPMDSharding, not an ordinary NamedSharding -- + # a fundamentally different representation Shardy's conversion path does not handle). device_put + # performs the actual data placement/reshard directly and accepts a maximal-sharded source array, + # so it works for BOTH the already-well-formed NamedSharding case (the sliced x-sharded input this + # function was originally written for) and the pure_callback-output case (its raw mode-solver result). + return jax.device_put(arr, replicated) + + def get_dtype_bytes(dtype: jnp.dtype) -> int: return jnp.dtype(dtype).itemsize diff --git a/src/fdtdx/core/physics/modes.py b/src/fdtdx/core/physics/modes.py index 2f2ca3b3..d406b148 100644 --- a/src/fdtdx/core/physics/modes.py +++ b/src/fdtdx/core/physics/modes.py @@ -10,6 +10,7 @@ from tidy3d.components.mode.solver import compute_modes as _compute_modes from fdtdx.core.axis import get_transverse_axes +from fdtdx.core.jax.sharding import pin_to_single_device from fdtdx.core.misc import expand_to_3x3 from fdtdx.core.physics.metrics import normalize_by_poynting_flux @@ -361,6 +362,14 @@ def mode_helper(permittivity, permeability, c0_um, c1_um): jax.lax.stop_gradient(c0_um), jax.lax.stop_gradient(c1_um), ) + # pure_callback executes on a single host device ({maximal device=0}); its OUTPUT, like the sliced + # INPUT above, needs an explicit sharding reconciliation before further ops run it through the + # surrounding x-sharded multi-device computation -- otherwise XLA/Shardy cannot convert the callback's + # single-device GSPMDSharding into an SdyArray at the next op (observed: jnp.expand_dims below raises + # "Cannot convert GSPMDSharding {maximal device=0} into SdyArray" on >1 device). + mode_E_raw = pin_to_single_device(mode_E_raw) + mode_H_raw = pin_to_single_device(mode_H_raw) + eff_idx = pin_to_single_device(eff_idx) # third callback output -- same fix, easy to miss mode_E = jnp.expand_dims(mode_E_raw, axis=propagation_axis + 1) mode_H = jnp.expand_dims(mode_H_raw, axis=propagation_axis + 1) diff --git a/src/fdtdx/objects/detectors/mode.py b/src/fdtdx/objects/detectors/mode.py index 73734391..6af70f0e 100644 --- a/src/fdtdx/objects/detectors/mode.py +++ b/src/fdtdx/objects/detectors/mode.py @@ -7,6 +7,7 @@ from fdtdx import constants from fdtdx.config import SimulationConfig from fdtdx.core.jax.pytrees import autoinit, frozen_field, private_field +from fdtdx.core.jax.sharding import pin_to_single_device from fdtdx.core.null import Null from fdtdx.core.physics.modes import compute_mode from fdtdx.dispersion import effective_complex_inv_permittivity @@ -156,9 +157,10 @@ def apply( dispersive_c4: jax.Array | None = None, ) -> Self: del key - inv_permittivity_slice = inv_permittivities[:, *self.grid_slice] + # Pin the sliced plane to a single device (see pin_to_single_device docstring for why). + inv_permittivity_slice = pin_to_single_device(inv_permittivities[:, *self.grid_slice]) if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0: - inv_permeability_slice = inv_permeabilities[:, *self.grid_slice] + inv_permeability_slice = pin_to_single_device(inv_permeabilities[:, *self.grid_slice]) else: inv_permeability_slice = inv_permeabilities diff --git a/src/fdtdx/objects/sources/mode.py b/src/fdtdx/objects/sources/mode.py index 29ca026f..b4569a4a 100644 --- a/src/fdtdx/objects/sources/mode.py +++ b/src/fdtdx/objects/sources/mode.py @@ -10,6 +10,7 @@ from fdtdx.core.axis import get_transverse_axes from fdtdx.core.grid import calculate_time_offset_yee from fdtdx.core.jax.pytrees import autoinit, frozen_field, private_field +from fdtdx.core.jax.sharding import pin_to_single_device from fdtdx.core.linalg import get_wave_vector_raw from fdtdx.core.physics.metrics import compute_energy from fdtdx.core.physics.modes import compute_mode @@ -112,10 +113,11 @@ def apply( raise NotImplementedError() # inv_permittivities shape: (3, Nx, Ny, Nz) - slice with component dimension - inv_permittivity_slice = inv_permittivities[:, *self.grid_slice] + # Pin the sliced plane to a single device (see pin_to_single_device docstring for why). + inv_permittivity_slice = pin_to_single_device(inv_permittivities[:, *self.grid_slice]) if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0: # inv_permeabilities shape: (3, Nx, Ny, Nz) - slice with component dimension - inv_permeability_slice = inv_permeabilities[:, *self.grid_slice] + inv_permeability_slice = pin_to_single_device(inv_permeabilities[:, *self.grid_slice]) else: inv_permeability_slice = inv_permeabilities diff --git a/tests/unit/core/jax/test_sharding.py b/tests/unit/core/jax/test_sharding.py index 734c31c4..5144a688 100644 --- a/tests/unit/core/jax/test_sharding.py +++ b/tests/unit/core/jax/test_sharding.py @@ -1,5 +1,9 @@ """Unit tests for fdtdx.core.jax.sharding module.""" +import os +import subprocess +import sys +import textwrap from unittest import mock from unittest.mock import MagicMock @@ -10,9 +14,11 @@ import fdtdx.core.jax.sharding as sharding_module from fdtdx.constants import SHARD_STR from fdtdx.core.jax.sharding import ( + _replicated_mesh_sharding, create_named_sharded_matrix, get_dtype_bytes, get_named_sharding_from_shape, + pin_to_single_device, pretty_print_sharding, ) @@ -198,3 +204,184 @@ def test_sharding_axis_fallback_when_dim_is_one(self): def test_has_named_sharding(self): result = create_named_sharded_matrix(shape=(4, 6), value=1.0, sharding_axis=0, dtype=jnp.float32, backend="cpu") assert isinstance(result.sharding, jax.sharding.NamedSharding) + + +# ---- _replicated_mesh_sharding ---- + + +class TestReplicatedMeshSharding: + """Tests for the _replicated_mesh_sharding helper, split out of pin_to_single_device so this + (pure, cheap) mesh/sharding construction can be covered in-process -- unlike pin_to_single_device's + own data placement, it does not require genuinely distinct physical devices, so a MOCKED duplicated + device list is a faithful, subprocess-free test (the corresponding data placement is exercised + separately by TestPinToSingleDevice's subprocess-based multi-device tests below).""" + + def test_single_device_replicated_spec(self): + cpu = CPU_DEVICES[0] + result = _replicated_mesh_sharding([cpu]) + assert isinstance(result, jax.sharding.NamedSharding) + assert tuple(result.spec) == () + assert result.mesh.devices.shape == (1,) + + def test_duplicated_device_list_replicated_spec(self): + # A duplicated physical device (simulating >1 "device" without needing real distinct hardware) + # is accepted by mesh/sharding CONSTRUCTION -- only the actual data placement in + # pin_to_single_device requires genuine distinctness (see module docstring above). + cpu = CPU_DEVICES[0] + result = _replicated_mesh_sharding([cpu, cpu]) + assert isinstance(result, jax.sharding.NamedSharding) + assert tuple(result.spec) == () + assert result.mesh.devices.shape == (2,) + assert SHARD_STR in result.mesh.axis_names + + def test_result_usable_as_pin_to_single_device_input(self): + """pin_to_single_device calls this helper directly with jax.devices() -- confirm the two agree + on shape/axis-name conventions for a real (single, non-mocked) device list.""" + expected = _replicated_mesh_sharding(jax.devices()) + assert isinstance(expected, jax.sharding.NamedSharding) + assert expected.mesh.devices.shape == (len(jax.devices()),) + + +# ---- pin_to_single_device ---- + + +def _run_in_subprocess_with_devices(num_devices: int, body: str) -> None: + """Run ``body`` in a fresh interpreter with ``num_devices`` virtual CPU devices. + + ``jax.devices()`` reflects the platform's device count only at JAX's first backend + initialisation, which has already happened (with a single forced CPU device, see + ``tests/conftest.py``) by the time this test module runs. A real multi-device mesh can + therefore only be exercised in a separate process started with + ``XLA_FLAGS=--xla_force_host_platform_device_count=N`` set beforehand. Mocking + ``jax.devices()`` to return the same physical device twice is not a substitute here: + unlike the mesh-shape/spec checks elsewhere in this file, ``pin_to_single_device`` places + data onto a real device mesh eagerly, which requires genuinely distinct physical devices + (a mesh built from a duplicated device raises at construction/placement time). + """ + env = { + **os.environ, + "JAX_PLATFORMS": "cpu", + "XLA_FLAGS": f"--xla_force_host_platform_device_count={num_devices}", + } + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(body)], + capture_output=True, + text=True, + env=env, + timeout=60, + ) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +class TestPinToSingleDevice: + """Tests for the pin_to_single_device helper.""" + + def test_single_device_is_noop_identity(self): + # The default test session forces a single CPU device (see tests/conftest.py). + assert len(jax.devices()) == 1 + arr = jnp.arange(12.0).reshape(3, 4) + result = pin_to_single_device(arr) + assert result is arr + + def test_multi_device_preserves_shape_dtype_and_values(self): + _run_in_subprocess_with_devices( + 2, + """ + import jax, jax.numpy as jnp + from fdtdx.core.jax.sharding import pin_to_single_device + + assert len(jax.devices()) == 2 + arr = jnp.arange(24.0, dtype=jnp.float32).reshape(4, 6) + out = pin_to_single_device(arr) + assert out.shape == arr.shape + assert out.dtype == arr.dtype + assert bool(jnp.array_equal(out, arr)) + """, + ) + + def test_multi_device_result_is_replicated_named_sharding(self): + _run_in_subprocess_with_devices( + 2, + """ + import jax, jax.numpy as jnp + from fdtdx.core.jax.sharding import pin_to_single_device + + out = pin_to_single_device(jnp.zeros((4, 6), dtype=jnp.float32)) + assert isinstance(out.sharding, jax.sharding.NamedSharding) + # Fully replicated: an empty PartitionSpec, not sharded along any axis. + assert tuple(out.sharding.spec) == () + assert "shard" in out.sharding.mesh.axis_names + """, + ) + + def test_multi_device_compatible_with_sliced_x_sharded_parent(self): + """Regression test for the exact usage in ModePlaneSource / ModeOverlapDetector: + pinning a plane sliced out of an x-sharded array inside a jit must not raise JAX's + "Received incompatible devices for jitted computation" error, which is what happens + if the sliced plane were instead placed on a plain SingleDeviceSharding (device 0) + while the surrounding computation is compiled over the full device mesh. + """ + _run_in_subprocess_with_devices( + 2, + """ + import jax, jax.numpy as jnp + from fdtdx.core.jax.sharding import create_named_sharded_matrix, pin_to_single_device + + # (component, Nx, Ny, Nz) permittivity-like array, x-sharded on axis=1. + arr = create_named_sharded_matrix( + (3, 8, 6, 6), value=2.0, dtype=jnp.float32, sharding_axis=1, backend="cpu" + ) + + @jax.jit + def slice_and_pin(a): + plane = a[:, 2:3, :, :] + return pin_to_single_device(plane) + + out = slice_and_pin(arr) + assert out.shape == (3, 1, 6, 6) + assert bool(jnp.allclose(out, 2.0)) + """, + ) + + def test_multi_device_compatible_with_pure_callback_output(self): + """Regression test for compute_mode's usage: pinning the RAW OUTPUT of a + jax.pure_callback (as returned by the Tidy3D mode solver) must not raise "Cannot + convert GSPMDSharding {maximal device=0} into SdyArray". A pure_callback's result + carries a maximal/single-device GSPMDSharding -- a fundamentally different + representation than the ordinary NamedSharding case covered by the sliced-input + test above, which XLA's Shardy partitioner cannot convert via + jax.lax.with_sharding_constraint (this failed before this fix; device_put succeeds + because it performs the placement directly instead of a Shardy-mediated re-shard). + """ + _run_in_subprocess_with_devices( + 2, + """ + import numpy as np + import jax, jax.numpy as jnp + from fdtdx.core.jax.sharding import create_named_sharded_matrix, pin_to_single_device + + arr = create_named_sharded_matrix( + (3, 8, 6, 6), value=2.0, dtype=jnp.float32, sharding_axis=1, backend="cpu" + ) + + def host_fn(x): + # Stand-in for the Tidy3D mode-solver call in compute_mode: an arbitrary + # host-side computation returning a NEW array (not just a passthrough). + return np.asarray(x) * 3.0 + + @jax.jit + def slice_callback_and_pin(a): + plane = a[:, 2:3, :, :] + raw = jax.pure_callback( + host_fn, jax.ShapeDtypeStruct(plane.shape, plane.dtype), plane + ) + pinned = pin_to_single_device(raw) + # Mirrors compute_mode's jnp.expand_dims(mode_E_raw, ...) right after pinning -- + # the exact next op that raised before this fix. + return jnp.expand_dims(pinned, axis=0) + + out = slice_callback_and_pin(arr) + assert out.shape == (1, 3, 1, 6, 6) + assert bool(jnp.allclose(out, 6.0)) + """, + )