perf(fastsac): weight-snapshot transport — fused publish + CUDA-IPC slots with size gate

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

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:]
 

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