From 60784877199887b048dce74131786ff28f616ae7 Mon Sep 17 00:00:00 2001 From: khatwanimohit Date: Wed, 26 Aug 2026 20:43:06 +0000 Subject: [PATCH] feat(weight_sync): Raiden weight synchronization with progressive memory reclamation --- .../weight_sync/raiden_synchronizer_test.py | 8 ++ .../weight_sync/raiden_synchronizer.py | 86 +++++++++++++++---- tunix/experimental/weight_sync/weight_sync.py | 73 ++++++++++++++++ .../weight_sync/weight_sync_coordinator.py | 12 ++- 4 files changed, 158 insertions(+), 21 deletions(-) diff --git a/tests/experimental/weight_sync/raiden_synchronizer_test.py b/tests/experimental/weight_sync/raiden_synchronizer_test.py index 156d8c533..d94c01537 100644 --- a/tests/experimental/weight_sync/raiden_synchronizer_test.py +++ b/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() \ No newline at end of file diff --git a/tunix/experimental/weight_sync/raiden_synchronizer.py b/tunix/experimental/weight_sync/raiden_synchronizer.py index 8f5d4405d..cb927da7d 100644 --- a/tunix/experimental/weight_sync/raiden_synchronizer.py +++ b/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 {} diff --git a/tunix/experimental/weight_sync/weight_sync.py b/tunix/experimental/weight_sync/weight_sync.py index 6a13adcae..fd78fcc94 100644 --- a/tunix/experimental/weight_sync/weight_sync.py +++ b/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: diff --git a/tunix/experimental/weight_sync/weight_sync_coordinator.py b/tunix/experimental/weight_sync/weight_sync_coordinator.py index 73f35de35..800034f45 100644 --- a/tunix/experimental/weight_sync/weight_sync_coordinator.py +++ b/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)}"