diff --git a/tpu_sync/api/torch/BUILD b/tpu_sync/api/torch/BUILD index ab28054c..7a9f965d 100644 --- a/tpu_sync/api/torch/BUILD +++ b/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", ], diff --git a/tpu_sync/api/torch/weight_synchronizer.py b/tpu_sync/api/torch/weight_synchronizer.py index d317c806..3f7054ba 100644 --- a/tpu_sync/api/torch/weight_synchronizer.py +++ b/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() diff --git a/tpu_sync/api/torch/weight_synchronizer_test.py b/tpu_sync/api/torch/weight_synchronizer_test.py index ed8ba1ff..8be933a6 100644 --- a/tpu_sync/api/torch/weight_synchronizer_test.py +++ b/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,77 +33,83 @@ 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() @@ -114,14 +117,157 @@ def test_e2e_3node_distributed_weight_push(self, dtype): # 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]]: diff --git a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc index 151e9f6f..ee55d976 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc +++ b/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()) + .def( + "bind_weights", + [](WeightSynchronizer& self, + const std::vector>& 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()) .def( "D2h", diff --git a/tpu_sync/frameworks/torch/weight_synchronizer.cc b/tpu_sync/frameworks/torch/weight_synchronizer.cc index 01e7b2b7..f6c3a16b 100644 --- a/tpu_sync/frameworks/torch/weight_synchronizer.cc +++ b/tpu_sync/frameworks/torch/weight_synchronizer.cc @@ -68,6 +68,53 @@ NumaAwareWeightSynchronizer::NumaAwareWeightSynchronizer( InitSubManagers(unpacked.buffers, local_port, unsafe_skip_buffer_lock, parallelism, listener_port, bind_ip, auto_h2d); } + +absl::Status NumaAwareWeightSynchronizer::BindWeights( + const std::vector>& device_tensors) { + try { + UnpackedTensors unpacked = + UnpackTorchTensors(device_tensors, unsafe_skip_buffer_lock_); + const auto& layer_buffers = unpacked.buffers; + if (layer_buffers.empty()) { + return absl::InvalidArgumentError( + "Empty layer buffers provided to BindWeights"); + } + if (layer_buffers.size() != num_layers_) { + return absl::InvalidArgumentError( + absl::StrCat("Layer count mismatch in BindWeights: expected ", + num_layers_, ", got ", layer_buffers.size())); + } + if (layer_buffers[0].size() != total_num_shards_) { + return absl::InvalidArgumentError( + absl::StrCat("Shard count mismatch in BindWeights: expected ", + total_num_shards_, ", got ", layer_buffers[0].size())); + } + + for (size_t s = 0; s < sub_synchronizers_.size(); ++s) { + if (!sub_synchronizers_[s]) continue; + const auto& local_shards = (s < submanager_to_local_shards_.size()) + ? submanager_to_local_shards_[s] + : std::vector{}; + std::vector> sub_buffers( + num_layers_); + for (size_t l = 0; l < num_layers_; ++l) { + sub_buffers[l].reserve(local_shards.size()); + for (int lsh : local_shards) { + if (lsh < 0 || lsh >= static_cast(layer_buffers[l].size())) { + return absl::OutOfRangeError("Local shard index out of range"); + } + sub_buffers[l].push_back(layer_buffers[l][lsh]); + } + } + auto status = sub_synchronizers_[s]->BindWeights(sub_buffers); + if (!status.ok()) return status; + } + buffer_refs_ = std::move(unpacked.refs); + return absl::OkStatus(); + } catch (const std::exception& e) { + return absl::InternalError(e.what()); + } +} #endif NumaAwareWeightSynchronizer::NumaAwareWeightSynchronizer( @@ -877,6 +924,11 @@ WeightSynchronizer::WeightSynchronizer( device_tensors, local_port, parallelism, listener_port, bind_ip, unsafe_skip_buffer_lock, auto_h2d); } + +absl::Status WeightSynchronizer::BindWeights( + const std::vector>& device_tensors) { + return numa_manager_->BindWeights(device_tensors); +} #endif WeightSynchronizer::WeightSynchronizer( diff --git a/tpu_sync/frameworks/torch/weight_synchronizer.h b/tpu_sync/frameworks/torch/weight_synchronizer.h index 2bed2e67..57168a69 100644 --- a/tpu_sync/frameworks/torch/weight_synchronizer.h +++ b/tpu_sync/frameworks/torch/weight_synchronizer.h @@ -62,6 +62,9 @@ class NumaAwareWeightSynchronizer std::optional listener_port = std::nullopt, std::optional bind_ip = std::nullopt, bool unsafe_skip_buffer_lock = true, bool auto_h2d = false); + + absl::Status BindWeights( + const std::vector>& device_tensors); #endif // CPU / Mock metadata constructor for tests without PJRT TPU devices @@ -169,6 +172,9 @@ class WeightSynchronizer { std::optional listener_port = std::nullopt, std::optional bind_ip = std::nullopt, bool unsafe_skip_buffer_lock = true, bool auto_h2d = false); + + absl::Status BindWeights( + const std::vector>& device_tensors); #endif // CPU / Mock metadata constructor for tests without PJRT TPU devices