Expand PyTorch WeightSynchronizer test suite to reach parity with JAX:

tpu_sync/api/torch/BUILD

@@ -169,7 +169,6 @@ py_test(
         ":weight_synchronizer_torch_py",
         "@com_google_absl_py//absl/testing:absltest",
         "@com_google_absl_py//absl/testing:parameterized",
-        "@pypi//numpy",
         "@torch_tpu//shims/torch:pytorch",
         "@torch_tpu//torch_tpu",
     ],

tpu_sync/api/torch/weight_synchronizer.py

@@ -77,6 +77,10 @@ def test_only_set_skip_tiling(self, skip: bool | List[bool]) -> None:
     else:
       self._impl.set_skip_tiling(list(skip))
 
+  def bind_weights(self, device_tensors: List[List[torch.Tensor]]) -> None:
+    """Dynamically re-binds new device weights in-place without daemon restart."""
+    self._impl.bind_weights(device_tensors)
+
   def d2h(self) -> None:
     """Triggers asynchronous D2H copy of current weights to Host buffer."""
     self._impl.D2h()

tpu_sync/api/torch/weight_synchronizer_test.py

@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""E2E physical integration tests for PyTorch WeightSynchronizer on XLA TPUs."""
-
 import os
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import numpy as np
 import torch
 import torch_tpu  # pylint: disable=unused-import
 
@@ -36,92 +33,241 @@ def setUp(self):
     self.num_layers = 2
     self.num_shards = 1
     self.block_size = 2
-    self.slice_byte_size = 16384 // 4  # float32 capacity
+
+  def _run_push_sync(
+      self,
+      src_tensors: list[list[torch.Tensor]],
+      dst_tensors: list[list[torch.Tensor]],
+  ):
+    ws_source = WeightSynchronizer(
+        src_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
+    )
+    ws_dest = WeightSynchronizer(
+        dst_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
+    )
+    self.assertIsNotNone(ws_source.local_port)
+    self.assertIsNotNone(ws_dest.local_port)
+
+    peer_dest = f"127.0.0.1:{ws_dest.local_port}"
+    ws_source.push_weights([peer_dest])
+    ws_dest.h2d()
+
+    for l in range(len(src_tensors)):
+      for sh in range(len(src_tensors[l])):
+        self.assertTrue(
+            torch.equal(dst_tensors[l][sh].cpu(), src_tensors[l][sh].cpu())
+        )
 
   @parameterized.named_parameters(
       ("fp32", torch.float32),
+      ("bf16", torch.bfloat16),
       ("int32", torch.int32),
   )
   def test_e2e_3node_distributed_weight_push(self, dtype):
-    shape = (self.block_size, 128, 8)  # 16384 bytes capacity per layer shard
-
-    # 1. Allocate source (Trainer) weights on Device TPU
-    src_tensors = []
-    for l in range(self.num_layers):
-      shards = []
-      for sh in range(self.num_shards):
-        t = torch.zeros(shape, dtype=dtype, device=self.device)
-        shards.append(t)
-      src_tensors.append(shards)
-
-    # Allocate destination 1 (Inference Peer 1) weights
-    dst1_tensors = []
-    for l in range(self.num_layers):
-      shards = []
-      for sh in range(self.num_shards):
-        t = torch.zeros(shape, dtype=dtype, device=self.device)
-        shards.append(t)
-      dst1_tensors.append(shards)
+    shape = (self.block_size, 128, 8)
 
-    # Allocate destination 2 (Inference Peer 2) weights
-    dst2_tensors = []
-    for l in range(self.num_layers):
-      shards = []
-      for sh in range(self.num_shards):
-        t = torch.zeros(shape, dtype=dtype, device=self.device)
-        shards.append(t)
-      dst2_tensors.append(shards)
+    src_tensors = [
+        [
+            torch.zeros(shape, dtype=dtype, device=self.device)
+            for _ in range(self.num_shards)
+        ]
+        for _ in range(self.num_layers)
+    ]
+    dst1_tensors = [
+        [
+            torch.zeros(shape, dtype=dtype, device=self.device)
+            for _ in range(self.num_shards)
+        ]
+        for _ in range(self.num_layers)
+    ]
+    dst2_tensors = [
+        [
+            torch.zeros(shape, dtype=dtype, device=self.device)
+            for _ in range(self.num_shards)
+        ]
+        for _ in range(self.num_layers)
+    ]
 
-    # 2. Instantiate destination WeightSynchronizers on ephemeral ports!
     ws_dest1 = WeightSynchronizer(
         dst1_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
     )
     ws_dest2 = WeightSynchronizer(
         dst2_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
     )
-
     self.assertIsNotNone(ws_dest1.local_port)
     self.assertIsNotNone(ws_dest2.local_port)
 
-    peer_dest1 = f"localhost:{ws_dest1.local_port}"
-    peer_dest2 = f"localhost:{ws_dest2.local_port}"
+    peer_dest1 = f"127.0.0.1:{ws_dest1.local_port}"
+    peer_dest2 = f"127.0.0.1:{ws_dest2.local_port}"
 
-    # ==========================================================================
-    # Scenario A: Test the Push API E2E (1 Source pushes to 2 Destinations!)
-    # ==========================================================================
-    # Trainer fills source weights with distinct values per layer
     for l in range(self.num_layers):
       for sh in range(self.num_shards):
-        val = float(l + 10.0)  # Layer 0=10.0, Layer 1=11.0
+        val = int(l + 10) if dtype == torch.int32 else float(l + 10.0)
         src_tensors[l][sh].fill_(val)
 
-    # Force execution of fill_ on source tensors to ensure TPU memory is updated
-    for l in range(self.num_layers):
-      for sh in range(self.num_shards):
-        _ = src_tensors[l][sh].cpu()
-
-    # Recreate/Instantiate ws_source to capture filled buffers!
     ws_source = WeightSynchronizer(
         src_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
     )
     self.assertIsNotNone(ws_source.local_port)
 
-    # Source pushes weights to both dest1 and dest2 socket servers E2E!
     ws_source.push_weights([peer_dest1, peer_dest2])
     ws_dest1.h2d()
     ws_dest2.h2d()
 
     # Assert both destinations have received the trainer's weights on TPU HBM!
     for l in range(self.num_layers):
       for sh in range(self.num_shards):
-        expected_val = float(l + 10.0)
-        np.testing.assert_allclose(
-            dst1_tensors[l][sh].cpu().numpy(), expected_val, atol=1e-5
+        self.assertTrue(
+            torch.equal(dst1_tensors[l][sh].cpu(), src_tensors[l][sh].cpu())
         )
-        np.testing.assert_allclose(
-            dst2_tensors[l][sh].cpu().numpy(), expected_val, atol=1e-5
+        self.assertTrue(
+            torch.equal(dst2_tensors[l][sh].cpu(), src_tensors[l][sh].cpu())
         )
 
+  @parameterized.named_parameters(
+      ("fp32", torch.float32),
+      ("bf16", torch.bfloat16),
+  )
+  def test_bind_weights(self, dtype):
+    shape = (self.block_size, 128, 8)
+
+    src_tensors = [
+        [torch.full(shape, fill_value=5.0, dtype=dtype, device=self.device)]
+        for _ in range(self.num_layers)
+    ]
+    dst_tensors = [
+        [torch.zeros(shape, dtype=dtype, device=self.device)]
+        for _ in range(self.num_layers)
+    ]
+
+    ws_source = WeightSynchronizer(
+        src_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
+    )
+    ws_dest = WeightSynchronizer(
+        dst_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
+    )
+    peer_dest = f"127.0.0.1:{ws_dest.local_port}"
+
+    # --- Sync 1 (V1: 5.0 -> 0.0) ---
+    ws_source.push_weights([peer_dest])
+    ws_dest.h2d()
+
+    for l in range(self.num_layers):
+      self.assertTrue(
+          torch.equal(dst_tensors[l][0].cpu(), src_tensors[l][0].cpu())
+      )
+
+    # --- Bind weights to V2 ---
+    new_src_tensors = [
+        [torch.full(shape, fill_value=10.0, dtype=dtype, device=self.device)]
+        for _ in range(self.num_layers)
+    ]
+    ws_source.bind_weights(new_src_tensors)
+    ws_source.d2h()
+
+    new_dst_tensors = [
+        [torch.full(shape, fill_value=-1.0, dtype=dtype, device=self.device)]
+        for _ in range(self.num_layers)
+    ]
+    ws_dest.bind_weights(new_dst_tensors)
+
+    # --- Sync 2 (V2: 10.0 -> -1.0) ---
+    ws_source.push_weights([peer_dest])
+    ws_dest.h2d()
+
+    # Verify Sync 2 updated new_dst_tensors to 10.0
+    for l in range(self.num_layers):
+      self.assertTrue(
+          torch.equal(new_dst_tensors[l][0].cpu(), new_src_tensors[l][0].cpu())
+      )
+
+    # Verify original V1 dst_tensors were NOT overwritten (still 5.0)
+    for l in range(self.num_layers):
+      self.assertTrue(
+          torch.equal(dst_tensors[l][0].cpu(), src_tensors[l][0].cpu())
+      )
+
+  @parameterized.named_parameters(
+      ("fp32", torch.float32),
+      ("bf16", torch.bfloat16),
+  )
+  def test_heterogeneous_layers_small_first(self, dtype):
+    shapes = [(1024,), (1024, 3072), (2048, 2048)]
+    src_tensors = [
+        [
+            torch.full(
+                shape,
+                fill_value=float(i + 1.0),
+                dtype=dtype,
+                device=self.device,
+            )
+        ]
+        for i, shape in enumerate(shapes)
+    ]
+    dst_tensors = [
+        [torch.zeros(shape, dtype=dtype, device=self.device)]
+        for shape in shapes
+    ]
+    self._run_push_sync(src_tensors, dst_tensors)
+
+  @parameterized.named_parameters(
+      ("fp32", torch.float32),
+      ("bf16", torch.bfloat16),
+  )
+  def test_heterogeneous_layers_large_first(self, dtype):
+    shapes = [(1024, 3072), (1024,), (128,)]
+    src_tensors = [
+        [
+            torch.full(
+                shape,
+                fill_value=float(i + 1.0),
+                dtype=dtype,
+                device=self.device,
+            )
+        ]
+        for i, shape in enumerate(shapes)
+    ]
+    dst_tensors = [
+        [torch.zeros(shape, dtype=dtype, device=self.device)]
+        for shape in shapes
+    ]
+    self._run_push_sync(src_tensors, dst_tensors)
+
+  @parameterized.named_parameters(
+      ("fp32", torch.float32),
+      ("bf16", torch.bfloat16),
+  )
+  def test_heterogeneous_layers_local_roundtrip(self, dtype):
+    shapes = [(1024,), (1024, 3072), (2048, 2048)]
+    src_tensors = [
+        [
+            torch.full(
+                shape,
+                fill_value=float(i + 10.0),
+                dtype=dtype,
+                device=self.device,
+            )
+        ]
+        for i, shape in enumerate(shapes)
+    ]
+    ws = WeightSynchronizer(
+        src_tensors, local_port=0, parallelism=1, bind_ip="127.0.0.1"
+    )
+    ws.d2h()
+
+    # Zero out new destination tensors and bind them
+    zero_tensors = [
+        [torch.zeros(shape, dtype=dtype, device=self.device)]
+        for shape in shapes
+    ]
+    ws.bind_weights(zero_tensors)
+    ws.h2d()
+
+    for i in range(len(shapes)):
+      self.assertTrue(
+          torch.equal(zero_tensors[i][0].cpu(), src_tensors[i][0].cpu())
+      )
+
   def _make_tensors(
       self, num_layers: int, num_shards: int
   ) -> list[list[torch.Tensor]]:

tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc

@@ -542,6 +542,18 @@ NB_MODULE(_tpu_raiden_torch, m) {
             }
           },
           nb::arg("peers"), nb::call_guard<nb::gil_scoped_release>())
+      .def(
+          "bind_weights",
+          [](WeightSynchronizer& self,
+             const std::vector<std::vector<at::Tensor>>& device_tensors) {
+            absl::Status s = self.BindWeights(device_tensors);
+            if (!s.ok()) {
+              throw std::runtime_error(
+                  "WeightSynchronizer bind_weights failed: " +
+                  std::string(s.message()));
+            }
+          },
+          nb::arg("device_tensors"), nb::call_guard<nb::gil_scoped_release>())
 
       .def(
           "D2h",