diff --git a/tpu_raiden/api/jax/BUILD b/tpu_raiden/api/jax/BUILD index 1a7c7e18..9ea6c0ed 100644 --- a/tpu_raiden/api/jax/BUILD +++ b/tpu_raiden/api/jax/BUILD @@ -75,10 +75,15 @@ py_test( py_library( name = "weight_synchronizer_jax_py", - srcs = ["weight_synchronizer.py"], + srcs = [ + "utils.py", + "weight_synchronizer.py", + ], visibility = ["//visibility:public"], deps = [ "//tpu_raiden/frameworks/jax:_tpu_raiden_jax", + "@jax//jax", + "@pypi//numpy", ], ) diff --git a/tpu_raiden/api/jax/utils.py b/tpu_raiden/api/jax/utils.py new file mode 100644 index 00000000..1ffc2656 --- /dev/null +++ b/tpu_raiden/api/jax/utils.py @@ -0,0 +1,149 @@ +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright 2026 The TPU Raiden Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Utility functions for JAX Weight Synchronizer.""" + +from typing import List +import jax +import numpy as np + + +def get_shard_sorting_permutation(arr: jax.Array) -> List[int]: + """Computes the permutation to sort JAX shards for Raiden Controller.""" + sharding = arr.sharding + if not isinstance(sharding, jax.sharding.NamedSharding): + return list(range(len(arr.addressable_shards))) + + mesh = sharding.mesh + spec = sharding.spec + + # 1. Reconstruct logical_mesh_shape and layout + logical_mesh_shape = [mesh.shape[ax] for ax in mesh.axis_names] + + sharded_axes = [] + for axis in spec: + if axis is None: + continue + if isinstance(axis, str): + sharded_axes.append(axis) + else: + sharded_axes.extend(axis) + + major_to_minor = sorted([mesh.axis_names.index(ax) for ax in sharded_axes]) + + # 2. Compute controller's expected global indices + num_shards = len(arr.addressable_shards) + num_physical_hosts = jax.process_count() + replica_id = jax.process_index() + + phys_mesh = [logical_mesh_shape[d] for d in major_to_minor] + + host_axis_logical = None + for d, size in enumerate(logical_mesh_shape): + if size == num_physical_hosts: + host_axis_logical = d + break + + non_host_axes = [ + d for d in range(len(logical_mesh_shape)) if d != host_axis_logical + ] + + controller_global_indices = [] + for j in range(num_shards): + local_coords = {} + temp = j + for d in reversed(non_host_axes): + size = logical_mesh_shape[d] + local_coords[d] = temp % size + temp = temp // size + + full_coords = [0] * len(logical_mesh_shape) + for d in range(len(logical_mesh_shape)): + if d == host_axis_logical: + full_coords[d] = replica_id + else: + full_coords[d] = local_coords.get(d, 0) + + tensor_coords = [full_coords[m_axis] for m_axis in major_to_minor] + + global_idx = 0 + stride = 1 + for val, size in zip(reversed(tensor_coords), reversed(phys_mesh)): + global_idx += val * stride + stride *= size + controller_global_indices.append(global_idx) + + # 3. Compute ACTUAL JAX global shard indices + jax_shard_global_indices = [] + for shard in arr.addressable_shards: + device = shard.device + coords = np.argwhere(mesh.devices == device) + if coords.size == 0: + raise ValueError(f"Device {device} not found in mesh") + m_coords = coords[0] + full_coords = list(m_coords) + + # Map array dimensions to mesh axes using spec + tensor_coords = [] + tensor_shape = [] + for axis in spec: + if axis is None: + tensor_coords.append(0) + tensor_shape.append(1) + elif isinstance(axis, str): + tensor_coords.append(full_coords[mesh.axis_names.index(axis)]) + tensor_shape.append(mesh.shape[axis]) + else: + for ax in axis: + tensor_coords.append(full_coords[mesh.axis_names.index(ax)]) + tensor_shape.append(mesh.shape[ax]) + + # Compute flat index (row-major) + global_idx = 0 + stride = 1 + for val, size in zip(reversed(tensor_coords), reversed(tensor_shape)): + global_idx += val * stride + stride *= size + jax_shard_global_indices.append(global_idx) + + # 4. Sort indices + indices = list(range(len(arr.addressable_shards))) + + def sort_key(idx): + g = jax_shard_global_indices[idx] + occurrence = jax_shard_global_indices[:idx].count(g) + matching_indices = [ + i for i, val in enumerate(controller_global_indices) if val == g + ] + if occurrence < len(matching_indices): + target_j = matching_indices[occurrence] + else: + target_j = matching_indices[-1] + return target_j + + sorted_indices = sorted(indices, key=sort_key) + return sorted_indices diff --git a/tpu_raiden/api/jax/weight_synchronizer_test.py b/tpu_raiden/api/jax/weight_synchronizer_test.py index 881bb563..9c9329e7 100644 --- a/tpu_raiden/api/jax/weight_synchronizer_test.py +++ b/tpu_raiden/api/jax/weight_synchronizer_test.py @@ -22,6 +22,7 @@ import jax.numpy as jnp import numpy as np +from tpu_raiden.api.jax import utils from tpu_raiden.api.jax import weight_synchronizer from tpu_raiden.rpc import raiden_service_pb2 @@ -43,6 +44,9 @@ def setUp(self): self.sharding = jax.sharding.NamedSharding( self.mesh, jax.sharding.PartitionSpec("data") ) + self.mesh_2d = jax.sharding.Mesh( + np.array(self.devices[:4]).reshape(2, 2), ("x", "y") + ) self.shape = (8, 128) self.dtype = jnp.float32 @@ -216,6 +220,110 @@ def test_bind_weights(self): for arr in dst_arrs: np.testing.assert_array_equal(np.asarray(arr), 5.0) + def _run_resharding_test(self, src_sharding, dst_sharding, shape): + src_arrs = [ + jax.device_put( + jnp.arange(np.prod(shape), dtype=self.dtype).reshape(shape), + src_sharding, + ) + ] + dst_arrs = [ + jax.device_put(jnp.zeros(shape, dtype=self.dtype), dst_sharding) + ] + + for arr in src_arrs: + arr.block_until_ready() + for arr in dst_arrs: + arr.block_until_ready() + + ws_source = WeightSynchronizer( + jax_arrays=src_arrs, + local_port=0, + unsafe_skip_buffer_lock=True, + listener_port=0, + bind_ip="127.0.0.1", + ) + ws_dest = WeightSynchronizer( + jax_arrays=dst_arrs, + local_port=0, + unsafe_skip_buffer_lock=True, + bind_ip="127.0.0.1", + ) + + req = raiden_service_pb2.ControlRequest( + command=raiden_service_pb2.ControlRequest.COMMAND_START_TRANSFER, + peers=[ + f"127.0.0.1:{ws_dest.local_port}", + ], + start_transfer_request=raiden_service_pb2.StartTransferRequest( + is_sender=True + ), + ) + payload = req.SerializeToString() + + sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM, 0) + sock.connect(("::1", ws_source.listener_port)) + sock.sendall(len(payload).to_bytes(4, "big") + payload) + + resp_len = int.from_bytes(sock.recv(4), "big") + resp_bytes = sock.recv(resp_len) + resp = raiden_service_pb2.ControlResponse() + resp.ParseFromString(resp_bytes) + self.assertTrue(resp.success) + sock.close() + + ws_dest.h2d() + + # Verify data integrity + np.testing.assert_array_equal( + np.asarray(dst_arrs[0]), np.asarray(src_arrs[0]) + ) + + def test_push_sync_aligned_to_aligned(self): + src_sharding = jax.sharding.NamedSharding( + self.mesh_2d, jax.sharding.PartitionSpec("x", "y") + ) + dst_sharding = jax.sharding.NamedSharding( + self.mesh_2d, jax.sharding.PartitionSpec("x", "y") + ) + self._run_resharding_test(src_sharding, dst_sharding, (8, 8)) + +class ShardSortingUtilTest(absltest.TestCase): + + def setUp(self): + super().setUp() + try: + self.devices = jax.devices("tpu") + except RuntimeError: + self.devices = jax.devices("cpu") + self.mesh_2d = jax.sharding.Mesh( + np.array(self.devices[:4]).reshape(2, 2), ("x", "y") + ) + + def test_aligned_sharding_permutation(self): + sharding = jax.sharding.NamedSharding( + self.mesh_2d, jax.sharding.PartitionSpec("x", "y") + ) + arr = jax.device_put(jnp.zeros((8, 8)), sharding) + perm = utils.get_shard_sorting_permutation(arr) + self.assertEqual(perm, [0, 1, 2, 3]) + + def test_transposed_sharding_permutation(self): + sharding = jax.sharding.NamedSharding( + self.mesh_2d, jax.sharding.PartitionSpec("y", "x") + ) + arr = jax.device_put(jnp.zeros((8, 8)), sharding) + perm = utils.get_shard_sorting_permutation(arr) + self.assertEqual(perm, [0, 2, 1, 3]) + + def test_replicated_sharding_permutation(self): + sharding = jax.sharding.NamedSharding( + self.mesh_2d, jax.sharding.PartitionSpec("x") + ) + arr = jax.device_put(jnp.zeros((8, 8)), sharding) + perm = utils.get_shard_sorting_permutation(arr) + self.assertEqual(perm, [0, 1, 2, 3]) + if __name__ == "__main__": absltest.main() diff --git a/tpu_raiden/frameworks/jax/utils.h b/tpu_raiden/frameworks/jax/utils.h index c7ca871a..a8b92f13 100644 --- a/tpu_raiden/frameworks/jax/utils.h +++ b/tpu_raiden/frameworks/jax/utils.h @@ -51,6 +51,29 @@ inline std::vector> UnpackJaxArrays( "Number of shards mismatch across layers during unpack"); } + try { + nanobind::object utils_mod = nanobind::module_::import_( + "tpu_raiden.api.jax.utils"); + nanobind::object get_permutation_fn = + utils_mod.attr("get_shard_sorting_permutation"); + nanobind::list permutation = + nanobind::cast(get_permutation_fn(dst)); + + std::vector sorted_shard_buffers; + sorted_shard_buffers.reserve(num_shards); + for (size_t i = 0; i < num_shards; ++i) { + size_t src_idx = nanobind::cast(permutation[i]); + if (src_idx >= shard_buffers.size()) { + throw std::runtime_error("Permutation index out of bounds"); + } + sorted_shard_buffers.push_back(std::move(shard_buffers[src_idx])); + } + shard_buffers = std::move(sorted_shard_buffers); + } catch (const std::exception& e) { + throw std::runtime_error(std::string("Failed to sort shards: ") + + e.what()); + } + layer_buffers.push_back(std::move(shard_buffers)); } return layer_buffers;