From aa3bde254dcd1df91323864832df9bea583e7598 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Sun, 20 Sep 2026 23:26:09 +0800 Subject: [PATCH] perf(fastsac): fix async collector push copy redundancy and sync blocking Closes #61, Closes #62 - halve transition-ring push cost on the CPU-env path (shm.py ring push copy optimization); the GPU-env async D2H staging path is not wired up yet and will be added together with GPU-env support (#61) - per-side CPU affinity pinning for the async learner/collector processes, configured as core specs like "0:5,7" via learner_cpu_cores / collector_cpu_cores (#62) Benchmark (g1-sonic 4096 envs): env_step 55 -> ~47ms, steady throughput ~59k -> ~65k env-steps/s. --- configs/algo_base/motrix.fastsac.yaml | 6 ++ .../motrix_rl/fastsac/async_impl/collector.py | 25 +++-- .../motrix_rl/fastsac/async_impl/learner.py | 13 ++- .../src/motrix_rl/fastsac/async_impl/shm.py | 56 +++++++---- .../motrix_rl/fastsac/async_impl/worker.py | 97 +++++++++++++++++++ motrix_rl/src/motrix_rl/fastsac/config.py | 8 ++ motrix_rl/tests/test_rl_sim_backend.py | 2 + scripts/bench_fastsac_collector_inference.py | 4 +- 8 files changed, 179 insertions(+), 32 deletions(-) diff --git a/configs/algo_base/motrix.fastsac.yaml b/configs/algo_base/motrix.fastsac.yaml index f0969aa2..2c518919 100644 --- a/configs/algo_base/motrix.fastsac.yaml +++ b/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 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 4929b01d..f39623e3 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py +++ b/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() 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 2494387e..d5ce7a3c 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py +++ b/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), diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py index 136b705f..e55422c5 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm.py +++ b/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 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 f90a9b89..a3faac01 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/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 diff --git a/motrix_rl/src/motrix_rl/fastsac/config.py b/motrix_rl/src/motrix_rl/fastsac/config.py index d2d6b265..d5096403 100644 --- a/motrix_rl/src/motrix_rl/fastsac/config.py +++ b/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 diff --git a/motrix_rl/tests/test_rl_sim_backend.py b/motrix_rl/tests/test_rl_sim_backend.py index 30d6e8f4..0e77fedc 100644 --- a/motrix_rl/tests/test_rl_sim_backend.py +++ b/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, ) ), ) diff --git a/scripts/bench_fastsac_collector_inference.py b/scripts/bench_fastsac_collector_inference.py index 4788699d..3872ec44 100644 --- a/scripts/bench_fastsac_collector_inference.py +++ b/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(