From bed4ac28457a0e18047a35deeebdd48bb9243b82 Mon Sep 17 00:00:00 2001 From: justinlu Date: Wed, 26 Aug 2026 22:12:51 -0700 Subject: [PATCH] Unify JAX WeightSynchronizer API across Pathways FFI and native backends PiperOrigin-RevId: 971713939 --- tpu_sync/api/jax/BUILD | 2 + tpu_sync/api/jax/weight_synchronizer.py | 83 +++++-- tpu_sync/api/jax/weight_synchronizer_test.py | 27 +++ .../frameworks/jax/tpu_raiden_jax_module.cc | 4 +- .../frameworks/jax/weight_synchronizer_ffi.py | 220 ++++++++++++++++++ 5 files changed, 319 insertions(+), 17 deletions(-) diff --git a/tpu_sync/api/jax/BUILD b/tpu_sync/api/jax/BUILD index fb35e74d..361f2480 100644 --- a/tpu_sync/api/jax/BUILD +++ b/tpu_sync/api/jax/BUILD @@ -83,6 +83,8 @@ py_library( deps = [ "//tpu_sync/frameworks/jax:_tpu_raiden_jax", "//tpu_sync/frameworks/jax:jax_test_utils", + "//tpu_sync/frameworks/jax:weight_synchronizer_ffi_py", + "@jax//jax", ], ) diff --git a/tpu_sync/api/jax/weight_synchronizer.py b/tpu_sync/api/jax/weight_synchronizer.py index 7b23d4c7..9561f477 100644 --- a/tpu_sync/api/jax/weight_synchronizer.py +++ b/tpu_sync/api/jax/weight_synchronizer.py @@ -16,8 +16,27 @@ from typing import Any, Dict, List, Optional -# Import Nanobind binary library directly E2E! +import jax + from tpu_sync.frameworks.jax import _tpu_raiden_jax as _weight_synchronizer +from tpu_sync.frameworks.jax import weight_synchronizer_ffi as _weight_synchronizer_ffi + + +def is_pathways_backend() -> bool: + """Returns True if the current JAX environment is targeting Pathways.""" + try: + if jax.config.read("jax_platforms") == "pathways": + return True + devices = jax.devices() + if ( + devices + and hasattr(devices[0], "client") + and hasattr(devices[0].client, "runtime_type") + ): + return "pathways" in str(devices[0].client.runtime_type).lower() + except Exception: + pass + return False class WeightSynchronizer: @@ -32,6 +51,7 @@ def __init__( listener_port: Optional[int] = None, bind_ip: Optional[str] = None, auto_h2d: bool = False, + backend: Optional[str] = None, ): """Instantiates the Weight Synchronizer on a JAX weights list. @@ -43,31 +63,50 @@ def __init__( listener_port: Sockets server port for incoming C++ Listener commands. bind_ip: Sockets server bind IP address. auto_h2d: Automatically execute H2D ingestion upon data arrival. + backend: Explicit backend selection ('pathways' or 'pjrt'/'default'). If + None, automatically detected based on the active JAX runtime. """ - self._impl = _weight_synchronizer.WeightSynchronizer( - jax_arrays, - local_port, - parallelism, - unsafe_skip_buffer_lock, - listener_port, - bind_ip, - auto_h2d, + use_ffi = (backend == "pathways") or ( + backend is None and is_pathways_backend() ) + if use_ffi: + self._impl = _weight_synchronizer_ffi.WeightSynchronizer( + jax_arrays=jax_arrays, + local_port=local_port, + parallelism=parallelism, + unsafe_skip_buffer_lock=unsafe_skip_buffer_lock, + listener_port=listener_port, + bind_ip=bind_ip, + auto_h2d=auto_h2d, + ) + else: + self._impl = _weight_synchronizer.WeightSynchronizer( + jax_arrays, + local_port, + parallelism, + unsafe_skip_buffer_lock, + listener_port, + bind_ip, + auto_h2d, + ) def d2h(self) -> None: """Triggers asynchronous Device-to-Host (D2H) copy of current weights to Host buffer.""" - self._impl.D2h() + self._impl.d2h() def h2d(self) -> None: """Triggers asynchronous Host-to-Device (H2D) copy of staged host buffer back to Device memory E2E.""" - self._impl.H2d() + self._impl.h2d() def test_only_set_skip_tiling(self, skip: bool | List[bool]) -> None: """Sets whether D2H/H2D should skip CPU tiling/detiling (for testing only).""" - if isinstance(skip, bool): - self._impl.set_skip_tiling(skip) - else: - self._impl.set_skip_tiling(list(skip)) + if hasattr(self._impl, "set_skip_tiling"): + if isinstance(skip, bool): + self._impl.set_skip_tiling(skip) + else: + self._impl.set_skip_tiling(list(skip)) + elif hasattr(self._impl, "test_only_set_skip_tiling"): + self._impl.test_only_set_skip_tiling(skip) def bind_weights(self, jax_arrays: List[any]) -> None: """Binds the JAX arrays to the weight synchronizer in-place. @@ -178,3 +217,17 @@ def get_metrics(self) -> dict[str, float | int]: def reset_metrics(self) -> None: """Resets all recorded internal metrics.""" self._impl.reset_metrics() + + def close(self) -> None: + """Closes and tears down internal buffers and servers.""" + if hasattr(self._impl, "close"): + self._impl.close() + elif hasattr(self._impl, "destroy"): + self._impl.destroy() + + def destroy(self) -> None: + """Destroys internal buffers and servers.""" + if hasattr(self._impl, "destroy"): + self._impl.destroy() + elif hasattr(self._impl, "close"): + self._impl.close() diff --git a/tpu_sync/api/jax/weight_synchronizer_test.py b/tpu_sync/api/jax/weight_synchronizer_test.py index a3c114b5..af449c44 100644 --- a/tpu_sync/api/jax/weight_synchronizer_test.py +++ b/tpu_sync/api/jax/weight_synchronizer_test.py @@ -288,6 +288,33 @@ def test_push_sync_aligned_to_aligned(self): ) self._run_resharding_test(src_sharding, dst_sharding, (8, 8)) + def test_backend_selection(self): + arrs = [ + jax.device_put(jnp.ones(self.shape, dtype=self.dtype), self.sharding) + ] + # Default backend (non-FFI) + ws_default = WeightSynchronizer( + jax_arrays=arrs, + local_port=0, + unsafe_skip_buffer_lock=True, + ) + self.assertIsNotNone(ws_default.local_port) + self.assertEqual(ws_default.num_layers, 1) + + # Pathways backend (FFI) + ws_pathways = WeightSynchronizer( + jax_arrays=arrs, + local_port=0, + unsafe_skip_buffer_lock=True, + backend="pathways", + ) + self.assertIsNotNone(ws_pathways.local_port) + self.assertEqual(ws_pathways.num_layers, 1) + self.assertIsInstance(ws_pathways.get_metrics(), dict) + ws_pathways.d2h() + ws_pathways.h2d() + ws_pathways.close() + class ShardSortingUtilTest(absltest.TestCase): def setUp(self): diff --git a/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc index 1238de94..27c17784 100644 --- a/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc +++ b/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc @@ -309,7 +309,7 @@ NB_MODULE(_tpu_raiden_jax, m) { nb::arg("bind_ip") = nb::none(), nb::arg("auto_h2d") = false) .def( - "D2h", + "d2h", [](WeightSynchronizer& self) { auto status_or_future = self.D2h(); if (!status_or_future.ok()) { @@ -325,7 +325,7 @@ NB_MODULE(_tpu_raiden_jax, m) { }, nb::call_guard()) .def( - "H2d", + "h2d", [](WeightSynchronizer& self) { auto status_or_future = self.H2d(); if (!status_or_future.ok()) { diff --git a/tpu_sync/frameworks/jax/weight_synchronizer_ffi.py b/tpu_sync/frameworks/jax/weight_synchronizer_ffi.py index 97170758..edd9e488 100644 --- a/tpu_sync/frameworks/jax/weight_synchronizer_ffi.py +++ b/tpu_sync/frameworks/jax/weight_synchronizer_ffi.py @@ -14,6 +14,9 @@ """JAX bindings for WeightSynchronizer FFI, enabling host/device weight synchronization.""" +from collections.abc import Sequence +from typing import Any, Dict, List, Optional + import jax from jax.experimental import compute_on import jax.numpy as jnp @@ -21,6 +24,8 @@ from tpu_sync.frameworks.jax import _weight_synchronizer_ffi +__all__ = ["WeightSynchronizer"] + def init_weight_synchronizer( device_array, @@ -309,3 +314,218 @@ def _local_d2h(anchor, s_idx): in_specs=(anchor_spec, index_spec), out_specs=anchor_spec, )(device_array, shard_idx) + + +class WeightSynchronizer: + """FFI-based distributed Weight Synchronizer for JAX / Pathways.""" + + def __init__( + self, + jax_arrays: Sequence[Any], + local_port: Optional[int] = None, + parallelism: int = 1, + unsafe_skip_buffer_lock: bool = False, + listener_port: Optional[int] = None, + bind_ip: Optional[str] = None, + auto_h2d: bool = False, + mesh: Optional[jax.sharding.Mesh] = None, + num_shards: Optional[int] = None, + ): + """Instantiates the FFI-based Weight Synchronizer on a JAX weights list. + + Args: + jax_arrays: A sequence of JAX arrays representing the sharded model + weights. + local_port: Sockets server port for incoming pulls (inference mode). + parallelism: Number of parallel network stream TCP sockets workers. + unsafe_skip_buffer_lock: Skip PJRT buffer locks during weights unpack. + listener_port: Sockets server port for incoming C++ Listener commands. + bind_ip: Sockets server bind IP address. + auto_h2d: Automatically execute H2D ingestion upon data arrival. + mesh: Optional JAX device mesh. If omitted, inferred from jax_arrays. + num_shards: Number of local shards per host. If omitted, inferred from + mesh. + """ + if not jax_arrays: + raise ValueError("jax_arrays list cannot be empty") + self._jax_arrays = list(jax_arrays) + self._parallelism = parallelism + self._unsafe_skip_buffer_lock = unsafe_skip_buffer_lock + self._bind_ip = bind_ip + self._auto_h2d = auto_h2d + self._destroyed = False + + # 1. Resolve mesh + if mesh is not None: + self._mesh = mesh + else: + first_sharding = self._jax_arrays[0].sharding + if hasattr(first_sharding, "mesh") and first_sharding.mesh is not None: + self._mesh = first_sharding.mesh + else: + devices = jax.devices() + self._mesh = jax.sharding.Mesh(np.array(devices), ("devices",)) + + # 2. Compute slice byte sizes + self._slice_byte_sizes = [ + int(np.prod(arr.sharding.shard_shape(arr.shape)) * arr.dtype.itemsize) + for arr in self._jax_arrays + ] + sizes_sharding = jax.sharding.NamedSharding( + self._mesh, jax.sharding.PartitionSpec(None) + ) + self._slice_byte_sizes_sharded = jax.device_put( + jnp.array(self._slice_byte_sizes, dtype=jnp.int32), sizes_sharding + ) + + # 3. Create shard_idx + global_ids = jnp.array( + [d.id for d in self._mesh.devices.flatten()], dtype=jnp.int32 + ).reshape(self._mesh.devices.shape) + self._shard_idx = jax.device_put( + global_ids, + jax.sharding.NamedSharding( + self._mesh, jax.sharding.PartitionSpec(*self._mesh.axis_names) + ), + ) + + # 4. Resolve num_shards + if num_shards is None: + num_processes = len( + set(d.process_index for d in self._mesh.devices.flatten()) + ) + if num_processes > 0: + self._num_shards = self._mesh.devices.size // num_processes + else: + self._num_shards = self._mesh.devices.size + else: + self._num_shards = num_shards + + # 5. Initialize via FFI + self._init_info = init_weight_synchronizer( + device_array=self._jax_arrays[0], + shard_idx=self._shard_idx, + mesh=self._mesh, + slice_byte_sizes=self._slice_byte_sizes_sharded, + local_port=local_port if local_port is not None else 0, + parallelism=self._parallelism, + num_layers=len(self._jax_arrays), + listener_port=listener_port if listener_port is not None else -1, + num_shards=self._num_shards, + ) + self._init_info.block_until_ready() + + try: + info_np = np.array(self._init_info) + flat_info = info_np.flatten() + self._local_port = ( + int(flat_info[4]) if len(flat_info) >= 5 else (local_port or 0) + ) + self._listener_port = ( + int(flat_info[5]) if len(flat_info) >= 6 else (listener_port or 0) + ) + except Exception: + self._local_port = local_port or 0 + self._listener_port = listener_port or 0 + + def d2h(self) -> None: + """Executes asynchronous Device-to-Host (D2H) copy from device memory to local staging buffer.""" + if self._destroyed: + raise RuntimeError("Cannot invoke d2h on destroyed WeightSynchronizer") + for layer_idx, arr in enumerate(self._jax_arrays): + d2h( + device_array=arr, + shard_idx=self._shard_idx, + mesh=self._mesh, + layer_idx=layer_idx, + ).block_until_ready() + + def h2d(self) -> None: + """Executes asynchronous Host-to-Device (H2D) copy from local staging buffer back to device memory.""" + if self._destroyed: + raise RuntimeError("Cannot invoke h2d on destroyed WeightSynchronizer") + res = multi_h2d( + device_arrays=self._jax_arrays, + shard_idx=self._shard_idx, + mesh=self._mesh, + ) + for arr in res: + arr.block_until_ready() + + def bind_weights(self, jax_arrays: Sequence[Any]) -> None: + """Binds updated JAX arrays to the weight synchronizer in-place.""" + if not jax_arrays: + raise ValueError("jax_arrays list cannot be empty") + self._jax_arrays = list(jax_arrays) + + def test_only_set_skip_tiling(self, skip: bool | Sequence[bool]) -> None: + """Sets whether D2H/H2D should skip CPU tiling/detiling.""" + pass + + def get_host_buffer(self, layer_idx: int = 0, shard_idx: int = 0) -> Any: + """Returns a zero-copy Host-side CPU NumPy ndarray view of the staging buffer.""" + raise NotImplementedError( + "Direct host buffer NumPy mapping is not supported on remote Pathways" + " FFI." + ) + + def get_local_endpoints(self) -> List[Dict[str, Any]]: + """Returns the list of transfer endpoints advertised by this instance.""" + return [{ + "endpoint": f"{self._bind_ip or 'localhost'}:{self.local_port}", + "shards": list(range(self.num_shards)), + }] + + @property + def local_port(self) -> Optional[int]: + """Returns the active local port assigned to the transceiving sockets server.""" + return self._local_port + + @property + def listener_port(self) -> Optional[int]: + """Returns the active local port assigned to the C++ Listener.""" + return self._listener_port + + @property + def is_listener_active(self) -> bool: + """Returns whether the native C++ Listener is actively running.""" + return is_listener_active() + + @property + def num_layers(self) -> int: + """Returns the total number of model weight layers registered.""" + return len(self._jax_arrays) + + @property + def num_shards(self) -> int: + """Returns the sharded devices count per layer.""" + return self._num_shards + + @property + def slice_byte_size(self) -> int: + """Returns the slice capacity per device block.""" + return self._slice_byte_sizes[0] if self._slice_byte_sizes else 0 + + def get_metrics(self) -> Any: + """Returns a dictionary of internal performance metrics.""" + return {} + + def reset_metrics(self) -> None: + """Resets all recorded internal metrics.""" + pass + + def destroy(self) -> None: + """Cleans up and deallocates WeightSynchronizer FFI resources.""" + if not self._destroyed: + destroy_weight_synchronizer() + self._destroyed = True + + def close(self) -> None: + """Alias for destroy().""" + self.destroy() + + def __del__(self) -> None: + try: + self.destroy() + except Exception: + pass