perf(fastsac): async collector push 的拷贝冗余与同步阻塞

configs/algo_base/motrix.fastsac.yaml

@@ -95,3 +95,9 @@ trainer:
     collector_amp: true
     # Collector autocast dtype: fp16 or bf16.
     collector_amp_dtype: fp16
+    # CPU affinity pinning spec per side, e.g. "0:5,7" pins cores 0-5 and 7.
+    # Cores outside the process's available set are dropped; null disables
+    # that side's pinning. Defaults off: the optimal split depends on the
+    # machine's CPU/GPU balance (see issue #62's sweep table).
+    learner_cpu_cores: null
+    collector_cpu_cores: null

motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py

@@ -157,8 +157,8 @@ def __init__(
                 inductor_config.compile_threads = 1
                 self._policy_runtime = torch.compile(self._policy, mode="reduce-overhead")
 
-        self._action_scale_cpu = action_scale.detach().cpu()
-        self._action_bias_cpu = action_bias.detach().cpu()
+        self._action_scale_cpu = self.actor.action_scale.detach().cpu()
+        self._action_bias_cpu = self.actor.action_bias.detach().cpu()
 
         # rollout state
         self.obs = None
@@ -244,6 +244,7 @@ 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)
 
+    # ------------------------------------------------------------------ ring handoff
     # ------------------------------------------------------------------ step
     def step_once(self) -> bool:
         """Run one env-step batch and push it to the ring.
@@ -270,15 +271,19 @@ def step_once(self) -> bool:
         next_obs, next_critic_obs, rewards, terminated, truncated = self.env.step(actions)
         t_push = time.perf_counter()
 
+        # The transition is (obs, action, reward, done) with obs being the
+        # pre-step observation; next_obs of step t is the stored obs of t+1.
+        # The learner derives next_obs from the successor slot, so the newest
+        # batch only becomes ingestible after the next push; the final batch
+        # pushed before training stops is intentionally dropped (one batch of
+        # num_envs transitions out of a full training run).
         pushed = self.ring.push(
-            self.obs.detach().cpu(),
-            self.critic_obs.detach().cpu(),
-            actions.detach().cpu(),
-            rewards.detach().cpu(),
-            terminated.detach().long().cpu(),
-            truncated.detach().long().cpu(),
-            next_obs.detach().cpu(),
-            next_critic_obs.detach().cpu(),
+            self.obs.detach(),
+            self.critic_obs.detach(),
+            actions.detach(),
+            rewards.detach(),
+            terminated.detach().long(),
+            truncated.detach().long(),
         )
         assert pushed, "ring became full after is_full() check — single-producer invariant violated"
         t_bookkeep = time.perf_counter()

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

@@ -56,15 +56,20 @@ 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.
+        GPU copy, so the collector cannot clobber an in-flight slot. Each slot's
+        ``next_obs``/``next_critic_obs`` come from the successor slot
+        (``ring.peek_next``), so a slot is only ingested once its successor is
+        committed; the last slot of a drain wave waits for the next one.
         """
         device = self.agent.device
         ingested = 0
         for _ in range(max(self.async_options.max_ingest_per_iter, 1)):
-            slot = self.ring.read_slot()
-            if slot is None:
+            if not self.ring.has_next():
                 break
-            obs, critic_obs, actions, rewards, dones, truncations, next_obs, next_critic_obs = slot
+            slot = self.ring.read_slot()
+            assert slot is not None  # has_next implies a readable slot
+            obs, critic_obs, actions, rewards, dones, truncations = slot
+            next_obs, next_critic_obs = self.ring.peek_next()
             self.agent.rb.extend(
                 obs.to(device),
                 critic_obs.to(device),

motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py

@@ -61,19 +61,24 @@ def _shared(shape, dtype) -> torch.Tensor:
 class SharedTransitionRing:
     """SPSC ring of transition batches with bounded backpressure.
 
-    Each slot holds one env-step batch: the eight tensors that
-    :meth:`motrix_rl.fastsac.buffer.SimpleReplayBuffer.extend` consumes, with the
-    leading dimension being ``num_envs`` (so a slot is ``(num_envs, dim)``).
+    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, 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).
+    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
     ~~~~~~~~~~~~~~~
@@ -83,7 +88,9 @@ class SharedTransitionRing:
     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).
+    (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 = (
@@ -93,8 +100,6 @@ class SharedTransitionRing:
         "rewards",
         "dones",
         "truncations",
-        "next_obs",
-        "next_critic_obs",
     )
 
     def __init__(
@@ -105,6 +110,10 @@ def __init__(
         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
@@ -114,8 +123,6 @@ def __init__(
         self.rewards = _shared((capacity, num_envs), f32)
         self.dones = _shared((capacity, num_envs), i64)
         self.truncations = _shared((capacity, num_envs), i64)
-        self.next_obs = _shared((capacity, num_envs, obs_dim), f32)
-        self.next_critic_obs = _shared((capacity, num_envs, critic_obs_dim), f32)
         # cursors are shared so the two processes see each other's progress.
         self._write = _shared((1,), i64)
         self._read = _shared((1,), i64)
@@ -140,7 +147,7 @@ def size(self) -> int:
     def is_full(self) -> bool:
         return self.size() >= self.capacity
 
-    def push(self, obs, critic_obs, actions, rewards, dones, truncations, next_obs, next_critic_obs) -> bool:
+    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``
@@ -155,8 +162,6 @@ def push(self, obs, critic_obs, actions, rewards, dones, truncations, next_obs,
         self.rewards[slot].copy_(rewards)
         self.dones[slot].copy_(dones)
         self.truncations[slot].copy_(truncations)
-        self.next_obs[slot].copy_(next_obs)
-        self.next_critic_obs[slot].copy_(next_critic_obs)
         # 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).
@@ -179,10 +184,27 @@ def read_slot(self):
             self.rewards[slot],
             self.dones[slot],
             self.truncations[slot],
-            self.next_obs[slot],
-            self.next_critic_obs[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

motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py

@@ -16,7 +16,9 @@
 from __future__ import annotations
 
 import logging
+import os
 import random
+import sys
 import time
 import traceback
 from pathlib import Path
@@ -116,6 +118,97 @@ def build_agent(cfg: FastSacCfg, dims, num_envs, device, action_scale, action_bi
 
 
 # ------------------------------------------------------------------ collector process
+def _available_cpu_ids() -> set[int]:
+    """CPU ids this process may run on (Linux affinity mask, Windows process mask)."""
+    if hasattr(os, "sched_getaffinity"):
+        return set(os.sched_getaffinity(0))
+    if sys.platform == "win32":
+        import ctypes
+
+        process_mask = ctypes.c_ulonglong()
+        system_mask = ctypes.c_ulonglong()
+        kernel32 = ctypes.windll.kernel32
+        got = kernel32.GetProcessAffinityMask(
+            kernel32.GetCurrentProcess(), ctypes.byref(process_mask), ctypes.byref(system_mask)
+        )
+        if got:
+            return {i for i in range(64) if process_mask.value & (1 << i)}
+    return set(range(os.cpu_count() or 1))
+
+
+def _set_cpu_affinity(cpus: set[int]) -> None:
+    """Pin this process to ``cpus``; best-effort and never fatal.
+
+    Raises OSError (or skips with a warning on platforms without an affinity
+    API) so the caller can decide; permission-constrained environments
+    (cgroups/cpusets) degrade to a warning instead of killing the worker.
+    """
+    if hasattr(os, "sched_setaffinity"):
+        os.sched_setaffinity(0, cpus)
+    elif sys.platform == "win32":
+        import ctypes
+
+        mask = ctypes.c_ulonglong(sum(1 << c for c in cpus))
+        current = ctypes.windll.kernel32.GetCurrentProcess()
+        if not ctypes.windll.kernel32.SetProcessAffinityMask(current, mask):
+            raise OSError(f"SetProcessAffinityMask failed for cpus {sorted(cpus)}")
+    else:
+        logging.getLogger(__name__).warning(
+            "CPU affinity is unsupported on this platform; skipping pinning to %s", sorted(cpus)
+        )
+
+
+def _resolve_cpu_set(spec: str | None, field: str) -> set[int]:
+    """Resolve one worker's CPU affinity from a spec like ``"0:5,7"``.
+
+    Each comma-separated item is a single core id or an inclusive ``A:B``
+    range; cores outside the process's available set are dropped. An empty
+    spec or an empty effective set disables pinning.
+    """
+    if not spec:
+        return set()
+    available = _available_cpu_ids()
+    cpus: set[int] = set()
+    for item in spec.split(","):
+        item = item.strip()
+        if not item:
+            continue
+        try:
+            if ":" in item:
+                start_s, _, end_s = item.partition(":")
+                start, end = int(start_s), int(end_s)
+            else:
+                start = end = int(item)
+        except ValueError as exc:
+            raise ValueError(
+                f"{field}='{spec}': invalid core spec item '{item}' "
+                "(expected core ids or inclusive A:B ranges, e.g. '0:5,7')"
+            ) from exc
+        if start > end:
+            raise ValueError(f"{field}='{spec}': range '{item}' has start > end")
+        # Clamp to the available core span instead of enumerating the raw
+        # range: a typo like 0:1000000000 must not iterate billions of ids.
+        lo, hi = min(available), max(available)
+        cpus.update(range(max(start, lo), min(end, hi) + 1))
+    return cpus & available
+
+
+def _pin_worker_cpus(cpus: set[int]) -> None:
+    """Pin this worker process and cap its torch threads to its CPU slice.
+
+    The thread cap applies even when the affinity call fails (cgroup/cpuset
+    constraints): limiting torch to the configured core count still reduces
+    CPU contention when running unpinned.
+    """
+    if not cpus:
+        return
+    try:
+        _set_cpu_affinity(cpus)
+    except OSError as exc:
+        logging.getLogger(__name__).warning("CPU pinning to %s failed (%s); continuing unpinned", sorted(cpus), exc)
+    torch.set_num_threads(max(len(cpus), 1))
+
+
 def _configure_process_logging() -> None:
     """Surface INFO logs (e.g. manager env startup) from spawned worker processes.
 
@@ -150,6 +243,8 @@ def run_collector_process(
     try:
         _configure_process_logging()
         set_seed(seed)
+        opts = cfg.trainer.async_options
+        _pin_worker_cpus(_resolve_cpu_set(opts.collector_cpu_cores, "collector_cpu_cores"))
         async_options = cfg.trainer.async_options
         obs_dim, critic_obs_dim, act_dim = dims
         device = torch.device("cpu")
@@ -216,6 +311,8 @@ def run_learner_process(
     console, live = open_training_live()
     try:
         set_seed(seed)
+        opts = cfg.trainer.async_options
+        _pin_worker_cpus(_resolve_cpu_set(opts.learner_cpu_cores, "learner_cpu_cores"))
         async_options = cfg.trainer.async_options
         device = torch.device(cfg.device or ("cuda" if torch.cuda.is_available() else "cpu"))
         writer = None

motrix_rl/src/motrix_rl/fastsac/config.py

@@ -69,6 +69,14 @@ class FastSacAsyncOptionsCfg:
     collector_compile: bool = MISSING
     collector_amp: bool = MISSING
     collector_amp_dtype: str = MISSING
+    # CPU affinity pinning for the async learner and collector processes, as a
+    # core spec like "0:5,7" (cores 0-5 and 7). Out-of-range and duplicate
+    # cores are dropped; an empty effective set or null disables pinning. The
+    # env step is CPU-bound while the learner is GPU-bound, so isolating them
+    # keeps learner-side drain/ingest work from stealing cores from the
+    # collector.
+    learner_cpu_cores: str | None = None
+    collector_cpu_cores: str | None = None
 
 
 @dataclass

motrix_rl/tests/test_rl_sim_backend.py

@@ -241,6 +241,8 @@ def _async_cfg():
                 collector_compile=False,
                 collector_amp=False,
                 collector_amp_dtype="fp16",
+                learner_cpu_cores=None,
+                collector_cpu_cores=None,
             )
         ),
     )

scripts/bench_fastsac_collector_inference.py

@@ -93,7 +93,9 @@ def _collector(args, source_actor, source_normalizer, action_scale, action_bias)
         ),
     )
     env = _BenchmarkEnv(args.num_envs, args.obs_dim, args.critic_obs_dim)
-    ring = SharedTransitionRing(1, args.num_envs, args.obs_dim, args.critic_obs_dim, args.act_dim)
+    # Stub ring (never pushed: the bench only exercises inference/sync); capacity
+    # 2 satisfies the ring's successor-slot contract.
+    ring = SharedTransitionRing(2, args.num_envs, args.obs_dim, args.critic_obs_dim, args.act_dim)
     weights = WeightSnapshot(sum(p.numel() for p in source_actor.parameters()), args.obs_dim)
     weights.publish(source_actor, source_normalizer)
     collector = Collector(