From 02b37babbf0bd3e409881edf441462c2dd0fefe9 Mon Sep 17 00:00:00 2001 From: Loki Chen Date: Fri, 4 Sep 2026 15:01:08 -0700 Subject: [PATCH] Forward skip_tiling across the controller boundary The remote-coordination branch of `_execute_slice_broadcast` (`tpu_sync/rpc/raiden_controller.py:1569`) forwards `skip_d2h=plan.skip_d2h` to `register_transfer_schedule` but does not forward `plan.skip_tiling`. Both flags configure how the host staging buffer is interpreted, and they are set together. The sender's D2h does see the map -- it reaches the local source worker through the `start_transfer(s_node, plan)` call just below -- so after the boundary the sender detiles according to a populated `skip_tiling` while the receiver's H2d re-tiles according to an empty one. Nothing catches the disagreement: no bounds check fires, the chunk counts still match, the transfer reports success, and the landed weights are silently wrong. Adds the missing kwarg plus a regression test that records the kwargs handed to a faked `RaidenControllerClientFacade` and asserts `skip_tiling` arrives intact alongside `skip_d2h`. Measured, on an unmodified tree: without the fix the new test fails, AssertionError: None != {7: True, 9: False} rest of raiden_controller_test.py: 35/35 pass with the fix the new test passes raiden_controller_test.py: 36/36 pass The test discriminates the one changed line rather than the setup: its `assertTrue(kwargs.get("skip_d2h"))` assertion passes on both arms, and only the `skip_tiling` assertion moves. Scope note, deliberately narrow: this is a forwarding omission found while investigating a separate replicated-leaf drop. It is **not** claimed to be the cause of that drop -- that root cause is still unknown, and the earlier hypotheses (collision, rank-1 tiling, C++ `(slice, dst_device)` keying) were each tested and refuted. This change stands on its own. Signed-off-by: Loki Chen --- tpu_sync/rpc/raiden_controller.py | 11 ++++ tpu_sync/rpc/raiden_controller_test.py | 70 ++++++++++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/tpu_sync/rpc/raiden_controller.py b/tpu_sync/rpc/raiden_controller.py index 018b9ad0..1f2a10e6 100644 --- a/tpu_sync/rpc/raiden_controller.py +++ b/tpu_sync/rpc/raiden_controller.py @@ -1738,6 +1738,17 @@ async def _run_single_transfer(s_node, d_node, plan): plan.shard_push_schedules, dst_mem_type, skip_d2h=plan.skip_d2h, + # skip_tiling must cross the controller boundary with + # skip_d2h. The sender's D2h honours this map (it + # reaches the local source worker via the + # start_transfer(s_node, plan) call below), so + # omitting it here leaves the receiver's H2d with an + # empty active_skip while the sender detiled + # according to a populated one. The two sides then + # disagree about the host staging buffer's format and + # the landed weights are silently wrong: no bounds + # check fires and the chunk counts still match. + skip_tiling=plan.skip_tiling, ), ) if not success: diff --git a/tpu_sync/rpc/raiden_controller_test.py b/tpu_sync/rpc/raiden_controller_test.py index f4fa3497..5924ec7d 100644 --- a/tpu_sync/rpc/raiden_controller_test.py +++ b/tpu_sync/rpc/raiden_controller_test.py @@ -711,6 +711,76 @@ def mock_start_transfer(*args, **kwargs): controller.start_transfer = original_start_transfer server.stop() + def test_slice_broadcast_forwards_skip_tiling_to_remote_controller(self): + """skip_tiling must cross the controller boundary alongside skip_d2h. + + Both configure how the host staging buffer is interpreted. The + remote-coordination branch of _execute_slice_broadcast forwarded only + skip_d2h, so the sender's D2h detiled according to a populated map while + the receiver's H2d re-tiled according to an empty one. Nothing detects + that: no bounds check fires, the chunk counts still match, the transfer + reports success, and the landed weights are silently wrong. + """ + controller = raiden_controller.RaidenController(port=0) + + src = raiden_controller.RaidenId("trainer", "0", "weights") + dst = raiden_controller.RaidenId("sampler", "0", "weights") + + recorded = [] + + class RecordingFacade: + + def __init__(self, address, name_resolver=None): + del address, name_resolver + + def register_transfer_schedule(self, *args, **kwargs): + recorded.append((args, kwargs)) + return True + + skip_tiling = {7: True, 9: False} + final_plan = raiden_controller.TransferPlan( + src_units=[src], + dst_units=[dst], + plan=None, + shard_push_schedules={}, + worker_rpc_addresses={}, + worker_data_addresses={}, + uuid=1, + dst_mem_type=raiden_controller.RaidenMemoryType.DRAM, + use_block_chunks=True, + is_sender=True, + expected_block_count=0, + req_id="req-bcast", + skip_d2h=True, + skip_tiling=skip_tiling, + ) + + # key: (src_unit, shard_idx, src_block_id, src_block_offset, size, + # src_stride, count, layer_idx, pool_group) + key = (src, 0, 0, 0, 4096, 0, 1, 7, 0) + # target: (dst_unit, dst_peer, dst_shard_idx, dst_block_id, + # dst_block_offset, dst_stride) + targets = [(dst, "10.0.0.2:8000", 0, 0, 0, 0)] + + with mock.patch.object( + raiden_controller, "RaidenControllerClientFacade", RecordingFacade + ): + asyncio.run( + controller._execute_slice_broadcast( + keys_and_targets=[(key, targets)], + final_plan=final_plan, + fanout_k=1, + req_id="req-bcast", + dst_mem_type=raiden_controller.RaidenMemoryType.DRAM, + dst_controller_address="10.0.0.2:9000", + ) + ) + + self.assertLen(recorded, 1) + _, kwargs = recorded[0] + self.assertTrue(kwargs.get("skip_d2h")) + self.assertEqual(kwargs.get("skip_tiling"), skip_tiling) + def test_multi_variable_resharding_planning(self): """Verifies resharding planning for multiple variables using absolute offsets.""" client = RecordingWorkerRpcClient()