feat(weight_sync): Raiden weight synchronization with progressive memory reclamation

tests/experimental/weight_sync/raiden_synchronizer_test.py

@@ -294,6 +294,14 @@ def test_host_stage_pulls_state_to_host(self):
       )
     pull.assert_called_once()
 
+  def test_release_host_arrays(self):
+    sync = raiden_synchronizer.RaidenSynchronizer("trainer", self._state())
+    self.assertTrue(sync.bound)
+    self.assertNotEmpty(sync.arrays)
+    sync.release_host_arrays()
+    self.assertEmpty(sync.arrays)
+    self.assertTrue(sync.bound)
+
 
 if __name__ == "__main__":
   absltest.main()

tunix/experimental/weight_sync/raiden_synchronizer.py

@@ -17,6 +17,8 @@
 from __future__ import annotations
 
 import collections
+import gc
+import resource
 import socket
 from typing import Any, List, Optional, Tuple
 
@@ -25,6 +27,14 @@
 import jax.numpy as jnp
 from tunix.experimental.weight_sync import weight_sync
 
+
+def _log_rss(tag: str) -> None:
+  """Logs process peak RSS (GB) -- pinpoints which bind() stage spikes host
+  memory, since ru_maxrss is a high-water mark that only grows.
+  """
+  rss_gb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6
+  logging.info("raiden bind rss checkpoint [%s]: %.1f GB (peak)", tag, rss_gb)
+
 _ws_lib: Any = None
 try:
   from tpu_sync.api.jax import weight_synchronizer as _ws_lib  # pytype: disable=import-error  pylint: disable=g-import-not-at-top
@@ -53,14 +63,24 @@ def local_ip() -> str:
 def to_host_cpu_state(state: Any) -> Any:
   """Pulls arrays to client host memory; proxy arrays cannot bind directly."""
   cpu = jax.local_devices(backend="cpu")[0]
-
-  def pull(leaf):
+  leaves, treedef = jax.tree_util.tree_flatten(state)
+  del state
+  new_leaves = []
+  for i in range(len(leaves)):
+    leaf = leaves[i]
+    leaves[i] = None
     arr = getattr(leaf, "value", leaf)
     if hasattr(arr, "shape") and hasattr(arr, "dtype"):
-      return jax.device_put(jax.device_get(arr), cpu)
-    return leaf
-
-  return jax.tree_util.tree_map(pull, state)
+      new_leaves.append(jax.device_put(jax.device_get(arr), cpu))
+    else:
+      new_leaves.append(leaf)
+    del leaf, arr
+    # Periodically run GC to release Pathways proxy transit buffers incrementally
+    if i % 4 == 3:
+      gc.collect()
+  del leaves
+  gc.collect()
+  return jax.tree_util.tree_unflatten(treedef, new_leaves)
 
 
 def flatten_weights(state: Any) -> Tuple[List[str], List[Any]]:
@@ -75,22 +95,33 @@ def flatten_weights(state: Any) -> Tuple[List[str], List[Any]]:
 
 
 def _bindable(arr: Any) -> bool:
-  """True if the native layer can bind this leaf."""
+  """True if the native layer can bind this leaf.
+
+  Binding an unsupported leaf (e.g. RNG keys) can SIGSEGV, so only
+  floating-point, rank>=1, fully TPU- or CPU-resident arrays qualify. CPU is
+  allowed because host_stage deliberately copies proxy-backed (Pathways)
+  arrays to host CPU memory before bind() gets here -- rejecting "not TPU"
+  would drop every leaf it just staged.
+  """
   try:
+    if not hasattr(arr, "shape") or not hasattr(arr, "dtype"):
+      return False
+    if arr.ndim < 1:
+      return False
+    if not jnp.issubdtype(arr.dtype, jnp.floating):
+      return False
     devices = arr.devices()
-  except AttributeError:
+    if not devices:
+      return False
+    return all(getattr(d, "platform", "?") in ("tpu", "cpu") for d in devices)
+  except Exception:
     return False
-  on_local_hw = all(
-      getattr(d, "platform", "?") in ("cpu", "tpu") for d in devices
-  )
-  return on_local_hw and jnp.issubdtype(arr.dtype, jnp.number)
 
 
 def _filter_bindable(
     names: List[str], arrays: List[Any]
 ) -> Tuple[List[str], List[Any]]:
-  """Drops leaves the native layer cannot bind; binding them is undefined
-  behavior (observed: random RuntimeError or SIGSEGV on RNG-key arrays)."""
+  """Drops leaves _bindable rejects."""
   logging.vlog(
       1,
       "raiden bind census: %s",
@@ -103,6 +134,12 @@ def _filter_bindable(
   dropped = []
   for name, arr in zip(names, arrays):
     if _bindable(arr):
+      if hasattr(arr, "block_until_ready"):
+        try:
+          # binding an in-flight buffer is part of what SIGSEGVs
+          arr.block_until_ready()
+        except Exception:
+          pass
       keep_names.append(name)
       keep_arrays.append(arr)
     else:
@@ -184,14 +221,18 @@ def active(self) -> bool:
     return self._sync is not None
 
   def bind(self, state: Any) -> None:
-    """Binds this host's weights, or rebinds them after a training step.
-
-    With host_stage the arrays are copied to local CPU memory first; arrays
-    backed by the pathways proxy cannot bind in place.
-    """
+    """Binds or rebinds weights to the Raiden transport."""
+    _log_rss("bind:start")
+    # Clear previous buffers before staging to avoid holding duplicate weight
+    # copies in host memory during rebinds.
+    self.names = []
+    self.arrays = []
     if self._host_stage:
       state = to_host_cpu_state(state)
+    _log_rss("bind:after_host_stage")
     self.names, self.arrays = _filter_bindable(*flatten_weights(state))
+    del state
+    _log_rss("bind:after_flatten")
     if _ws_lib is None:
       return
     if self._sync is None:
@@ -206,8 +247,10 @@ def bind(self, state: Any) -> None:
           bind_ip=None,
           auto_h2d=self._auto_h2d,
       )
+      _log_rss("bind:after_native_construct")
     else:
       self._sync.bind_weights(self.arrays)
+      _log_rss("bind:after_native_rebind")
 
   def _require_sync(self, op: str) -> Any:
     if self._sync is None:
@@ -219,6 +262,11 @@ def d2h(self) -> None:
 
   def h2d(self) -> None:
     self._require_sync("h2d()").h2d()
+    jax.block_until_ready(self.arrays)
+
+  def release_host_arrays(self) -> None:
+    """Drops host-staged array references to reclaim memory between sync rounds."""
+    self.arrays = []
 
   def metrics(self) -> dict:
     return self._sync.get_metrics() if self._sync else {}

tunix/experimental/weight_sync/weight_sync.py

@@ -186,6 +186,79 @@ class WorkUnitMetadata:
   variables: tuple[TensorMetadata, ...] = ()
   mesh_axes: Optional[tuple[str, ...]] = None
 
+  @classmethod
+  def from_dict(cls, d: Any) -> WorkUnitMetadata:
+    """Reconstructs WorkUnitMetadata from a dictionary or returns metadata directly."""
+    if isinstance(d, cls):
+      return d
+    if not isinstance(d, dict):
+      raise TypeError(f"Expected WorkUnitMetadata or dict, got {type(d)}")
+
+    unit_raw = d.get("unit")
+    if isinstance(unit_raw, dict):
+      unit = WorkUnitId(**unit_raw)
+    elif isinstance(unit_raw, WorkUnitId):
+      unit = unit_raw
+    else:
+      unit = WorkUnitId(job_name=str(unit_raw or "destination"))
+
+    variables_raw = d.get("variables", ())
+    variables = []
+    for v in variables_raw:
+      if isinstance(v, TensorMetadata):
+        variables.append(v)
+      elif isinstance(v, dict):
+        variables.append(
+            TensorMetadata(
+                name=v["name"],
+                shape=tuple(v["shape"]),
+                mesh_shape=tuple(v["mesh_shape"]),
+                layout=tuple(v["layout"]),
+                item_size=int(v["item_size"]),
+                layer_idx=int(v.get("layer_idx", 0)),
+                sharding_spec=tuple(v.get("sharding_spec", ())),
+            )
+        )
+      elif hasattr(v, "name"):
+        variables.append(
+            TensorMetadata(
+                name=v.name,
+                shape=tuple(v.shape),
+                mesh_shape=tuple(v.mesh_shape),
+                layout=tuple(v.layout),
+                item_size=int(v.item_size),
+                layer_idx=int(getattr(v, "layer_idx", 0)),
+                sharding_spec=tuple(getattr(v, "sharding_spec", ())),
+            )
+        )
+
+    return cls(
+        unit=unit,
+        shards=tuple(d.get("shards", ())),
+        control_plane_rpc_address=str(d.get("control_plane_rpc_address", "")),
+        global_shape=(
+            tuple(d["global_shape"])
+            if d.get("global_shape") is not None
+            else None
+        ),
+        mesh_shape=(
+            tuple(d["mesh_shape"]) if d.get("mesh_shape") is not None else None
+        ),
+        layout=tuple(d["layout"]) if d.get("layout") is not None else None,
+        item_size=(
+            int(d["item_size"]) if d.get("item_size") is not None else None
+        ),
+        variables=tuple(variables),
+        mesh_axes=(
+            tuple(d["mesh_axes"]) if d.get("mesh_axes") is not None else None
+        ),
+    )
+
+
+def dict_to_metadata(d: Any) -> WorkUnitMetadata:
+  """Reconstructs WorkUnitMetadata from a dictionary (delegates to WorkUnitMetadata.from_dict)."""
+  return WorkUnitMetadata.from_dict(d)
+
 
 @dataclasses.dataclass(frozen=True)
 class TransferResult:

tunix/experimental/weight_sync/weight_sync_coordinator.py

@@ -889,8 +889,16 @@ async def record_workers(final_error: str = "") -> None:
             " quiesced; no rollback needed"
         ) from e
 
-      src_metadata = [m for per_source in src_meta_lists for m in per_source]
-      dst_metadata = [m for per_dest in dst_meta_lists for m in per_dest]
+      src_metadata = [
+          weight_sync.dict_to_metadata(m)
+          for per_source in src_meta_lists
+          for m in per_source
+      ]
+      dst_metadata = [
+          weight_sync.dict_to_metadata(m)
+          for per_dest in dst_meta_lists
+          for m in per_dest
+      ]
       if not src_metadata or not dst_metadata:
         failures.append(
             f"metadata: {len(src_metadata)} source, {len(dst_metadata)}"