perf(fastsac): fast async learner pipeline — batched drain, CUDA-IPC ring, fused updates

configs/algo_base/motrix.fastsac.yaml

@@ -101,11 +101,19 @@ trainer:
     # machine's CPU/GPU balance (see issue #62's sweep table).
     learner_cpu_cores: null
     collector_cpu_cores: null
+    # Transition-ring transport: "auto" (default) puts the ring slots in
+    # CUDA-IPC device memory when learner and collector inference share one
+    # GPU (fused single H2D on the collector, D2D-only learner ingest) and
+    # falls back to host shared memory otherwise; "on"/"off" force the
+    # device/host path ("on" warns and falls back to the host ring when the
+    # two sides are not on one GPU). Keep the values quoted: unquoted on/off
+    # parse as booleans in YAML.
+    transition_ipc: "auto"
     # 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
+    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/agent.py

@@ -171,7 +171,12 @@ def __init__(
             inductor_config.compile_threads = 1
             # B3 experiment: reduce-overhead wraps each compiled update in a
             # CUDA graph (trees), removing per-kernel launch and region-gap CPU
-            # time inside the hot learner loop.
+            # time inside the hot learner loop. (Coarser units were measured at
+            # microduck scale and do NOT help: torch.compile splits regions
+            # containing backward()/optimizer.step() regardless of the outer
+            # boundary, and module-only/default-mode compilation is ~50% slower
+            # or crashes on cudagraph output pools. The per-step boundary is
+            # the measured optimum.)
             self._update_main_runtime = torch.compile(self._update_main, mode="reduce-overhead")
             self._update_pol_runtime = torch.compile(self._update_pol, mode="reduce-overhead")
 
@@ -226,6 +231,19 @@ def _update_main(self, b: dict):
             torch.nn.utils.clip_grad_norm_(self.qnet.parameters(), cfg.max_grad_norm)
         self.q_optimizer.step()
 
+        # Target-network soft update, fused into the compiled region: the
+        # foreach ops write the same kind of param storage the optimizer step
+        # above already mutates in place, so the whole thing replays as one
+        # CUDA graph instead of two eager kernel launches per step. Ordering
+        # constraints: after the target read at the top of this function and
+        # after q_optimizer.step() (both satisfied here); _update_pol reads
+        # neither qnet_target nor writes qnet, so it stays order-independent.
+        with torch.no_grad():
+            tau = cfg.tau
+            tgt = [p.data for p in self.qnet_target.parameters()]
+            torch._foreach_mul_(tgt, 1.0 - tau)
+            torch._foreach_add_(tgt, [p.data for p in self.qnet.parameters()], alpha=tau)
+
         alpha_loss = torch.zeros((), device=self.device)
         if cfg.use_autotune:
             alpha_loss = (-self.log_alpha.exp() * (next_logp.detach() + self.target_entropy)).mean()
@@ -255,14 +273,6 @@ def _update_pol(self, b: dict):
         self.actor_optimizer.step()
         return actor_loss.detach().float(), (-log_probs.mean()).detach().float()
 
-    @torch.no_grad()
-    def _soft_update(self):
-        tau = self.cfg.tau
-        src = [p.data for p in self.qnet.parameters()]
-        tgt = [p.data for p in self.qnet_target.parameters()]
-        torch._foreach_mul_(tgt, 1.0 - tau)
-        torch._foreach_add_(tgt, src, alpha=tau)
-
     def update(self, num_updates: int):
         """Run ``num_updates`` gradient steps, each on a fresh batch.
 
@@ -288,7 +298,7 @@ def update(self, num_updates: int):
             return None
         batch_per_env = max(cfg.batch_size // self.num_envs, 1)
         last = (torch.zeros((), device=self.device),) * 5
-        timing_s = {key: 0.0 for key in ("sample_normalize", "critic_alpha", "actor", "soft_update")}
+        timing_s = {key: 0.0 for key in ("sample_normalize", "critic_alpha", "actor")}
         update_started = time.perf_counter()
         # Batched data preparation (Holosoma-style): sample once and normalize
         # once per update() call, then slice views into per-gradient-step
@@ -319,22 +329,28 @@ def update(self, num_updates: int):
             # Required with reduce-overhead (CUDA graph trees): open a new graph
             # generation for this update. NOTE it does not preserve anything --
             # it *invalidates* the previous generation's outputs, which is why
-            # every output that outlives its own iteration goes through `_own`.
+            # any output that outlives its own iteration goes through `_own`.
             torch.compiler.cudagraph_mark_step_begin()
             stage_started = time.perf_counter()
-            qf_loss, alpha_loss, qf_max, qf_min = _own(self._update_main_runtime(b))
+            outputs = self._update_main_runtime(b)
             timing_s["critic_alpha"] += time.perf_counter() - stage_started
 
-            actor_loss, entropy = last[3], last[4]
+            actor_pair = (last[3], last[4])
             if (self.update_idx + i) % cfg.policy_frequency == 0:
                 stage_started = time.perf_counter()
-                actor_loss, entropy = _own(self._update_pol_runtime(b))
+                pol_outputs = self._update_pol_runtime(b)
                 timing_s["actor"] += time.perf_counter() - stage_started
+                # Always own: the pair is carried across later generations
+                # within this call AND the returned metrics must stay readable
+                # after future update() calls replay the pol graph.
+                actor_pair = _own(pol_outputs)
+
+            # Only the final step's main outputs feed the returned metrics; own
+            # them so they survive future update() calls' replays. Earlier
+            # steps' outputs are never read — only replaced below.
+            main_outputs = _own(outputs) if i == num_updates - 1 else outputs
+            last = (*main_outputs, *actor_pair)
 
-            stage_started = time.perf_counter()
-            self._soft_update()
-            timing_s["soft_update"] += time.perf_counter() - stage_started
-            last = (qf_loss, alpha_loss, qf_max, actor_loss, entropy)
         self.update_idx += num_updates
         timing_s["total"] = time.perf_counter() - update_started
         self._last_update_timing_ms = {key: value * 1000.0 for key, value in timing_s.items()}

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 its :class:`~motrix_rl.fastsac.async_impl.shm.WeightReceiver` endpoint. Each step
+the learner via its :class:`~motrix_rl.fastsac.async_impl.transport.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,8 +20,8 @@
 import torch
 from torch import nn
 
-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.async_impl.transport import Control, IpcTransitionRing, SharedTransitionRing, bind_flat_params
+from motrix_rl.fastsac.async_impl.transport.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 +76,7 @@ def __init__(
         act_dim: int,
         action_scale: torch.Tensor,
         action_bias: torch.Tensor,
-        ring: SharedTransitionRing,
+        ring: SharedTransitionRing | IpcTransitionRing,
         weights: WeightReceiver,
         control: Control,
         is_resume: bool = False,
@@ -350,6 +350,11 @@ def snapshot_stats(self) -> dict:
         env_perf = self._env_perf
         if env_perf is not None:
             for path, mean_ms in env_perf.stage_mean_ms("step").items():
+                if "." in path:
+                    # Only the first sub-stage level is reported: a parent's
+                    # total already includes its children, so deeper paths
+                    # would duplicate time without adding actionable signal.
+                    continue
                 stats["timing_ms"][f"env_step.{path}"] = mean_ms
             env_perf.reset()
         self.term_accum, self.term_count = {}, 0

motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py

@@ -4,10 +4,11 @@
 """Learner: owns a full ``FastSacAgent`` and drives training off the shared ring.
 
 Unlike the sync trainer it does NOT step the env. It drains raw transitions from
-: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 its :class:`~motrix_rl.fastsac.async_impl.shm.WeightSender` endpoint.
+the shared transition ring (host fields or CUDA-IPC device fields, see
+``transport/ring.py`` / ``transport/ipc_ring.py``) 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 its
+:class:`~motrix_rl.fastsac.async_impl.transport.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
@@ -18,9 +19,11 @@
 
 import time
 
+import torch
+
 from motrix_rl.fastsac.agent import FastSacAgent
-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.async_impl.transport import Control, IpcTransitionRing, SharedTransitionRing
+from motrix_rl.fastsac.async_impl.transport.weight_channel import WeightSender
 from motrix_rl.fastsac.config import FastSacCfg
 
 
@@ -29,7 +32,7 @@ def __init__(
         self,
         agent: FastSacAgent,
         cfg: FastSacCfg,
-        ring: SharedTransitionRing,
+        ring: SharedTransitionRing | IpcTransitionRing,
         weights: WeightSender,
         control: Control,
     ):
@@ -42,6 +45,31 @@ def __init__(
         self._learning_starts = agent.cfg.learning_starts
         self._last_publish_ms = 0.0
 
+        # Host-ring async-ingest plumbing (CUDA only): contiguous runs are
+        # staged through pinned buffers and moved to the GPU with non-blocking
+        # H2D copies on a dedicated stream, so ingestion overlaps gradient
+        # updates. The CUDA-IPC device ring needs none of this — its slots are
+        # already on the device and ingest is a same-stream D2D copy.
+        self._device_ring = isinstance(ring, IpcTransitionRing)
+        self._copy_stream = None
+        self._copy_event = None
+        self._pending_copy = False
+        self._staging = None
+        if agent.device.type == "cuda" and not self._device_ring:
+            self._copy_stream = torch.cuda.Stream(device=agent.device)
+            self._copy_event = torch.cuda.Event()
+            rb = agent.rb
+            chunk = max(self.async_options.max_ingest_per_iter, 1)
+            pin = lambda *shape: torch.empty(*shape, pin_memory=True)  # noqa: E731
+            self._staging = (
+                pin(chunk, rb.n_env, rb.n_obs),
+                pin(chunk, rb.n_env, rb.n_critic_obs),
+                pin(chunk, rb.n_env, rb.n_act),
+                pin(chunk, rb.n_env),
+                torch.empty(chunk, rb.n_env, dtype=torch.int64, pin_memory=True),
+                torch.empty(chunk, rb.n_env, dtype=torch.int64, pin_memory=True),
+            )
+
         # keep normalizers/actor in train mode: the learner is the update side.
         self.agent.set_train_mode()
 
@@ -56,33 +84,62 @@ def update_idx(self) -> int:
     def drain(self) -> int:
         """Move up to ``max_ingest_per_iter`` ring slots into the replay buffer.
 
-        Returns the number of slots ingested. Read cursor advances only after the
-        GPU copy, so the collector cannot clobber an in-flight slot. The replay
-        buffer derives each transition's ``next_obs`` from the following slot's
-        stored observation, so no successor peek is needed and a slot is
-        ingested as soon as it is committed.
+        Returns the number of slots ingested. Slots are consumed in contiguous
+        runs (``ring.read_span()``):
+
+        * CUDA-IPC device ring: the strided device views go straight into the
+          replay buffer with D2D copies on the current stream (ordered before
+          any subsequent sample by stream order); the read cursor is released
+          behind an event once those copies complete.
+        * host ring on CUDA: each run is memcpy'd into pinned staging and moved
+          to the GPU as one non-blocking H2D copy per field on the copy stream
+          (overlapping the next gradient update), then the read cursor
+          advances — the producer cannot clobber in-flight data because the
+          ring slot was already fully copied to staging.
+        * host ring on CPU: the views copy directly into the buffer.
+
+        The replay buffer derives each transition's ``next_obs`` from the
+        following slot's stored observation, so no successor peek is needed and
+        a slot is ingested as soon as it is committed.
         """
-        device = self.agent.device
+        budget = max(self.async_options.max_ingest_per_iter, 1)
+        if self._staging is not None:
+            # Staging may still be the source of the previous drain's in-flight
+            # H2D copies; wait before overwriting it. Any pending copy was
+            # already joined by the intervening maybe_train(), so this is
+            # normally a no-op sync.
+            self._copy_stream.synchronize()
         ingested = 0
-        for _ in range(max(self.async_options.max_ingest_per_iter, 1)):
-            if not self.ring.has_next():
-                break
-            slot = self.ring.read_slot()
-            assert slot is not None  # has_next implies a readable slot
-            obs, critic_obs, actions, rewards, dones, truncations = slot
-            self.agent.rb.extend(
-                obs.to(device),
-                critic_obs.to(device),
-                actions.to(device),
-                rewards.to(device),
-                dones.to(device),
-                truncations.to(device),
-            )
-            self.ring.commit_read()
-            ingested += 1
+        while ingested < budget and self.ring.has_next():
+            k, views = self.ring.read_span()
+            k = min(k, budget - ingested)
+            if k < views[0].shape[0]:
+                views = tuple(v[:k] for v in views)
+            if self._staging is not None:
+                for stage, view in zip(self._staging, views):
+                    stage[:k].copy_(view)  # ring -> pinned (plain CPU memcpy)
+                with torch.cuda.stream(self._copy_stream):
+                    self.agent.rb.extend_batch(*(stage[:k] for stage in self._staging))
+                self._pending_copy = True
+            else:
+                # CUDA-IPC device ring or CPU host ring: consume the views
+                # directly (D2D strided copy, or plain CPU copy).
+                self.agent.rb.extend_batch(*views)
+            self.ring.commit_reads(k)
+            ingested += k
+        if self._pending_copy:
+            self._copy_event.record(self._copy_stream)
         return ingested
 
     # ------------------------------------------------------------------ update
+    def wait_ingest(self) -> None:
+        """Block until every issued ingest copy has landed in the buffer."""
+        if self._copy_stream is not None:
+            self._copy_stream.synchronize()
+        elif self._device_ring:
+            torch.cuda.synchronize(self.agent.device)
+        self._pending_copy = False
+
     def _ready(self) -> bool:
         return self.control.collector_steps >= self._learning_starts and self.agent.rb.num_stored > 0
 
@@ -101,6 +158,13 @@ def maybe_train(self, ingested: int) -> dict | None:
         """Run ratio-governed updates. Returns last metrics dict or ``None``."""
         if not self._ready():
             return None
+        if self._pending_copy:
+            # Sampling reads slots the copy stream may still be filling; make
+            # the compute stream wait for the in-flight H2D ingest copies.
+            # (The device-ring path needs no event: its D2D copies are on the
+            # same stream as sampling.)
+            self._copy_event.wait()
+            self._pending_copy = False
         n = self._num_updates_for(ingested)
         # Delegate the per-step work to the agent; this module no longer keeps
         # its own update-loop / update_idx / _last_actor — the agent's