From dd4021a2a03ee7832af593e1c00151ff9c7adb3b Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Mon, 21 Sep 2026 17:19:56 +0800 Subject: [PATCH] =?UTF-8?q?perf(fastsac):=20weight-snapshot=20transport=20?= =?UTF-8?q?=E2=80=94=20fused=20publish=20+=20CUDA-IPC=20slots=20with=20siz?= =?UTF-8?q?e=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- configs/algo_base/motrix.fastsac.yaml | 8 + .../motrix_rl/fastsac/async_impl/collector.py | 48 +- .../motrix_rl/fastsac/async_impl/learner.py | 7 +- .../src/motrix_rl/fastsac/async_impl/shm.py | 457 ------------------ .../fastsac/async_impl/shm/__init__.py | 53 ++ .../fastsac/async_impl/shm/common.py | 132 +++++ .../motrix_rl/fastsac/async_impl/shm/ring.py | 165 +++++++ .../fastsac/async_impl/shm/weight_channel.py | 423 ++++++++++++++++ .../src/motrix_rl/fastsac/async_impl/train.py | 18 +- .../motrix_rl/fastsac/async_impl/worker.py | 107 +++- motrix_rl/src/motrix_rl/fastsac/config.py | 11 + motrix_rl/src/motrix_rl/fastsac/sync/train.py | 14 +- motrix_rl/tests/test_fastsac_collector.py | 80 ++- motrix_rl/tests/test_rl_sim_backend.py | 10 +- 14 files changed, 989 insertions(+), 544 deletions(-) delete mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py diff --git a/configs/algo_base/motrix.fastsac.yaml b/configs/algo_base/motrix.fastsac.yaml index 2c518919..d43b6a07 100644 --- a/configs/algo_base/motrix.fastsac.yaml +++ b/configs/algo_base/motrix.fastsac.yaml @@ -101,3 +101,11 @@ trainer: # machine's CPU/GPU balance (see issue #62's sweep table). learner_cpu_cores: null collector_cpu_cores: null + # Weight-snapshot transport: "auto" (default) enables CUDA-IPC device slots + # only when learner and collector share one GPU and the actor params reach + # weight_ipc_min_bytes; "on"/"off" force the device/host path ("on" warns + # and falls back to the host path when learner and collector are not on one + # GPU). Keep the values quoted: unquoted on/off parse as booleans in YAML. + weight_ipc: auto + # Minimum flattened parameter bytes for the CUDA-IPC path under "auto". + weight_ipc_min_bytes: 16777216 diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py index f39623e3..2529d448 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py @@ -5,7 +5,7 @@ Holds its own :class:`~motrix_rl.fastsac.networks.Actor` and read-only :class:`~motrix_rl.fastsac.buffer.EmpiricalNormalization`, both refreshed from -the learner via :class:`~motrix_rl.fastsac.async_impl.shm.WeightSnapshot`. Each step +the learner via its :class:`~motrix_rl.fastsac.async_impl.shm.WeightReceiver` endpoint. Each step mirrors the sync collector phase (``agent.py`` collect phase) exactly: decide action -> ``env.step`` -> push the transition batch to the shared ring -> update episode bookkeeping. The normalizer is used read-only (``update=False``), matching @@ -20,7 +20,8 @@ import torch from torch import nn -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot, bind_flat_params +from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, bind_flat_params +from motrix_rl.fastsac.async_impl.shm.weight_channel import WeightReceiver from motrix_rl.fastsac.buffer import EmpiricalNormalization from motrix_rl.fastsac.config import FastSacAgentCfg, FastSacCfg from motrix_rl.fastsac.networks import Actor @@ -76,7 +77,7 @@ def __init__( action_scale: torch.Tensor, action_bias: torch.Tensor, ring: SharedTransitionRing, - weights: WeightSnapshot, + weights: WeightReceiver, control: Control, is_resume: bool = False, ): @@ -95,7 +96,6 @@ def __init__( self.control = control self.is_resume = is_resume self._learning_starts = acfg.learning_starts - self._local_version = 0 self.actor = Actor( n_obs=obs_dim, @@ -125,25 +125,6 @@ def __init__( raise ValueError("collector_amp_dtype must be 'fp16' or 'bf16'") from exc self._flat_params = bind_flat_params(self.actor) if self.device.type == "cuda" else None - param_numel = ( - self._flat_params.numel() - if self._flat_params is not None - else sum(param.numel() for param in self.actor.parameters()) - ) - pin_weight_staging = self.device.type == "cuda" - self._weight_param_staging = torch.empty( - param_numel, - dtype=torch.float32, - pin_memory=pin_weight_staging, - ) - self._weight_normalizer_staging = tuple( - torch.empty( - (1, obs_dim), - dtype=torch.float32, - pin_memory=pin_weight_staging, - ) - for _ in range(3) - ) self._obs_host = None self._obs_device = None self._actions_host = None @@ -228,12 +209,8 @@ def sync_weights(self, *, record_timing: bool = False) -> None: version, wait_writer_s, host_snapshot_s, actor_load_s = self.weights.maybe_load( self.actor, self.obs_normalizer, - self._local_version, - param_staging=self._weight_param_staging, - normalizer_staging=self._weight_normalizer_staging, flat_params=self._flat_params, ) - self._local_version = version if record_timing: self._sync_wait_writer_t += wait_writer_s self._sync_host_snapshot_t += host_snapshot_s @@ -242,7 +219,7 @@ def sync_weights(self, *, record_timing: bool = False) -> None: @property def policy_lag(self) -> int: """How many published versions behind the collector's local policy is.""" - return max(0, self.weights.version - self._local_version) + return self.weights.lag # ------------------------------------------------------------------ ring handoff # ------------------------------------------------------------------ step @@ -292,12 +269,15 @@ def step_once(self) -> bool: self.ep_return += rewards self.ep_len += 1 done_idx = torch.nonzero(terminated | truncated, as_tuple=False).flatten() - for j in done_idx.tolist(): - self.recent_returns.append(float(self.ep_return[j])) - self.recent_lengths.append(float(self.ep_len[j])) - self.ep_return[j] = 0.0 - self.ep_len[j] = 0.0 - self.n_episodes += 1 + if done_idx.numel(): + # Vectorized: one batched gather + clear instead of per-episode + # Python-level scalar indexing — early in training hundreds of + # dones per step make the scalar loop dominate bookkeeping. + self.recent_returns.extend(self.ep_return[done_idx].tolist()) + self.recent_lengths.extend(self.ep_len[done_idx].tolist()) + self.ep_return[done_idx] = 0.0 + self.ep_len[done_idx] = 0.0 + self.n_episodes += int(done_idx.numel()) self.recent_returns = self.recent_returns[-100:] self.recent_lengths = self.recent_lengths[-100:] diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py index d5ce7a3c..186b3d94 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py @@ -7,7 +7,7 @@ :class:`~motrix_rl.fastsac.async_impl.shm.SharedTransitionRing` into the agent's GPU replay buffer, runs gradient updates governed by ``utd_mode`` (§6 of the design), and periodically publishes actor weights + obs-normalizer stats to the -collector via :class:`~motrix_rl.fastsac.async_impl.shm.WeightSnapshot`. +collector via its :class:`~motrix_rl.fastsac.async_impl.shm.WeightSender` endpoint. The update math is reused unchanged from the sync agent: this module delegates the per-step gradient work to ``agent.update(n)`` and only owns the @@ -19,7 +19,8 @@ import time from motrix_rl.fastsac.agent import FastSacAgent -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot +from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.shm.weight_channel import WeightSender from motrix_rl.fastsac.config import FastSacCfg @@ -29,7 +30,7 @@ def __init__( agent: FastSacAgent, cfg: FastSacCfg, ring: SharedTransitionRing, - weights: WeightSnapshot, + weights: WeightSender, control: Control, ): self.agent = agent diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py deleted file mode 100644 index e55422c5..00000000 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py +++ /dev/null @@ -1,457 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Shared-memory primitives for the async FastSAC trainer. - -All tensors are CPU tensors marked ``share_memory_()`` so they can be handed to a -``torch.multiprocessing`` ``spawn`` child in M1 unchanged. In M0 everything runs -in one process; the same objects work in-process, which lets us validate the -collector/learner decomposition before introducing real concurrency. - -Three primitives: - -* :class:`SharedTransitionRing` — single-producer / single-consumer ring of raw - transition batches (collector -> learner) with bounded backpressure. -* :class:`WeightSnapshot` — double-buffered actor weights + obs-normalizer stats - (learner -> collector) guarded by a seqlock so readers always see a complete, - consistent snapshot even when the writer publishes twice during a read. -* :class:`Control` — a few shared scalars (stop flag, global_step, ...). - -Memory ordering ---------------- -All cross-process safety here rests on two invariants: - -1. **8-byte aligned loads/stores of int64 are hardware-atomic** on every ISA we - target (x86-64, ARM64). Each cursor/counter below has a *single writer*, so - an aligned 8-byte store is enough to publish a coherent value — no atomic - RMW primitive (CAS / fetch_add) is needed, and a reader can never see a - torn counter. - -2. **Ordering between the counter and the data holds on strong-memory ISAs.** - The "publish data, then bump counter" pattern is free on x86-64 (TSO: stores - are not reordered with stores, loads not with loads), so a consumer that - reads a bumped cursor is guaranteed to also see the slot's data stores. - - .. warning:: - x86-64 is the **only** platform class currently supported. On weak-memory - ISAs (ARM64 — Jetson, Grace, Apple Silicon) the CPU *may* reorder a data - store after the counter store, letting a consumer observe a bumped cursor - while the slot's data stores are still in flight — a torn snapshot. This - code inserts no memory barriers (there is no portable standalone fence in - the Python stdlib, and the ``atomics`` package offers only ordered - load/store on an ``atomicview``, with no x86 wheel). Supporting ARM would - mean routing every cursor / ``_seq`` store and load through an - ``atomicview`` with ``MemoryOrder.RELEASE`` / ``.ACQUIRE``. -""" - -from __future__ import annotations - -import time - -import torch -from torch import nn - - -def _shared(shape, dtype) -> torch.Tensor: - """Allocate a zero CPU tensor in shared memory.""" - return torch.zeros(shape, dtype=dtype).share_memory_() - - -# ---------------------------------------------------------------- transition ring -class SharedTransitionRing: - """SPSC ring of transition batches with bounded backpressure. - - Each slot holds one env-step batch: the six tensors produced by one - collector step, with the leading dimension being ``num_envs`` (so a slot - is ``(num_envs, dim)``). ``next_obs``/``next_critic_obs`` are NOT stored: - with auto-reset envs the observation returned by step ``t`` is exactly the - stored observation of step ``t+1`` (at episode ends it is the reset - observation), so the consumer derives them from the successive slot via - :meth:`has_next`/:meth:`peek_next`. This halves the per-slot copy volume - and shared-memory footprint. - - Producer (collector) calls :meth:`push`; when the ring is full it returns - ``False`` and the caller must retry/backoff — that is the backpressure that - keeps the collector from outrunning the learner and flooding memory. - - Consumer (learner) calls :meth:`read_slot` to get zero-copy CPU views of the - oldest unread slot plus :meth:`peek_next` for the derived next-observation - views, moves them to its device, then calls :meth:`commit_read`. The read - cursor only advances after the copy, so the producer can never clobber a - slot that is still being ingested (``push`` blocks while the ring is full). - - Memory ordering - ~~~~~~~~~~~~~~~ - Only ``_write`` is read by the consumer, only ``_read`` is read by the - producer, and each cursor has a single writer — so an aligned int64 store - is enough to publish progress and no atomic RMW is needed. Correctness also - needs the consumer to not observe the ``_write`` bump before the slot's data - stores have landed (and symmetrically for ``_read``); on x86/TSO that - ordering is free, so no memory barrier is used and this path is x86-only - (see the module "Memory ordering" note). The same guarantee covers - :meth:`peek_next`: the producer wrote the successor slot's data before - publishing it, which is a precondition of the consumer seeing it committed. - """ - - FIELDS = ( - "obs", - "critic_obs", - "actions", - "rewards", - "dones", - "truncations", - ) - - def __init__( - self, - capacity: int, - num_envs: int, - obs_dim: int, - critic_obs_dim: int, - act_dim: int, - ): - # The consumer derives next_obs from the successor slot, so capacity 1 - # can never satisfy has_next() and would deadlock the pipeline. - if capacity < 2: - raise ValueError(f"ring_capacity must be >= 2 (consumer reads the successor slot), got {capacity}") - self.capacity = capacity - self.num_envs = num_envs - f32, i64 = torch.float32, torch.int64 - self.obs = _shared((capacity, num_envs, obs_dim), f32) - self.critic_obs = _shared((capacity, num_envs, critic_obs_dim), f32) - self.actions = _shared((capacity, num_envs, act_dim), f32) - self.rewards = _shared((capacity, num_envs), f32) - self.dones = _shared((capacity, num_envs), i64) - self.truncations = _shared((capacity, num_envs), i64) - # cursors are shared so the two processes see each other's progress. - self._write = _shared((1,), i64) - self._read = _shared((1,), i64) - - @property - def write_idx(self) -> int: - # Consumer reads this; single-writer producer, so a plain aligned int64 - # load is coherent. On x86/TSO the data loads that follow cannot be - # reordered ahead of it, so no acquire barrier is needed (x86-only). - return int(self._write[0]) - - @property - def read_idx(self) -> int: - # Producer reads this; single-writer consumer. On x86/TSO the following - # is_full() decision is based on an up-to-date value without a barrier. - return int(self._read[0]) - - def size(self) -> int: - """Number of unread slots currently buffered.""" - return self.write_idx - self.read_idx - - def is_full(self) -> bool: - return self.size() >= self.capacity - - def push(self, obs, critic_obs, actions, rewards, dones, truncations) -> bool: - """Copy one env-step batch into the next slot. Returns False if full. - - Inputs are CPU tensors shaped ``(num_envs, dim)``; ``dones``/``truncations`` - are int64 to match ``SimpleReplayBuffer.extend`` semantics. - """ - if self.is_full(): - return False - slot = self.write_idx % self.capacity - self.obs[slot].copy_(obs) - self.critic_obs[slot].copy_(critic_obs) - self.actions[slot].copy_(actions) - self.rewards[slot].copy_(rewards) - self.dones[slot].copy_(dones) - self.truncations[slot].copy_(truncations) - # Publish the slot. On x86/TSO the field copies above are guaranteed - # visible before this cursor bump, so a consumer that reads the new - # write_idx also sees the data (x86-only; ARM would need a release here). - self._write[0] += 1 - return True - - def read_slot(self): - """Return CPU views of the oldest unread slot, or ``None`` if empty. - - Does NOT advance the read cursor; call :meth:`commit_read` after the - consumer has finished copying the data elsewhere. - """ - if self.size() <= 0: - return None - slot = self.read_idx % self.capacity - return ( - self.obs[slot], - self.critic_obs[slot], - self.actions[slot], - self.rewards[slot], - self.dones[slot], - self.truncations[slot], - ) - - def has_next(self) -> bool: - """Whether the successor of the oldest unread slot is already committed. - - The consumer needs it to derive ``next_obs``/``next_critic_obs`` for the - oldest slot (see :meth:`peek_next`), so it must wait for one extra - committed slot before ingesting. - """ - return self.size() > 1 - - def peek_next(self): - """Return CPU views of the successor slot's obs/critic_obs. - - Valid only when :meth:`has_next` is true; the views alias ring memory - that stays untouched until the consumer's own ``commit_read`` calls - advance past it. - """ - slot = (self.read_idx + 1) % self.capacity - return self.obs[slot], self.critic_obs[slot] - - def commit_read(self) -> None: - # Free the slot. On x86/TSO our reads above complete before this cursor - # bump, so the producer's is_full() cannot reuse a slot we are still - # copying out (x86-only; ARM would need a release here). - self._read[0] += 1 - - -# ---------------------------------------------------------------- weight snapshot -def flatten_params(module: nn.Module) -> torch.Tensor: - """Flatten a module's parameters into a single CPU float vector (in order).""" - return torch.cat([p.detach().reshape(-1).float() for p in module.parameters()]).cpu() - - -def load_flat_params(module: nn.Module, flat: torch.Tensor) -> None: - """Inverse of :func:`flatten_params`; copies a flat vector into the params.""" - offset = 0 - for p in module.parameters(): - n = p.numel() - p.data.copy_(flat[offset : offset + n].view_as(p).to(p.device)) - offset += n - - -def bind_flat_params(module: nn.Module) -> torch.Tensor: - """Bind all module parameters to views of one device-contiguous flat tensor. - - The collector owns an inference-only actor, so its parameters do not need - optimizer storage. Binding them once lets every later CPU snapshot version - reach CUDA through one H2D copy instead of one transfer per parameter. - """ - params = list(module.parameters()) - flat = torch.empty(sum(p.numel() for p in params), dtype=params[0].dtype, device=params[0].device) - offset = 0 - with torch.no_grad(): - for param in params: - view = flat[offset : offset + param.numel()].view_as(param) - view.copy_(param) - param.data = view - offset += param.numel() - return flat - - -# obs-normalizer stat buffers that the collector needs (read-only) to reproduce -# the sync ``act()`` path (normalize with update=False, see agent.act). -_NORM_KEYS = ("_mean", "_std", "_var", "count") - - -class WeightSnapshot: - """Double-buffered actor weights + obs-normalizer stats, learner -> collector. - - The naive double buffer ("writer publishes to the other slot, reader reads - the current one") is **not** race-free on its own: if the learner publishes - twice while a collector is mid-``maybe_load``, the second publish reuses the - slot the collector is still copying from — producing a torn snapshot that - then drives the policy for thousands of steps. Hard to reproduce, expensive - to debug. - - We guard the buffer with a **seqlock**: a shared counter ``_seq`` that the - writer bumps to *odd* while writing and back to *even* when done. Readers - spin until they observe a stable even value that did not change across the - read, retrying on any inconsistency. Because the only writer is the learner - (one process), seqlock gives us: - - * **Lock-free, CAS-free reads.** Plain aligned int64 loads/stores suffice; - no atomic RMW primitive is needed. Readers almost never retry. - * **No collision on the common path.** A single publish during a read - targets the *other* slot, so the seq value doesn't even change and the - reader's read is consistent on the first try. - * **Safe detection of the rare race.** A second publish during a read bumps - ``_seq`` twice; the reader notices (``s1 != s2``) and retries from - scratch, never exposing a torn snapshot to the actor. - - Memory ordering - ~~~~~~~~~~~~~~~ - Correctness needs the writer's slot-data stores visible before the seq - change, and the reader's data loads not hoisted above the seq read. On - x86/TSO that holds for free, so no memory barrier is used and this path is - x86-only (see the module "Memory ordering" note); ARM would need real - release/acquire around the seq bumps and reads. - """ - - def __init__(self, param_numel: int, obs_dim: int): - self.param_numel = param_numel - # Two slots for params + normalizer stats (double buffer). - self._params = [_shared((param_numel,), torch.float32) for _ in range(2)] - self._mean = [_shared((1, obs_dim), torch.float32) for _ in range(2)] - self._std = [_shared((1, obs_dim), torch.float32) for _ in range(2)] - self._var = [_shared((1, obs_dim), torch.float32) for _ in range(2)] - self._count = [_shared((1,), torch.int64) for _ in range(2)] - # Seqlock counter. Even = quiescent, odd = write in progress. - # One writer (learner) → aligned int64 stores suffice; the seq value - # itself can never tear. Publish/observe ordering relies on x86/TSO (see - # the module "Memory ordering" note; no barriers, x86-only). - # Public version = seq // 2; starts at 0 (matching the collector's - # ``_local_version = 0`` initial state so the first publish is seen). - self._seq = _shared((1,), torch.int64) - - @property - def version(self) -> int: - """Number of completed publishes (= seq // 2). Best-effort; readers - that need a consistent snapshot must go through :meth:`maybe_load`.""" - return int(self._seq[0]) // 2 - - def publish(self, actor: nn.Module, obs_normalizer: nn.Module) -> None: - """Write current actor params + normalizer stats to the inactive slot - and flip the active pointer via the seqlock.""" - # Materialize the learner's CUDA state on CPU before making the seqlock - # odd. Collector readers may keep using the previous complete version - # while these device-to-host copies finish. - params = flatten_params(actor) - has_stats = all(hasattr(obs_normalizer, k) for k in _NORM_KEYS) - if has_stats: - mean = obs_normalizer._mean.detach().cpu() - std = obs_normalizer._std.detach().cpu() - var = obs_normalizer._var.detach().cpu() - count = int(obs_normalizer.count) - - prev = int(self._seq[0]) - # 1) Mark "write in progress" (odd). Readers seeing this retry. - self._seq[0] = prev + 1 - # 2) Write to the slot opposite the previously-active one - # (active slot before this publish was (prev // 2) % 2). Two - # publishes in a row therefore alternate slots, so a single - # publish during a read targets the *other* slot — no collision - # even without the seqlock guard. The guard exists for the case - # of two publishes during one read. - slot = ((prev // 2) + 1) % 2 - self._params[slot].copy_(params) - if has_stats: - self._mean[slot].copy_(mean) - self._std[slot].copy_(std) - self._var[slot].copy_(var) - self._count[slot][0] = count - # 3) Bump seq back to even. On x86/TSO the data stores above are visible - # before this store, so a reader that observes the new even value - # sees all writes above. Readers that observed the odd value retry. - self._seq[0] = prev + 2 - - def maybe_load( - self, - actor: nn.Module, - obs_normalizer: nn.Module, - local_version: int, - *, - param_staging: torch.Tensor, - normalizer_staging: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - flat_params: torch.Tensor | None = None, - ) -> tuple[int, float, float, float]: - """Load the latest snapshot into ``actor``/``obs_normalizer`` if newer. - - Returns the version actually loaded (== ``local_version`` if nothing - new), followed by wall-clock seconds spent waiting for an in-progress - writer, copying a stable host snapshot, and loading it onto the actor's - device. The breakdown lets the collector distinguish shared-memory - publication contention from actor-load enqueue work. - - Implements the seqlock read loop: read ``_seq``, ensure it is even and - stable, copy the data out, then re-check ``_seq``. If the writer - published in between we discard the (possibly torn) copy and retry. - In practice retries are extremely rare — publish cadence is bounded by - ``weight_publish_interval`` gradient steps, while a read is one CPU - memcpy. - """ - wait_writer_s = 0.0 - host_snapshot_s = 0.0 - while True: - s1 = int(self._seq[0]) - if s1 & 1: # writer is mid-publish; wait for a completed version. - wait_start = time.perf_counter() - while s1 & 1: - s1 = int(self._seq[0]) - wait_writer_s += time.perf_counter() - wait_start - version = s1 // 2 - if version <= local_version: # nothing new to load - return local_version, wait_writer_s, host_snapshot_s, 0.0 - slot = version % 2 # active slot at this seq - # Snapshot the slot into local tensors first. We do not load - # directly into the actor because load_flat_params performs - # many small per-param copies that could each see a different - # seq state; we want one consistent view of the whole slot. - snapshot_start = time.perf_counter() - param_staging.copy_(self._params[slot]) - has_stats = all(hasattr(obs_normalizer, k) for k in _NORM_KEYS) - if has_stats: - mean, std, var = normalizer_staging - mean.copy_(self._mean[slot]) - std.copy_(self._std[slot]) - var.copy_(self._var[slot]) - count = int(self._count[slot][0]) - # Re-check seq. If it changed, the writer published (at least - # once) during our copy and the slot may have been overwritten - # — discard and retry. This is the line that closes the race - # the old double-buffer had. - s2 = int(self._seq[0]) - host_snapshot_s += time.perf_counter() - snapshot_start - if s1 != s2: - continue - - actor_load_start = time.perf_counter() - if flat_params is None: - load_flat_params(actor, param_staging) - else: - flat_params.copy_(param_staging, non_blocking=True) - if has_stats: - obs_normalizer._mean.copy_(mean, non_blocking=True) - obs_normalizer._std.copy_(std, non_blocking=True) - obs_normalizer._var.copy_(var, non_blocking=True) - obs_normalizer.count.fill_(count) - # CUDA collectors keep the host staging buffers alive. The copies - # above are ordered before the next inference on the same stream; - # that inference already synchronizes after returning actions to - # the CPU environment, so a separate weight-load barrier only - # serializes collector and learner work unnecessarily. - actor_load_s = time.perf_counter() - actor_load_start - return version, wait_writer_s, host_snapshot_s, actor_load_s - - -# ---------------------------------------------------------------- control block -class Control: - """A handful of shared scalar controls / counters.""" - - def __init__(self): - self._stop = _shared((1,), torch.int64) - self._global_step = _shared((1,), torch.int64) # learner iteration counter - self._collector_steps = _shared((1,), torch.int64) # env-step batches produced - - @property - def stop(self) -> bool: - return bool(self._stop[0]) - - def set_stop(self) -> None: - self._stop[0] = 1 - - @property - def global_step(self) -> int: - return int(self._global_step[0]) - - @global_step.setter - def global_step(self, v: int) -> None: - self._global_step[0] = v - - @property - def collector_steps(self) -> int: - return int(self._collector_steps[0]) - - @collector_steps.setter - def collector_steps(self, v: int) -> None: - self._collector_steps[0] = v - - def inc_collector_steps(self) -> None: - self._collector_steps[0] += 1 diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py new file mode 100644 index 00000000..b6275c77 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py @@ -0,0 +1,53 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Shared-memory primitives for the async FastSAC trainer. + +This package is the established internal API surface of the former ``shm`` +module; the facade keeps the ``motrix_rl.fastsac.async_impl.shm`` import path +stable across the worker/collector/learner modules and tests. Import from the +defining submodules for new internal uses; import from here only when relying +on the stable package namespace. + +Submodules: + +* :mod:`.common` — the shared-memory allocator, shared scalars + (:class:`Control`) and flat-parameter helpers. +* :mod:`.ring` — the SPSC transition ring (collector -> learner). +* :mod:`.weight_channel` — the seqlock weight channel (learner -> collector) + with host-shm and CUDA-IPC endpoint implementations. +""" + +from motrix_rl.fastsac.async_impl.shm.common import ( + Control, + bind_flat_params, + flatten_params, + load_flat_params, +) +from motrix_rl.fastsac.async_impl.shm.ring import SharedTransitionRing +from motrix_rl.fastsac.async_impl.shm.weight_channel import ( + GpuIpcWeightReceiver, + GpuIpcWeightSender, + HostWeightReceiver, + HostWeightSender, + WeightChannelShared, + WeightReceiver, + WeightSender, + weight_receiver_for, +) + +__all__ = [ + "Control", + "GpuIpcWeightReceiver", + "GpuIpcWeightSender", + "HostWeightReceiver", + "HostWeightSender", + "SharedTransitionRing", + "WeightChannelShared", + "WeightReceiver", + "WeightSender", + "weight_receiver_for", + "bind_flat_params", + "flatten_params", + "load_flat_params", +] diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py new file mode 100644 index 00000000..39189914 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py @@ -0,0 +1,132 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Shared-memory primitives for the async FastSAC trainer. + +All tensors are CPU tensors marked ``share_memory_()`` so they can be handed to a +``torch.multiprocessing`` ``spawn`` child in M1 unchanged. In M0 everything runs +in one process; the same objects work in-process, which lets us validate the +collector/learner decomposition before introducing real concurrency. + +Three primitives: + +* :class:`SharedTransitionRing` — single-producer / single-consumer ring of raw + transition batches (collector -> learner) with bounded backpressure. +* the weight channel (see :mod:`motrix_rl.fastsac.async_impl.shm.weight_channel`) + — double-buffered actor weights + obs-normalizer stats (learner -> + collector) guarded by a seqlock so readers always see a complete, + consistent snapshot even when the writer publishes twice during a read. +* :class:`Control` — a few shared scalars (stop flag, global_step, ...). + +Memory ordering +--------------- +All cross-process safety here rests on two invariants: + +1. **8-byte aligned loads/stores of int64 are hardware-atomic** on every ISA we + target (x86-64, ARM64). Each cursor/counter below has a *single writer*, so + an aligned 8-byte store is enough to publish a coherent value — no atomic + RMW primitive (CAS / fetch_add) is needed, and a reader can never see a + torn counter. + +2. **Ordering between the counter and the data holds on strong-memory ISAs.** + The "publish data, then bump counter" pattern is free on x86-64 (TSO: stores + are not reordered with stores, loads not with loads), so a consumer that + reads a bumped cursor is guaranteed to also see the slot's data stores. + + .. warning:: + x86-64 is the **only** platform class currently supported. On weak-memory + ISAs (ARM64 — Jetson, Grace, Apple Silicon) the CPU *may* reorder a data + store after the counter store, letting a consumer observe a bumped cursor + while the slot's data stores are still in flight — a torn snapshot. This + code inserts no memory barriers (there is no portable standalone fence in + the Python stdlib, and the ``atomics`` package offers only ordered + load/store on an ``atomicview``, with no x86 wheel). Supporting ARM would + mean routing every cursor / ``_seq`` store and load through an + ``atomicview`` with ``MemoryOrder.RELEASE`` / ``.ACQUIRE``. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +def _shared(shape, dtype) -> torch.Tensor: + """Allocate a zero-initialized CPU tensor in shared memory (usable across spawn children).""" + return torch.zeros(shape, dtype=dtype).share_memory_() + + +# ---------------------------------------------------------------- control block +class Control: + """A handful of shared scalar controls / counters.""" + + def __init__(self): + self._stop = _shared((1,), torch.int64) + self._global_step = _shared((1,), torch.int64) # learner iteration counter + self._collector_steps = _shared((1,), torch.int64) # env-step batches produced + + @property + def stop(self) -> bool: + return bool(self._stop[0]) + + def set_stop(self) -> None: + self._stop[0] = 1 + + @property + def global_step(self) -> int: + return int(self._global_step[0]) + + @global_step.setter + def global_step(self, v: int) -> None: + self._global_step[0] = v + + @property + def collector_steps(self) -> int: + return int(self._collector_steps[0]) + + @collector_steps.setter + def collector_steps(self, v: int) -> None: + self._collector_steps[0] = v + + def inc_collector_steps(self) -> None: + self._collector_steps[0] += 1 + + +# ---------------------------------------------------------------- flat-param helpers +def flatten_params(module: nn.Module) -> torch.Tensor: + """Flatten module parameters into one contiguous float vector on the module's device.""" + return torch.cat([p.detach().reshape(-1).float() for p in module.parameters()]) + + +def load_flat_params(module: nn.Module, flat: torch.Tensor) -> None: + """Inverse of :func:`flatten_params`; copies a flat vector into the params.""" + offset = 0 + for p in module.parameters(): + n = p.numel() + p.data.copy_(flat[offset : offset + n].view_as(p).to(p.device)) + offset += n + + +def bind_flat_params(module: nn.Module) -> torch.Tensor: + """Bind all module parameters to views of one device-contiguous flat tensor. + + The collector owns an inference-only actor, so its parameters do not need + optimizer storage. Binding them once lets every later CPU snapshot version + reach CUDA through one H2D copy instead of one transfer per parameter. + """ + params = list(module.parameters()) + flat = torch.empty(sum(p.numel() for p in params), dtype=params[0].dtype, device=params[0].device) + offset = 0 + with torch.no_grad(): + for param in params: + view = flat[offset : offset + param.numel()].view_as(param) + view.copy_(param) + param.data = view + offset += param.numel() + return flat + + +# obs-normalizer stat buffers that the collector needs (read-only) to reproduce +# the sync ``act()`` path (normalize with update=False, see agent.act); consumed +# by the weight channel (weight_channel.py). +_NORM_KEYS = ("_mean", "_std", "_var", "count") diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py new file mode 100644 index 00000000..1a6007c6 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py @@ -0,0 +1,165 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""SPSC transition ring between the collector and learner processes.""" + +from __future__ import annotations + +import torch + +from motrix_rl.fastsac.async_impl.shm.common import _shared + + +# ---------------------------------------------------------------- transition ring +class SharedTransitionRing: + """SPSC ring of transition batches with bounded backpressure. + + Each slot holds one env-step batch: the six tensors produced by one + collector step, with the leading dimension being ``num_envs`` (so a slot + is ``(num_envs, dim)``). ``next_obs``/``next_critic_obs`` are NOT stored: + with auto-reset envs the observation returned by step ``t`` is exactly the + stored observation of step ``t+1`` (at episode ends it is the reset + observation), so the consumer derives them from the successive slot via + :meth:`has_next`/:meth:`peek_next`. This halves the per-slot copy volume + and shared-memory footprint. + + Producer (collector) calls :meth:`push`; when the ring is full it returns + ``False`` and the caller must retry/backoff — that is the backpressure that + keeps the collector from outrunning the learner and flooding memory. + + Consumer (learner) calls :meth:`read_slot` to get zero-copy CPU views of the + oldest unread slot plus :meth:`peek_next` for the derived next-observation + views, moves them to its device, then calls :meth:`commit_read`. The read + cursor only advances after the copy, so the producer can never clobber a + slot that is still being ingested (``push`` blocks while the ring is full). + + Memory ordering + ~~~~~~~~~~~~~~~ + Only ``_write`` is read by the consumer, only ``_read`` is read by the + producer, and each cursor has a single writer — so an aligned int64 store + is enough to publish progress and no atomic RMW is needed. Correctness also + needs the consumer to not observe the ``_write`` bump before the slot's data + stores have landed (and symmetrically for ``_read``); on x86/TSO that + ordering is free, so no memory barrier is used and this path is x86-only + (see the module "Memory ordering" note). The same guarantee covers + :meth:`peek_next`: the producer wrote the successor slot's data before + publishing it, which is a precondition of the consumer seeing it committed. + """ + + FIELDS = ( + "obs", + "critic_obs", + "actions", + "rewards", + "dones", + "truncations", + ) + + def __init__( + self, + capacity: int, + num_envs: int, + obs_dim: int, + critic_obs_dim: int, + act_dim: int, + ): + # The consumer derives next_obs from the successor slot, so capacity 1 + # can never satisfy has_next() and would deadlock the pipeline. + if capacity < 2: + raise ValueError(f"ring_capacity must be >= 2 (consumer reads the successor slot), got {capacity}") + self.capacity = capacity + self.num_envs = num_envs + f32, i64 = torch.float32, torch.int64 + self.obs = _shared((capacity, num_envs, obs_dim), f32) + self.critic_obs = _shared((capacity, num_envs, critic_obs_dim), f32) + self.actions = _shared((capacity, num_envs, act_dim), f32) + self.rewards = _shared((capacity, num_envs), f32) + self.dones = _shared((capacity, num_envs), i64) + self.truncations = _shared((capacity, num_envs), i64) + # cursors are shared so the two processes see each other's progress. + self._write = _shared((1,), i64) + self._read = _shared((1,), i64) + + @property + def write_idx(self) -> int: + # Consumer reads this; single-writer producer, so a plain aligned int64 + # load is coherent. On x86/TSO the data loads that follow cannot be + # reordered ahead of it, so no acquire barrier is needed (x86-only). + return int(self._write[0]) + + @property + def read_idx(self) -> int: + # Producer reads this; single-writer consumer. On x86/TSO the following + # is_full() decision is based on an up-to-date value without a barrier. + return int(self._read[0]) + + def size(self) -> int: + """Number of unread slots currently buffered.""" + return self.write_idx - self.read_idx + + def is_full(self) -> bool: + return self.size() >= self.capacity + + def push(self, obs, critic_obs, actions, rewards, dones, truncations) -> bool: + """Copy one env-step batch into the next slot. Returns False if full. + + Inputs are CPU tensors shaped ``(num_envs, dim)``; ``dones``/``truncations`` + are int64 to match ``SimpleReplayBuffer.extend`` semantics. + """ + if self.is_full(): + return False + slot = self.write_idx % self.capacity + self.obs[slot].copy_(obs) + self.critic_obs[slot].copy_(critic_obs) + self.actions[slot].copy_(actions) + self.rewards[slot].copy_(rewards) + self.dones[slot].copy_(dones) + self.truncations[slot].copy_(truncations) + # Publish the slot. On x86/TSO the field copies above are guaranteed + # visible before this cursor bump, so a consumer that reads the new + # write_idx also sees the data (x86-only; ARM would need a release here). + self._write[0] += 1 + return True + + def read_slot(self): + """Return CPU views of the oldest unread slot, or ``None`` if empty. + + Does NOT advance the read cursor; call :meth:`commit_read` after the + consumer has finished copying the data elsewhere. + """ + if self.size() <= 0: + return None + slot = self.read_idx % self.capacity + return ( + self.obs[slot], + self.critic_obs[slot], + self.actions[slot], + self.rewards[slot], + self.dones[slot], + self.truncations[slot], + ) + + def has_next(self) -> bool: + """Whether the successor of the oldest unread slot is already committed. + + The consumer needs it to derive ``next_obs``/``next_critic_obs`` for the + oldest slot (see :meth:`peek_next`), so it must wait for one extra + committed slot before ingesting. + """ + return self.size() > 1 + + def peek_next(self): + """Return CPU views of the successor slot's obs/critic_obs. + + Valid only when :meth:`has_next` is true; the views alias ring memory + that stays untouched until the consumer's own ``commit_read`` calls + advance past it. + """ + slot = (self.read_idx + 1) % self.capacity + return self.obs[slot], self.critic_obs[slot] + + def commit_read(self) -> None: + # Free the slot. On x86/TSO our reads above complete before this cursor + # bump, so the producer's is_full() cannot reuse a slot we are still + # copying out (x86-only; ARM would need a release here). + self._read[0] += 1 diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py new file mode 100644 index 00000000..6dbac0a5 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py @@ -0,0 +1,423 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Single-producer / single-consumer weight channel (learner -> collector). + +The learner publishes actor weights + obs-normalizer stats; the collector +loads the latest published version into its inference-only actor copy. The +protocol state (normalizer-stat double buffers + seqlock counter) is created +once by the parent process in :class:`WeightChannelShared` and pickled into +both worker processes, where the concrete endpoints are constructed: + +* the learner builds a :class:`WeightSender` subclass — host shared-memory + slots anywhere, or CUDA-IPC device slots in the learner process (which owns + the CUDA context and must keep the IPC tensors alive) — and the orchestration + ships the slot tensors to the collector over a one-shot queue; +* the collector builds the matching :class:`WeightReceiver` subclass from the + arrived tensors (:func:`weight_receiver_for` dispatches on their device). + +The endpoints themselves know nothing about queues. + +Seqlock protocol +---------------- +The naive double buffer ("writer publishes to the other slot, reader reads +the current one") is **not** race-free on its own: if the learner publishes +twice while the collector is mid-read, the second publish reuses the slot the +collector is still copying from — producing a torn snapshot that then drives +the policy for thousands of steps. The buffer is guarded by a **seqlock**: a +shared counter that the writer bumps to *odd* while writing and back to +*even* when done; readers retry on any inconsistency. Single writer + single +reader means plain aligned int64 stores suffice and no atomic RMW is needed. +On x86/TSO the data-before-counter ordering is free; ARM would need real +release/acquire (see the shm module's "Memory ordering" note). + +Transport choice +---------------- +Host shared memory is the default: it is sub-millisecond for small actors and +avoids any GPU synchronization. CUDA-IPC device slots win only for large +actors — both directions pay a stream/event synchronization to preserve the +seqlock's "data visible before version bump" ordering, which is a net win +once the avoided host transfer dominates that barrier (the orchestration +gates this on async_options.weight_ipc and the actor parameter size). +""" + +from __future__ import annotations + +import time +from abc import ABC, abstractmethod + +import torch +import torch.multiprocessing # noqa: F401 registers CUDA-IPC reducers in every importing process +from torch import nn + +from motrix_rl.fastsac.async_impl.shm.common import _NORM_KEYS, _shared, flatten_params, load_flat_params + +# obs-normalizer stat buffers that the collector needs (read-only) to reproduce +# the sync ``act()`` path (normalize with update=False, see agent.act). + + +class WeightChannelShared: + """Parent-created protocol state shared by both channel endpoints.""" + + def __init__(self, obs_dim: int): + # Two slots for normalizer stats (double buffer) + the seqlock counter. + self.mean = [_shared((1, obs_dim), torch.float32) for _ in range(2)] + self.std = [_shared((1, obs_dim), torch.float32) for _ in range(2)] + self.var = [_shared((1, obs_dim), torch.float32) for _ in range(2)] + self.count = [_shared((1,), torch.int64) for _ in range(2)] + # Even = quiescent, odd = write in progress. Public version = seq // 2; + # starts at 0 (matching the receiver's ``loaded_version = 0`` initial + # state so the first publish is seen). + self.seq = _shared((1,), torch.int64) + + +class WeightSender(ABC): + """Learner-side endpoint; publishes snapshots under the seqlock.""" + + def __init__(self, shared: WeightChannelShared): + self.shared = shared + + @property + def version(self) -> int: + """Number of completed publishes (= seq // 2).""" + return int(self.shared.seq[0]) // 2 + + #: The double-buffered slot tensors. Construction ships nothing by itself; + #: the orchestration sends these to the collector to build its receiver. + params: list[torch.Tensor] + + @abstractmethod + def _write_slot(self, slot: int, device_flat: torch.Tensor) -> None: + """Start writing the flattened params into a slot (may be asynchronous).""" + + @abstractmethod + def _complete_slot_write(self) -> None: + """Block until the last :meth:`_write_slot` has fully landed. + + Splitting the write lets the GPU copy overlap with the normalizer-stat + CPU writes in between; completion must still precede the seq-even bump + (a published version implies landed data — that is the seqlock's + publish contract). + """ + + def publish(self, actor: nn.Module, obs_normalizer: nn.Module) -> None: + """Write current actor params + normalizer stats to the inactive slot + and flip the active pointer via the seqlock.""" + device_flat = flatten_params(actor) + has_stats = all(hasattr(obs_normalizer, k) for k in _NORM_KEYS) + if has_stats: + mean = obs_normalizer._mean.detach().cpu() + std = obs_normalizer._std.detach().cpu() + var = obs_normalizer._var.detach().cpu() + count = int(obs_normalizer.count) + + prev = int(self.shared.seq[0]) + # 1) Mark "write in progress" (odd). Readers seeing this retry. + self.shared.seq[0] = prev + 1 + # 2) Write to the slot opposite the previously-active one + # (active slot before this publish was (prev // 2) % 2). Two + # publishes in a row therefore alternate slots, so a single + # publish during a read targets the *other* slot — no collision + # even without the seqlock guard. The guard exists for the case + # of two publishes during one read. + slot = ((prev // 2) + 1) % 2 + # Start the slot write; on GPU transports this only enqueues, letting + # the copy run under the normalizer-stat CPU writes below. + self._write_slot(slot, device_flat) + if has_stats: + self.shared.mean[slot].copy_(mean) + self.shared.std[slot].copy_(std) + self.shared.var[slot].copy_(var) + self.shared.count[slot][0] = count + # Land the slot write before announcing the version: a reader that + # observes the new even seq value must be guaranteed to see the data. + self._complete_slot_write() + # 3) Bump seq back to even. On x86/TSO the data stores above are visible + # before this store, so a reader that observes the new even value + # sees all writes above. Readers that observed the odd value retry. + self.shared.seq[0] = prev + 2 + + +class HostWeightSender(WeightSender): + """Host shared-memory slots; one fused pinned device->host transfer per publish.""" + + def __init__(self, shared: WeightChannelShared, param_numel: int): + super().__init__(shared) + self.params = [_shared((param_numel,), torch.float32) for _ in range(2)] + # Writer-private pinned staging, allocated lazily and keyed by size. + self._staging_cache: tuple[int, torch.Tensor] | None = None + + def _staging(self, device_flat: torch.Tensor) -> torch.Tensor: + size = device_flat.numel() + if self._staging_cache is None or self._staging_cache[0] != size: + pinned = device_flat.is_cuda + self._staging_cache = (size, torch.empty(size, dtype=torch.float32, pin_memory=pinned)) + return self._staging_cache[1] + + def _write_slot(self, slot: int, device_flat: torch.Tensor) -> None: + staging = self._staging(device_flat) + # Keep the source device (not just cuda-ness) so the completion sync + # below targets the stream the D2H copy was actually enqueued on. + self._pending = (slot, staging, device_flat.device if device_flat.is_cuda else None) + staging.copy_(device_flat, non_blocking=True) + + def _complete_slot_write(self) -> None: + # The shared-memory copy reads the staging buffer on the CPU, so a CUDA + # source's D2H transfer must land first. (Pinned-ness of the staging + # buffer is not the discriminator: staging is always a CPU tensor.) + # The device is passed explicitly: the process's current device may + # differ from the device the copy ran on. + slot, staging, source_device = self._pending + if source_device is not None: + torch.cuda.current_stream(source_device).synchronize() + self.params[slot].copy_(staging) + + +class GpuIpcWeightSender(WeightSender): + """CUDA-IPC device slots; publish is a device-to-device copy. + + Must be constructed in the learner process: the exporter of IPC handles + needs a CUDA context and must keep the slot tensors alive for the process + lifetime. The endpoint itself is not shippable — it holds a CUDA event; + only the ``params`` tensors cross processes. + """ + + def __init__(self, shared: WeightChannelShared, param_numel: int, device: torch.device): + super().__init__(shared) + if device.type != "cuda": + raise ValueError(f"GPU weight slots require a CUDA device, got {device}") + self.params = [torch.zeros(param_numel, device=device) for _ in range(2)] + # The event binds to the device current at creation time; create it + # under the slot device so record/synchronize always target the right + # context even when the learner device is not the process default. + with torch.cuda.device(device): + self._event = torch.cuda.Event() + + def _write_slot(self, slot: int, device_flat: torch.Tensor) -> None: + # Explicit device: the current device may differ from the slot device. + stream = torch.cuda.current_stream(self.params[slot].device) + self.params[slot].copy_(device_flat) + # Mark the copy; completion is awaited in _complete_slot_write so the + # copy can overlap with the normalizer-stat CPU writes in between. + self._event.record(stream) + + def _complete_slot_write(self) -> None: + # The seq-even bump is a CPU store, but the slot write is an + # asynchronous GPU copy: complete (and thus device-globally visible) it + # here so a reader that observes the new version sees the data. + self._event.synchronize() + + +class WeightReceiver(ABC): + """Collector-side endpoint; owns the reader's load cursor. + + maybe_load implements the seqlock read loop; the transport-specific + snapshot/load steps are the abstract hooks. + """ + + def __init__(self, shared: WeightChannelShared): + self.shared = shared + self._loaded_version = 0 + # Reader-private staging for the normalizer stats (they always travel + # through host shared memory); pinned when the actor lives on CUDA so + # the stat copies into it can be non-blocking. Allocated lazily. + self._norm_staging: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + + @property + def version(self) -> int: + """Number of completed publishes (= seq // 2). Best-effort; readers + that need a consistent snapshot must go through :meth:`maybe_load`.""" + return int(self.shared.seq[0]) // 2 + + @property + def loaded_version(self) -> int: + """The snapshot version this receiver last loaded (0 = nothing yet).""" + return self._loaded_version + + @property + def lag(self) -> int: + """How many published versions behind the loaded policy is.""" + return max(0, self.version - self._loaded_version) + + def _ensure_norm_staging(self, actor: nn.Module) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self._norm_staging is None: + obs_dim = self.shared.mean[0].shape[1] + pinned = next(actor.parameters()).is_cuda + self._norm_staging = tuple( + torch.empty((1, obs_dim), dtype=torch.float32, pin_memory=pinned) for _ in range(3) + ) + return self._norm_staging + + @abstractmethod + def _snapshot(self, slot: int, flat_params: torch.Tensor | None, actor: nn.Module) -> None: + """Make a private, completed copy of the slot for this reader.""" + + @abstractmethod + def _load(self, slot: int, flat_params: torch.Tensor | None, actor: nn.Module) -> None: + """Update the actor from the snapshot taken by :meth:`_snapshot`.""" + + def maybe_load( + self, + actor: nn.Module, + obs_normalizer: nn.Module, + flat_params: torch.Tensor | None = None, + ) -> tuple[int, float, float, float]: + """Load the latest snapshot into ``actor``/``obs_normalizer`` if newer. + + Args: + actor: The collector's inference-only actor copy. Only touched + when a new version is loaded; the host transport without + ``flat_params`` copies per-parameter into it, otherwise it is + updated through the bound flat-parameter tensor. + obs_normalizer: The collector's read-only observation normalizer + (``EmpiricalNormalization`` or ``Identity``). Its ``_mean`` / + ``_std`` / ``_var`` / ``count`` stats are refreshed alongside + the weights when the learner publishes them. + flat_params: The collector actor's parameters bound into one + contiguous CUDA tensor (see :func:`bind_flat_params`), letting + the load be a single copy. Required by the CUDA-IPC transport; + ``None`` falls back to per-parameter copies on the host path. + Host-path staging buffers are an internal detail of the + receiver, allocated lazily on first use. + + Returns: + A ``(version, wait_writer_s, host_snapshot_s, actor_load_s)`` + tuple: the + version actually loaded (== the receiver's previous + ``loaded_version`` if nothing new was available this poll), + followed by wall-clock seconds spent waiting for an in-progress + writer (always 0 — a mid-publish poll returns immediately), copying + a stable snapshot, and loading it onto the actor's device. The + breakdown lets the collector distinguish shared-memory publication + contention from actor-load work. + + Implements the seqlock read loop — non-blocking on the writer: read + ``seq``; if it is odd (writer mid-publish) return the current version + immediately and retry on the next poll. Otherwise ensure the value is + stable, copy the data out, then re-check ``seq``: if the writer + published in between we discard the (possibly torn) copy and retry. + In practice retries are extremely rare — publish cadence is bounded by + ``weight_publish_interval`` gradient steps, while a read is one CPU + memcpy. + """ + wait_writer_s = 0.0 + host_snapshot_s = 0.0 + local_version = self._loaded_version + while True: + s1 = int(self.shared.seq[0]) + if s1 & 1: + # Writer is mid-publish. Weights are eventually consistent — + # keep running the current policy and pick up the new version + # on the next poll. Busy-waiting here would burn a core on the + # collector's critical path: the odd window spans the learner's + # in-flight gradient kernels (publish runs right after the + # async update), so it can last tens of milliseconds. This + # also covers retries: a torn read re-reads seq fresh instead + # of spinning on the stale value. + return local_version, wait_writer_s, host_snapshot_s, 0.0 + version = s1 // 2 + if version <= local_version: # nothing new to load + return local_version, wait_writer_s, host_snapshot_s, 0.0 + slot = version % 2 # active slot at this seq + has_stats = all(hasattr(obs_normalizer, k) for k in _NORM_KEYS) + # Phase 1 — snapshot: the transport makes a private, completed copy + # of the slot (host staging, or the bound flat params on the device + # path — each implementation owns its ordering constraints). + snapshot_start = time.perf_counter() + self._snapshot(slot, flat_params, actor) + if has_stats: + mean, std, var = self._ensure_norm_staging(actor) + mean.copy_(self.shared.mean[slot]) + std.copy_(self.shared.std[slot]) + var.copy_(self.shared.var[slot]) + count = int(self.shared.count[slot][0]) + # Re-check seq. If it changed, the writer published (at least + # once) during our copy and the slot may have been overwritten + # — discard and retry. This is the line that closes the race + # the old double-buffer had. + s2 = int(self.shared.seq[0]) + host_snapshot_s += time.perf_counter() - snapshot_start + if s1 != s2: + continue + + # Phase 2 — load: update the actor from the stable snapshot. + actor_load_start = time.perf_counter() + self._load(slot, flat_params, actor) + if has_stats: + obs_normalizer._mean.copy_(mean, non_blocking=True) + obs_normalizer._std.copy_(std, non_blocking=True) + obs_normalizer._var.copy_(var, non_blocking=True) + obs_normalizer.count.fill_(count) + # CUDA collectors keep the host staging buffers alive. The copies + # above are ordered before the next inference on the same stream; + # that inference already synchronizes after returning actions to + # the CPU environment, so a separate weight-load barrier only + # serializes collector and learner work unnecessarily. + actor_load_s = time.perf_counter() - actor_load_start + self._loaded_version = version + return version, wait_writer_s, host_snapshot_s, actor_load_s + + +class HostWeightReceiver(WeightReceiver): + """Host shared-memory slots; stages through CPU pinned memory. + + The actor copy is deliberately deferred to :meth:`_load` so a seqlock + retry never reuses staging memory that an in-flight H2D still reads. + """ + + def __init__(self, shared: WeightChannelShared, params: list[torch.Tensor]): + super().__init__(shared) + self.params = params + self._param_staging: torch.Tensor | None = None + + def _ensure_param_staging(self, actor: nn.Module) -> torch.Tensor: + if self._param_staging is None or self._param_staging.numel() != self.params[0].numel(): + pinned = next(actor.parameters()).is_cuda + self._param_staging = torch.empty(self.params[0].numel(), dtype=torch.float32, pin_memory=pinned) + return self._param_staging + + def _snapshot(self, slot: int, flat_params: torch.Tensor | None, actor: nn.Module) -> None: + del flat_params + self._ensure_param_staging(actor).copy_(self.params[slot]) + + def _load(self, slot: int, flat_params: torch.Tensor | None, actor: nn.Module) -> None: + del slot # the snapshot is already staged + staging = self._ensure_param_staging(actor) + if flat_params is not None: + flat_params.copy_(staging, non_blocking=True) + else: + load_flat_params(actor, staging) + + +class GpuIpcWeightReceiver(WeightReceiver): + """CUDA-IPC device slots; reads device-to-device into the bound flat params. + + The copy must COMPLETE before the seq re-check (a republish during an + in-flight copy would tear the read), so it is synchronized rather than + left async. + """ + + def __init__(self, shared: WeightChannelShared, params: list[torch.Tensor]): + super().__init__(shared) + self.params = params + + def _snapshot(self, slot: int, flat_params: torch.Tensor | None, actor: nn.Module) -> None: + del actor + if flat_params is None: + raise ValueError("the CUDA-IPC weight path requires bound flat_params") + flat_params.copy_(self.params[slot]) + # Sync the stream(s) the copy was enqueued on. Same-device (the gated + # design) it is one stream; a cross-device edge case may use both. + torch.cuda.current_stream(self.params[slot].device).synchronize() + if flat_params.device != self.params[slot].device: + torch.cuda.current_stream(flat_params.device).synchronize() + + def _load(self, slot: int, flat_params: torch.Tensor | None, actor: nn.Module) -> None: + del slot, flat_params, actor # no-op: _snapshot already wrote the actor + + +def weight_receiver_for(shared: WeightChannelShared, params: list[torch.Tensor]) -> WeightReceiver: + """Build the receiver matching a shipped slot pair (dispatch on device).""" + if params[0].is_cuda: + return GpuIpcWeightReceiver(shared, params) + return HostWeightReceiver(shared, params) diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py index c88c407d..72006e8f 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py @@ -28,9 +28,9 @@ from motrix_env_core.renderer import RenderConfig from motrix_rl.fastsac.agent import FastSacAgent from motrix_rl.fastsac.async_impl.collector import resolve_collector_inference_device -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot +from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.shm.weight_channel import WeightChannelShared from motrix_rl.fastsac.async_impl.worker import ( - actor_param_numel, build_env, run_collector_process, run_learner_process, @@ -129,12 +129,11 @@ def train(self) -> None: learner_device = self._device() collector_device = resolve_collector_inference_device(async_options.collector_inference_device) - param_numel = actor_param_numel(cfg, dims, action_scale, action_bias) # shared-memory primitives allocated in the parent, inherited by children. num_envs = self._context.num_envs ring = SharedTransitionRing(async_options.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim) - weights = WeightSnapshot(param_numel=param_numel, obs_dim=obs_dim) + weights = WeightChannelShared(obs_dim=obs_dim) control = Control() resume_step = 0 @@ -148,9 +147,18 @@ def train(self) -> None: control.collector_steps = resume_step control.global_step = resume_step + # Learner -> collector handoff of the CUDA-IPC weight slots (one message). + # ``learner=`` without an index resolves to the current CUDA device and + # is compared against the explicit collector index; only a conflicting + # explicit index (or a CPU collector) disables the IPC path. ctx = mp.get_context("spawn") stats_queue = ctx.Queue(maxsize=8) error_queue = ctx.Queue(maxsize=8) + # One-shot handshake queue for the weight slot pair: the learner-side + # endpoint allocates the slots (host shm, or CUDA-IPC device slots per + # the weight_ipc mode and size threshold — decided inside the learner) + # and ships the tensors to the collector-side endpoint. + slot_queue = ctx.Queue(maxsize=1) reported_errors: set[tuple[str, str]] = set() seed = self._context.seed @@ -203,6 +211,7 @@ def _drain_child_errors() -> list[tuple[str, str]]: self._context.checkpoint_format, self._resume_from, seed, + slot_queue, ), name="fastsac-async-learner", ) @@ -224,6 +233,7 @@ def _drain_child_errors() -> list[tuple[str, str]]: logging_interval, is_resume, seed, + slot_queue, ), name="fastsac-async-collector", ) diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py index a3faac01..4fcbd120 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -21,6 +21,7 @@ import sys import time import traceback +from multiprocessing.queues import Queue from pathlib import Path from queue import Empty from typing import Any @@ -35,9 +36,16 @@ from motrix_rl import checkpoints from motrix_rl.console import TrainingPanelStats, emit_training_panel, open_training_live from motrix_rl.fastsac.agent import FastSacAgent -from motrix_rl.fastsac.async_impl.collector import Collector +from motrix_rl.fastsac.async_impl.collector import Collector, resolve_collector_inference_device from motrix_rl.fastsac.async_impl.learner import Learner -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot +from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.shm.weight_channel import ( + GpuIpcWeightSender, + HostWeightSender, + WeightChannelShared, + WeightSender, + weight_receiver_for, +) from motrix_rl.fastsac.config import FastSacCfg from motrix_rl.fastsac.wrap import FastSacEnvWrap from motrix_rl.fastsac.wrap_np import FastSacNpEnvWrap @@ -223,22 +231,74 @@ def _configure_process_logging() -> None: logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") +def _build_weight_sender( + shared: WeightChannelShared, + cfg: FastSacCfg, + dims: tuple[int, int, int], + action_scale: torch.Tensor, + action_bias: torch.Tensor, + device: torch.device, +) -> WeightSender: + """Construct the sender-side weight endpoint per the configured transport. + + CUDA-IPC device slots only when learner and collector inference share one + GPU and the actor parameters reach the configured size threshold; host + shared-memory slots otherwise. The learner process is the only place both + transports' requirements can be met (the IPC-handle exporter needs the + CUDA context and must keep the tensors alive). + """ + opts = cfg.trainer.async_options + mode = opts.weight_ipc + # YAML 1.1 parses unquoted ``on``/``off`` scalars as booleans; accept that + # form so ``weight_ipc: on`` in a config behaves like the documented string. + if isinstance(mode, bool): + mode = "on" if mode else "off" + if mode not in ("auto", "on", "off"): + raise ValueError(f"async_options.weight_ipc must be auto, on or off, got {mode!r}") + collector_device = resolve_collector_inference_device(opts.collector_inference_device) + same_gpu = device.type == "cuda" and collector_device.type == "cuda" + if same_gpu: + # ``learner=`` without an index means the current device; resolve it so + # the comparison never treats "cuda" as matching an explicit different + # index. Only resolve under the cuda branch: a CPU learner has no CUDA + # context and torch.cuda.current_device() would raise. + learner_index = device.index if device.index is not None else torch.cuda.current_device() + same_gpu = learner_index == collector_device.index + if mode == "on" and not same_gpu: + reason = ( + "collector inference device is not CUDA" + if collector_device.type != "cuda" + else f"learner device {device} and collector device {collector_device} are different GPUs" + ) + logging.getLogger(__name__).warning( + "async_options.weight_ipc=on requires learner and collector inference on the same GPU, " + "but %s; falling back to host shared-memory transport", + reason, + ) + param_numel = actor_param_numel(cfg, dims, action_scale, action_bias) + use_gpu = same_gpu and (mode == "on" or (mode == "auto" and param_numel * 4 >= opts.weight_ipc_min_bytes)) + if use_gpu: + return GpuIpcWeightSender(shared, param_numel, device) + return HostWeightSender(shared, param_numel) + + def run_collector_process( env_spec: EnvBuildSpec, cfg: FastSacCfg, num_envs: int, - dims, + dims: tuple[int, int, int], action_scale: torch.Tensor, action_bias: torch.Tensor, ring: SharedTransitionRing, - weights: WeightSnapshot, + weights: WeightChannelShared, control: Control, - stats_queue, - error_queue, + stats_queue: Queue, + error_queue: Queue, num_iterations: int, logging_interval: int, is_resume: bool, - seed, + seed: int | None, + slot_queue: Queue, ) -> None: try: _configure_process_logging() @@ -249,6 +309,17 @@ def run_collector_process( obs_dim, critic_obs_dim, act_dim = dims device = torch.device("cpu") env = build_env(env_spec, num_envs, device, seed=seed) + # Handshake: build the receiver from the slot tensors the learner + # shipped (host shm or CUDA-IPC), before the collector is wired up. + try: + slots = slot_queue.get(timeout=60.0) + except Empty as exc: + raise RuntimeError( + "timed out waiting for the learner to ship the weight-slot tensors " + "(learner startup — agent build / checkpoint load / CUDA warmup — " + "likely failed or took over 60s; check the learner process's error queue/log)" + ) from exc + weight_rx = weight_receiver_for(weights, slots) collector = Collector( env, cfg, @@ -258,7 +329,7 @@ def run_collector_process( action_scale, action_bias, ring, - weights, + weight_rx, control, is_resume=is_resume, ) @@ -289,14 +360,14 @@ def run_collector_process( def run_learner_process( cfg: FastSacCfg, num_envs: int, - dims, + dims: tuple[int, int, int], action_scale: torch.Tensor, action_bias: torch.Tensor, ring: SharedTransitionRing, - weights: WeightSnapshot, + weights: WeightChannelShared, control: Control, - stats_queue, - error_queue, + stats_queue: Queue, + error_queue: Queue, num_iterations: int, logging_interval: int, save_interval: int, @@ -305,7 +376,8 @@ def run_learner_process( checkpoint_dir: str, checkpoint_format: str, resume_from: str | None, - seed, + seed: int | None, + slot_queue: Queue, ) -> None: _configure_process_logging() console, live = open_training_live() @@ -328,7 +400,12 @@ def run_learner_process( ckpt = torch.load(resume_from, map_location=device, weights_only=False) agent.load_state_dict(ckpt, load_optimizers=True) - learner = Learner(agent, cfg, ring, weights, control) + # Build the sender endpoint and ship its slot tensors BEFORE the first + # publish so the collector (blocking on the handshake queue) builds the + # matching receiver and takes the agreed transport from step one. + weight_tx = _build_weight_sender(weights, cfg, dims, action_scale, action_bias, device) + slot_queue.put(weight_tx.params) + learner = Learner(agent, cfg, ring, weight_tx, control) learner.publish_weights() # give the collector an initial policy before it warms up start_time = time.time() @@ -493,7 +570,7 @@ def _drain_stats(): ) writer.add_scalar("async/policy_lag", last_stats["policy_lag"], step) writer.add_scalar("async/ring_fill", ring.size(), step) - writer.add_scalar("async/weight_version", weights.version, step) + writer.add_scalar("async/weight_version", weight_tx.version, step) writer.add_scalar("async/utd", utd, step) writer.add_scalar("perf/collect_ms_per_batch", collector_timing_ms.get("collect", 0.0), step) for k, v in collector_timing_detail_ms.items(): diff --git a/motrix_rl/src/motrix_rl/fastsac/config.py b/motrix_rl/src/motrix_rl/fastsac/config.py index d5096403..2f25e167 100644 --- a/motrix_rl/src/motrix_rl/fastsac/config.py +++ b/motrix_rl/src/motrix_rl/fastsac/config.py @@ -77,6 +77,17 @@ class FastSacAsyncOptionsCfg: # collector. learner_cpu_cores: str | None = None collector_cpu_cores: str | None = None + # Weight-snapshot transport between the async learner and collector: + # "auto" picks CUDA-IPC device slots when both sides share one GPU and the + # actor parameters are large enough (>= weight_ipc_min_bytes) for the + # device path to pay off, and host shared memory otherwise; "on"/"off" + # force the device/host path ("on" logs a warning and falls back to the + # host path when learner and collector are not on the same GPU). Small + # actors are faster on the host path — + # the IPC path's stream synchronizations cost more than the sub-millisecond + # host transfer they avoid. + weight_ipc: str = "auto" + weight_ipc_min_bytes: int = 16 * 1024 * 1024 @dataclass diff --git a/motrix_rl/src/motrix_rl/fastsac/sync/train.py b/motrix_rl/src/motrix_rl/fastsac/sync/train.py index 5189c958..1a69404b 100644 --- a/motrix_rl/src/motrix_rl/fastsac/sync/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/sync/train.py @@ -258,12 +258,14 @@ def emit_msg(msg: str) -> None: ep_return += rewards ep_len += 1 done_idx = torch.nonzero(terminated | truncated, as_tuple=False).flatten() - for j in done_idx.tolist(): - recent_returns.append(float(ep_return[j])) - recent_lengths.append(float(ep_len[j])) - ep_return[j] = 0.0 - ep_len[j] = 0.0 - n_episodes += 1 + if done_idx.numel(): + # Vectorized: one batched gather + clear instead of + # per-episode Python-level scalar indexing. + recent_returns.extend(ep_return[done_idx].tolist()) + recent_lengths.extend(ep_len[done_idx].tolist()) + ep_return[done_idx] = 0.0 + ep_len[done_idx] = 0.0 + n_episodes += int(done_idx.numel()) recent_returns = recent_returns[-100:] recent_lengths = recent_lengths[-100:] diff --git a/motrix_rl/tests/test_fastsac_collector.py b/motrix_rl/tests/test_fastsac_collector.py index 26ab124e..1c5a7b89 100644 --- a/motrix_rl/tests/test_fastsac_collector.py +++ b/motrix_rl/tests/test_fastsac_collector.py @@ -1,14 +1,16 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import sys +import time from types import SimpleNamespace import pytest import torch -import motrix_rl.fastsac.async_impl.shm as shm_module from motrix_rl.fastsac.async_impl.collector import Collector, resolve_collector_inference_device -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot +from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.shm.weight_channel import HostWeightReceiver, HostWeightSender, WeightChannelShared from motrix_rl.fastsac.buffer import EmpiricalNormalization from motrix_rl.fastsac.networks import Actor @@ -92,8 +94,10 @@ def _collector(device: str, *, compile: bool = False, amp: bool = False): cfg = _cfg(device, compile=compile, amp=amp) ring = SharedTransitionRing(2, _NUM_ENVS, _OBS_DIM, _CRITIC_OBS_DIM, _ACT_DIM) source_actor, source_normalizer = _source_policy() - weights = WeightSnapshot(sum(p.numel() for p in source_actor.parameters()), _OBS_DIM) - weights.publish(source_actor, source_normalizer) + shared = WeightChannelShared(_OBS_DIM) + weight_tx = HostWeightSender(shared, sum(p.numel() for p in source_actor.parameters())) + weight_rx = HostWeightReceiver(shared, weight_tx.params) + weight_tx.publish(source_actor, source_normalizer) collector = Collector( env, cfg, @@ -103,30 +107,62 @@ def _collector(device: str, *, compile: bool = False, amp: bool = False): source_actor.action_scale, source_actor.action_bias, ring, - weights, + weight_rx, Control(), ) collector.reset() collector.sync_weights() - return collector, source_actor, source_normalizer, weights + return collector, source_actor, source_normalizer, weight_tx -def test_weight_snapshot_prepares_params_before_opening_seqlock_write(monkeypatch) -> None: +def test_weight_publish_prepares_params_before_opening_seqlock_write(monkeypatch) -> None: actor, normalizer = _source_policy() - weights = WeightSnapshot(sum(param.numel() for param in actor.parameters()), _OBS_DIM) - original_flatten = shm_module.flatten_params + shared = WeightChannelShared(_OBS_DIM) + weight_tx = HostWeightSender(shared, sum(param.numel() for param in actor.parameters())) + weight_rx = HostWeightReceiver(shared, weight_tx.params) + weight_channel_module = sys.modules[WeightChannelShared.__module__] + original_flatten = weight_channel_module.flatten_params observed_sequences = [] def observe_flatten(module): - observed_sequences.append(int(weights._seq[0])) + observed_sequences.append(int(weight_tx.shared.seq[0])) return original_flatten(module) - monkeypatch.setattr(shm_module, "flatten_params", observe_flatten) - weights.publish(actor, normalizer) - weights.publish(actor, normalizer) + monkeypatch.setattr(weight_channel_module, "flatten_params", observe_flatten) + weight_tx.publish(actor, normalizer) + weight_tx.publish(actor, normalizer) assert observed_sequences == [0, 2] - assert weights.version == 2 + assert weight_tx.version == 2 + assert weight_rx.version == 2 + assert weight_rx.lag == 2 + + +def test_seqlock_read_is_nonblocking_during_publish() -> None: + """A mid-publish (odd seq) poll returns the current version immediately. + + Contract: the collector never busy-waits for the learner's publish window + (it spans the learner's in-flight gradient kernels); weights are + eventually consistent and the next poll picks up the new version. + """ + actor, normalizer = _source_policy() + shared = WeightChannelShared(_OBS_DIM) + weight_tx = HostWeightSender(shared, sum(param.numel() for param in actor.parameters())) + weight_rx = HostWeightReceiver(shared, weight_tx.params) + weight_tx.publish(actor, normalizer) + + shared.seq[0] = 3 # simulate a publish in progress (odd) + start = time.perf_counter() + version, wait_writer_s, _, _ = weight_rx.maybe_load(actor, normalizer) + elapsed = time.perf_counter() - start + + assert version == 0 # keeps the receiver's current version (nothing loaded yet) + assert wait_writer_s == 0.0 + assert elapsed < 1.0 # returned immediately, did not spin until seq closes + + shared.seq[0] = 4 # publish completed + version, _, _, _ = weight_rx.maybe_load(actor, normalizer) + assert version == 2 # picked up on the next poll def test_collector_explicit_cpu_placement_and_timing() -> None: @@ -135,8 +171,8 @@ def test_collector_explicit_cpu_placement_and_timing() -> None: assert collector.device.type == "cpu" assert collector.obs.device.type == "cpu" assert all(param.device.type == "cpu" for param in collector.actor.parameters()) - assert not collector._weight_param_staging.is_pinned() - assert all(not buffer.is_pinned() for buffer in collector._weight_normalizer_staging) + assert not collector.weights._param_staging.is_pinned() + assert all(not buffer.is_pinned() for buffer in collector.weights._norm_staging) torch.testing.assert_close( torch.cat([p.detach().flatten() for p in collector.actor.parameters()]), torch.cat([p.detach().flatten() for p in source_actor.parameters()]), @@ -176,7 +212,7 @@ def test_collector_rejects_unavailable_cuda_index() -> None: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA collector test requires a GPU") def test_cuda_collector_uses_flat_weight_copy_and_reuses_staging() -> None: - collector, source_actor, source_normalizer, weights = _collector("cuda") + collector, source_actor, source_normalizer, weight_tx = _collector("cuda") assert collector.device.type == "cuda" assert collector.obs.device.type == "cpu" @@ -196,7 +232,7 @@ def test_cuda_collector_uses_flat_weight_copy_and_reuses_staging() -> None: with torch.no_grad(): source_actor.fc_mu.bias.add_(0.25) source_normalizer._mean.add_(0.5) - weights.publish(source_actor, source_normalizer) + weight_tx.publish(source_actor, source_normalizer) assert collector.policy_lag == 1 collector.sync_weights() assert collector.policy_lag == 0 @@ -209,8 +245,8 @@ def test_cuda_collector_uses_flat_weight_copy_and_reuses_staging() -> None: collector._obs_host.data_ptr(), collector._obs_device.data_ptr(), collector._actions_host.data_ptr(), - collector._weight_param_staging.data_ptr(), - *(buffer.data_ptr() for buffer in collector._weight_normalizer_staging), + collector.weights._param_staging.data_ptr(), + *(buffer.data_ptr() for buffer in collector.weights._norm_staging), collector._flat_params.data_ptr(), ) first = collector._infer(torch.randn(_NUM_ENVS, _OBS_DIM)) @@ -220,8 +256,8 @@ def test_cuda_collector_uses_flat_weight_copy_and_reuses_staging() -> None: collector._obs_host.data_ptr(), collector._obs_device.data_ptr(), collector._actions_host.data_ptr(), - collector._weight_param_staging.data_ptr(), - *(buffer.data_ptr() for buffer in collector._weight_normalizer_staging), + collector.weights._param_staging.data_ptr(), + *(buffer.data_ptr() for buffer in collector.weights._norm_staging), collector._flat_params.data_ptr(), ) diff --git a/motrix_rl/tests/test_rl_sim_backend.py b/motrix_rl/tests/test_rl_sim_backend.py index 0e77fedc..c908f00c 100644 --- a/motrix_rl/tests/test_rl_sim_backend.py +++ b/motrix_rl/tests/test_rl_sim_backend.py @@ -19,7 +19,8 @@ from motrix_env_core.direct.env import DirectEnv from motrix_env_core.registry import EnvBuildSpec from motrix_env_motrixsim.torch_env import TorchEnv, TorchEnvState, TorchObs -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot +from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.shm.weight_channel import HostWeightSender, WeightChannelShared from motrix_rl.fastsac.async_impl.worker import actor_param_numel, run_collector_process from motrix_rl.fastsac.wrap import FastSacEnvWrap @@ -254,17 +255,20 @@ def _collect_in_spawn(sim_backend: str) -> tuple[torch.Tensor, ...]: action_scale = torch.ones(_ACT_DIM) action_bias = torch.zeros(_ACT_DIM) ring = SharedTransitionRing(2, _NUM_ENVS, *dims) - weights = WeightSnapshot(actor_param_numel(cfg, dims, action_scale, action_bias), _OBS_DIM) + weights = WeightChannelShared(_OBS_DIM) control = Control() ctx = mp.get_context("spawn") stats_queue = ctx.Queue(maxsize=2) error_queue = ctx.Queue(maxsize=2) + slot_queue = ctx.Queue(maxsize=1) + weight_tx = HostWeightSender(weights, actor_param_numel(cfg, dims, action_scale, action_bias)) + slot_queue.put(weight_tx.params) # ship before the collector process starts env_cls = _AsyncNpEnv if sim_backend == "np" else _AsyncTorchEnv env_spec = EnvBuildSpec(env_cls, EnvCfg(scene=SceneCfg())) ipc_resources = (ring, weights, control, stats_queue, error_queue) process = ctx.Process( target=run_collector_process, - args=(env_spec, cfg, _NUM_ENVS, dims, action_scale, action_bias, *ipc_resources, 1, 1, False, 7), + args=(env_spec, cfg, _NUM_ENVS, dims, action_scale, action_bias, *ipc_resources, 1, 1, False, 7, slot_queue), ) process.start()