Added auto_h2d for Torch weight synchronizer.

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",
     ],
 )
 

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.

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",
     ],
 )
 

tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc

@@ -501,10 +501,12 @@ NB_MODULE(_tpu_raiden_torch, m) {
   nb::class_<WeightSynchronizer>(m, "WeightSynchronizer")
       .def(nb::init<const std::vector<std::vector<at::Tensor>>&,
                     std::optional<int>, int, std::optional<int>,
-                    std::optional<std::string>>(),
+                    std::optional<std::string>, 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<std::string>& peers) {

tpu_sync/frameworks/torch/weight_synchronizer.cc

@@ -15,6 +15,7 @@
 #include "tpu_sync/frameworks/torch/weight_synchronizer.h"
 
 #include <optional>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -29,22 +30,21 @@ WeightSynchronizer::WeightSynchronizer(
     const std::vector<std::vector<at::Tensor>>& device_tensors,
     std::optional<int> local_port, int parallelism,
     std::optional<int> listener_port, std::optional<std::string> 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<int> local_port,
-                                       int parallelism,
-                                       std::optional<int> listener_port,
-                                       std::optional<std::string> bind_ip,
-                                       bool unsafe_skip_buffer_lock)
+          unsafe_skip_buffer_lock, auto_h2d) {}
+
+WeightSynchronizer::WeightSynchronizer(
+    UnpackedTensors unpacked, std::optional<int> local_port, int parallelism,
+    std::optional<int> listener_port, std::optional<std::string> 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;

tpu_sync/frameworks/torch/weight_synchronizer.h

@@ -16,10 +16,11 @@
 #define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_FRAMEWORKS_TORCH_WEIGHT_SYNCHRONIZER_H_
 
 #include <optional>
+#include <string>
 #include <vector>
 
 #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<int> listener_port = std::nullopt,
                      std::optional<std::string> 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<int> local_port,
                      int parallelism, std::optional<int> listener_port,
                      std::optional<std::string> bind_ip,
-                     bool unsafe_skip_buffer_lock);
+                     bool unsafe_skip_buffer_lock, bool auto_h2d);
 
   std::vector<torch_tpu::DeviceBufferRef> buffer_refs_;
 };