From a1d70ca16bae0b4fac86556e8e8204ef4ca24b26 Mon Sep 17 00:00:00 2001 From: datenglin Date: Sat, 26 Sep 2026 00:45:15 -0700 Subject: [PATCH] Fixed multi-shard relay scheduling and preserved transfer UUIDs in BroadcastEngine. PiperOrigin-RevId: 988733182 --- .../weight_sync/manager/broadcast_engine.py | 37 ++-- .../manager/broadcast_engine_test.py | 179 ++++++++++++++++++ 2 files changed, 198 insertions(+), 18 deletions(-) diff --git a/tpu_sync/weight_sync/manager/broadcast_engine.py b/tpu_sync/weight_sync/manager/broadcast_engine.py index 58e6f848f..a58949b87 100644 --- a/tpu_sync/weight_sync/manager/broadcast_engine.py +++ b/tpu_sync/weight_sync/manager/broadcast_engine.py @@ -498,15 +498,10 @@ def _dispatch_hop(hop: _HopTask) -> None: active_pushes[s] += 1 - ref_offset = group.node_slice_offsets[s][group.ref_key] - s_shard_idx = ref_offset[0] - - sub_schedule: dict[RaidenId, dict[int, list[Any]]] = { - s: {s_shard_idx: []} - } + sub_schedule: dict[RaidenId, dict[int, list[Any]]] = {s: {}} hop_expected_block_count = 0 for key, k_targets in group.keys_and_sorted_targets: - _, k_s_block_id, k_s_block_offset, k_s_stride = ( + k_s_shard_idx, k_s_block_id, k_s_block_offset, k_s_stride = ( group.node_slice_offsets[s][key] ) k_size = key[4] @@ -539,7 +534,7 @@ def _dispatch_hop(hop: _HopTask) -> None: k_layer_idx, k_pool_group, ) - sub_schedule[s][s_shard_idx].append(entry) + sub_schedule[s].setdefault(k_s_shard_idx, []).append(entry) is_contiguous = (k_count == 1) or ( k_s_stride == k_size and k_dst_stride == k_size @@ -548,17 +543,23 @@ def _dispatch_hop(hop: _HopTask) -> None: hop_expected_block_count += push_count if s != group.src_unit: - sub_schedule[s][s_shard_idx] = _coalesce_contiguous_relay_entries( - sub_schedule[s][s_shard_idx], - is_weight_sync=bool(final_plan.is_weight_sync), - ) - hop_expected_block_count = sum( - 1 if (e[9] == 1 or (e[7] == e[4] and e[8] == e[4])) else e[9] - for e in sub_schedule[s][s_shard_idx] - ) + hop_expected_block_count = 0 + for shard_idx, entries in list(sub_schedule[s].items()): + sub_schedule[s][shard_idx] = _coalesce_contiguous_relay_entries( + entries, + is_weight_sync=bool(final_plan.is_weight_sync), + ) + hop_expected_block_count += sum( + 1 if (e[9] == 1 or (e[7] == e[4] and e[8] == e[4])) else e[9] + for e in sub_schedule[s][shard_idx] + ) - hop_uuid = random.randint(1, 2**63 - 1) - hop_req_id = f"{req_id}_{group.group_idx}_{hop_uuid}" + hop_uuid = ( + final_plan.uuid + if (final_plan.uuid is not None and final_plan.uuid > 0) + else random.randint(1, 2**63 - 1) + ) + hop_req_id = f"{req_id}_{group.group_idx}_{dst_unit}_{hop_uuid}" sub_plan = type(final_plan)( src_units=[s], diff --git a/tpu_sync/weight_sync/manager/broadcast_engine_test.py b/tpu_sync/weight_sync/manager/broadcast_engine_test.py index fa21b5845..c276e5164 100644 --- a/tpu_sync/weight_sync/manager/broadcast_engine_test.py +++ b/tpu_sync/weight_sync/manager/broadcast_engine_test.py @@ -540,6 +540,185 @@ def test_execute_slice_broadcast_pipeline_invalid_fanout(self) -> None: ) ) + def test_execute_slice_broadcast_multishard_relay(self) -> None: + """Verifies relay nodes dispatch push schedules partitioned by actual shard index.""" + rpc_client = RecordingWorkerRpcClient() + self.addCleanup(rpc_client.close) + engine = broadcast_engine.BroadcastEngine(rpc_client) + + src = RaidenId(job_name="src", job_replica_id="0", data_name="w") + relay = RaidenId(job_name="relay", job_replica_id="0", data_name="w") + dst = RaidenId(job_name="dst", job_replica_id="0", data_name="w") + + # key0 maps to relay shard 0 (e.g. Host 0) + key0 = (src, 0, 0, 0, 1024, 0, 1, 0, 0) + # key1 maps to relay shard 1 (e.g. Host 1) + key1 = (src, 0, 1, 0, 1024, 0, 1, 0, 0) + + targets0: list[tuple[Any, ...]] = [ + (relay, "127.0.0.1:8001", 0, 0, 0, 0), + (dst, "127.0.0.1:8002", 0, 0, 0, 0), + ] + targets1: list[tuple[Any, ...]] = [ + (relay, "127.0.0.1:8001", 1, 1, 0, 0), + (dst, "127.0.0.1:8002", 1, 1, 0, 0), + ] + + final_plan = raiden_controller.TransferPlan( + src_units=[src], + dst_units=[relay, dst], + plan=None, + worker_data_addresses={ + src: ["127.0.0.1:8000"], + relay: ["127.0.0.1:8001"], + dst: ["127.0.0.1:8002"], + }, + is_weight_sync=True, + uuid=42, + ) + registered_shards = {src: ["s0"], relay: ["s0", "s1"], dst: ["s0", "s1"]} + + asyncio.run( + engine.execute_slice_broadcast( + keys_and_targets=[(key0, targets0), (key1, targets1)], + final_plan=final_plan, + fanout_k=1, + req_id="req_multishard_test", + dst_mem_type=raiden_controller.RaidenMemoryType.DRAM, + registered_shards=registered_shards, + ) + ) + + # Find the relay's forward transfer plan + # (where relay is sender, dst is receiver) + relay_forward_plans = [ + plan + for target_id, plan in rpc_client.invocations + if target_id == relay + and plan.src_units[0] == relay + and plan.dst_units[0] == dst + ] + self.assertLen(relay_forward_plans, 1) + relay_plan = relay_forward_plans[0] + self.assertEqual(relay_plan.uuid, 42, "Relay plan must preserve final_plan.uuid") + + # Verify both shard 0 and shard 1 are present in the relay's push schedule + self.assertIn(relay, relay_plan.shard_push_schedules) + relay_schedules = relay_plan.shard_push_schedules[relay] + self.assertIn( + 0, relay_schedules, "Relay must include push schedules for shard 0" + ) + self.assertIn( + 1, relay_schedules, "Relay must include push schedules for shard 1" + ) + + # Verify key0 is under shard 0 and key1 is under shard 1 + self.assertLen(relay_schedules[0], 1) + entry0 = relay_schedules[0][0] + self.assertEqual( + entry0[1], 0, "Destination shard index for key0 should be 0" + ) + self.assertEqual(entry0[5], 0, "Source block ID for key0 should be 0") + self.assertEqual(entry0[6], 0, "Destination block ID for key0 should be 0") + + self.assertLen(relay_schedules[1], 1) + entry1 = relay_schedules[1][0] + self.assertEqual( + entry1[1], 1, "Destination shard index for key1 should be 1" + ) + self.assertEqual(entry1[5], 1, "Source block ID for key1 should be 1") + self.assertEqual(entry1[6], 1, "Destination block ID for key1 should be 1") + + self.assertEqual(relay_plan.expected_block_count, 2) + + def test_execute_slice_broadcast_multishard_relay_coalescing(self) -> None: + """Verifies contiguous slices on multiple shards coalesce independently on relay.""" + rpc_client = RecordingWorkerRpcClient() + self.addCleanup(rpc_client.close) + engine = broadcast_engine.BroadcastEngine(rpc_client) + + src = RaidenId(job_name="src", job_replica_id="0", data_name="w") + relay = RaidenId(job_name="relay", job_replica_id="0", data_name="w") + dst = RaidenId(job_name="dst", job_replica_id="0", data_name="w") + + # Two contiguous slices for shard 0 + # (block_id 0, offsets 0 and 512, size 512) + key0_a = (src, 0, 0, 0, 512, 0, 1, 0, 0) + key0_b = (src, 0, 0, 512, 512, 0, 1, 0, 0) + # Two contiguous slices for shard 1 + # (block_id 1, offsets 0 and 512, size 512) + key1_a = (src, 0, 1, 0, 512, 0, 1, 0, 0) + key1_b = (src, 0, 1, 512, 512, 0, 1, 0, 0) + + targets0_a: list[tuple[Any, ...]] = [ + (relay, "127.0.0.1:8001", 0, 0, 0, 0), + (dst, "127.0.0.1:8002", 0, 0, 0, 0), + ] + targets0_b: list[tuple[Any, ...]] = [ + (relay, "127.0.0.1:8001", 0, 0, 512, 0), + (dst, "127.0.0.1:8002", 0, 0, 512, 0), + ] + targets1_a: list[tuple[Any, ...]] = [ + (relay, "127.0.0.1:8001", 1, 1, 0, 0), + (dst, "127.0.0.1:8002", 1, 1, 0, 0), + ] + targets1_b: list[tuple[Any, ...]] = [ + (relay, "127.0.0.1:8001", 1, 1, 512, 0), + (dst, "127.0.0.1:8002", 1, 1, 512, 0), + ] + + final_plan = raiden_controller.TransferPlan( + src_units=[src], + dst_units=[relay, dst], + plan=None, + worker_data_addresses={ + src: ["127.0.0.1:8000"], + relay: ["127.0.0.1:8001"], + dst: ["127.0.0.1:8002"], + }, + is_weight_sync=True, + ) + registered_shards = {src: ["s0"], relay: ["s0", "s1"], dst: ["s0", "s1"]} + + asyncio.run( + engine.execute_slice_broadcast( + keys_and_targets=[ + (key0_a, targets0_a), + (key0_b, targets0_b), + (key1_a, targets1_a), + (key1_b, targets1_b), + ], + final_plan=final_plan, + fanout_k=1, + req_id="req_coalesce_test", + dst_mem_type=raiden_controller.RaidenMemoryType.DRAM, + registered_shards=registered_shards, + ) + ) + + relay_forward_plans = [ + plan + for target_id, plan in rpc_client.invocations + if target_id == relay + and plan.src_units[0] == relay + and plan.dst_units[0] == dst + ] + self.assertLen(relay_forward_plans, 1) + relay_plan = relay_forward_plans[0] + relay_schedules = relay_plan.shard_push_schedules[relay] + + # Both shards must have their 2 slices coalesced into 1 contiguous transfer + # of 1024 bytes + self.assertLen(relay_schedules[0], 1) + self.assertEqual(relay_schedules[0][0][4], 1024) + self.assertEqual(relay_schedules[0][0][5], 0) + + self.assertLen(relay_schedules[1], 1) + self.assertEqual(relay_schedules[1][0][4], 1024) + self.assertEqual(relay_schedules[1][0][5], 1) + + self.assertEqual(relay_plan.expected_block_count, 2) + if __name__ == "__main__": absltest.main()