From c1b1bf8d3acc4c453e1036d9ba2d77cee1306bc2 Mon Sep 17 00:00:00 2001 From: datenglin Date: Wed, 19 Aug 2026 23:14:00 -0700 Subject: [PATCH] Added `auto_h2d` for Torch weight synchronizer. PiperOrigin-RevId: 967606021 --- tpu_sync/api/torch/BUILD | 3 ++- tpu_sync/api/torch/weight_synchronizer.py | 20 ++++++++++++++----- tpu_sync/frameworks/torch/BUILD | 4 +++- .../torch/tpu_raiden_torch_module.cc | 6 ++++-- .../frameworks/torch/weight_synchronizer.cc | 20 +++++++++---------- .../frameworks/torch/weight_synchronizer.h | 8 +++++--- 6 files changed, 39 insertions(+), 22 deletions(-) diff --git a/tpu_sync/api/torch/BUILD b/tpu_sync/api/torch/BUILD index bd4f2fa9..311ca28f 100644 --- a/tpu_sync/api/torch/BUILD +++ b/tpu_sync/api/torch/BUILD @@ -78,8 +78,9 @@ py_library( deps = [ ":torch_abi", ":torch_tpu_common_loader", - "//tpu_sync/frameworks/torch:_tpu_raiden_torch", "@torch_tpu//shims/torch:pytorch", + # buildcleaner: keep + "//tpu_sync/frameworks/torch:_tpu_raiden_torch", ], ) diff --git a/tpu_sync/api/torch/weight_synchronizer.py b/tpu_sync/api/torch/weight_synchronizer.py index fd9a80aa..4cf9481a 100644 --- a/tpu_sync/api/torch/weight_synchronizer.py +++ b/tpu_sync/api/torch/weight_synchronizer.py @@ -43,6 +43,8 @@ def __init__( parallelism: int = 1, listener_port: Optional[int] = None, bind_ip: Optional[str] = None, + unsafe_skip_buffer_lock: bool = True, + auto_h2d: bool = False, ): """Instantiates the PyTorch Weight Synchronizer shims. @@ -54,27 +56,35 @@ def __init__( parallelism: Parallel TCP sockets workers count. listener_port: RPC control listener port. bind_ip: Sockets server bind IP address. + unsafe_skip_buffer_lock: Whether to bypass buffer lock safety. + auto_h2d: Automatically execute H2D ingestion upon data arrival. """ self._impl = _weight_synchronizer.WeightSynchronizer( - device_tensors, local_port, parallelism, listener_port, bind_ip + device_tensors, + local_port, + parallelism, + listener_port, + bind_ip, + unsafe_skip_buffer_lock, + auto_h2d, ) def push_weights(self, peers: List[str]) -> None: - """Trainer pushing current model weights to peer inference server coordinates E2E.""" + """Trainer pushes model weights to peer inference server coordinates.""" self._impl.PushWeights(peers) def d2h(self) -> None: - """Triggers asynchronous Device-to-Host (D2H) copy of current weights to Host buffer.""" + """Triggers asynchronous D2H copy of current weights to Host buffer.""" self._impl.D2h() def h2d(self) -> None: - """Triggers asynchronous Host-to-Device (H2D) copy of weights from Host buffer to Device.""" + """Triggers asynchronous H2D copy of weights from Host buffer to Device.""" self._impl.H2d() def get_host_buffer( self, layer_idx: int = 0, shard_idx: int = 0 ) -> torch.Tensor: - """Returns a zero-copy Host-side CPU PyTorch Tensor view of the C++ staging buffer. + """Returns a zero-copy Host CPU PyTorch Tensor view of staging buffer. Args: layer_idx: Target layer index to fetch. diff --git a/tpu_sync/frameworks/torch/BUILD b/tpu_sync/frameworks/torch/BUILD index 70c0b2a4..3239b75f 100644 --- a/tpu_sync/frameworks/torch/BUILD +++ b/tpu_sync/frameworks/torch/BUILD @@ -55,10 +55,12 @@ cc_library( linkopts = ["-Wl,--allow-shlib-undefined"], visibility = ["//visibility:public"], deps = [ + # buildcleaner: keep + ":torch_tpu_device_buffer_headers", ":torch_tpu_tensor_to_buffer_headers", ":torch_utils", - "//tpu_sync/weight_sync:weight_synchronizer_base", "@torch_tpu//shims/torch:aten_headers", + "//tpu_sync/weight_sync:weight_synchronizer_base", ], ) diff --git a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc index ed86d710..fd355ece 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc +++ b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc @@ -501,10 +501,12 @@ NB_MODULE(_tpu_raiden_torch, m) { nb::class_(m, "WeightSynchronizer") .def(nb::init>&, std::optional, int, std::optional, - std::optional>(), + std::optional, bool, bool>(), nb::arg("device_tensors"), nb::arg("local_port") = nb::none(), nb::arg("parallelism") = 1, nb::arg("listener_port") = nb::none(), - nb::arg("bind_ip") = nb::none()) + nb::arg("bind_ip") = nb::none(), + nb::arg("unsafe_skip_buffer_lock") = true, + nb::arg("auto_h2d") = false) .def( "PushWeights", [](WeightSynchronizer& self, const std::vector& peers) { diff --git a/tpu_sync/frameworks/torch/weight_synchronizer.cc b/tpu_sync/frameworks/torch/weight_synchronizer.cc index f2ae4770..d4fa1775 100644 --- a/tpu_sync/frameworks/torch/weight_synchronizer.cc +++ b/tpu_sync/frameworks/torch/weight_synchronizer.cc @@ -15,6 +15,7 @@ #include "tpu_sync/frameworks/torch/weight_synchronizer.h" #include +#include #include #include @@ -29,22 +30,21 @@ WeightSynchronizer::WeightSynchronizer( const std::vector>& device_tensors, std::optional local_port, int parallelism, std::optional listener_port, std::optional bind_ip, - bool unsafe_skip_buffer_lock) + bool unsafe_skip_buffer_lock, bool auto_h2d) : WeightSynchronizer( UnpackTorchTensors(device_tensors, unsafe_skip_buffer_lock), local_port, parallelism, listener_port, bind_ip, - unsafe_skip_buffer_lock) {} - -WeightSynchronizer::WeightSynchronizer(UnpackedTensors unpacked, - std::optional local_port, - int parallelism, - std::optional listener_port, - std::optional bind_ip, - bool unsafe_skip_buffer_lock) + unsafe_skip_buffer_lock, auto_h2d) {} + +WeightSynchronizer::WeightSynchronizer( + UnpackedTensors unpacked, std::optional local_port, int parallelism, + std::optional listener_port, std::optional bind_ip, + bool unsafe_skip_buffer_lock, bool auto_h2d) : weight_sync::WeightSynchronizerBase( std::move(unpacked.buffers), local_port, /*external_host_ptrs=*/std::nullopt, unsafe_skip_buffer_lock, - parallelism, listener_port, bind_ip), + parallelism, listener_port, bind_ip, + /*layer_names=*/{}, auto_h2d), buffer_refs_(std::move(unpacked.refs)) {} WeightSynchronizer::~WeightSynchronizer() = default; diff --git a/tpu_sync/frameworks/torch/weight_synchronizer.h b/tpu_sync/frameworks/torch/weight_synchronizer.h index c71bc992..bd402112 100644 --- a/tpu_sync/frameworks/torch/weight_synchronizer.h +++ b/tpu_sync/frameworks/torch/weight_synchronizer.h @@ -16,10 +16,11 @@ #define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_FRAMEWORKS_TORCH_WEIGHT_SYNCHRONIZER_H_ #include +#include #include #include "ATen/core/TensorBody.h" -#include "torch_tpu/eager/tensor_to_buffer.h" +#include "torch_tpu/eager/device_buffer.h" #include "tpu_sync/frameworks/torch/torch_utils.h" #include "tpu_sync/weight_sync/weight_synchronizer_base.h" @@ -34,7 +35,8 @@ class WeightSynchronizer : public weight_sync::WeightSynchronizerBase { int parallelism = 1, std::optional listener_port = std::nullopt, std::optional bind_ip = std::nullopt, - bool unsafe_skip_buffer_lock = true); + bool unsafe_skip_buffer_lock = true, + bool auto_h2d = false); ~WeightSynchronizer() override; @@ -44,7 +46,7 @@ class WeightSynchronizer : public weight_sync::WeightSynchronizerBase { WeightSynchronizer(UnpackedTensors unpacked, std::optional local_port, int parallelism, std::optional listener_port, std::optional bind_ip, - bool unsafe_skip_buffer_lock); + bool unsafe_skip_buffer_lock, bool auto_h2d); std::vector buffer_refs_; };