From ae00452e7bcecf4435c9110e794c569464178f44 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Wed, 23 Sep 2026 13:29:08 +0800 Subject: [PATCH] feat(fastsac): support multiple collector processes with automatic NUMA binding --- configs/algo_base/motrix.fastsac.yaml | 11 + motrix_rl/src/motrix_rl/console.py | 37 +- motrix_rl/src/motrix_rl/fastsac/agent.py | 126 ++- .../motrix_rl/fastsac/async_impl/collector.py | 12 +- .../motrix_rl/fastsac/async_impl/learner.py | 477 ++++++++--- .../src/motrix_rl/fastsac/async_impl/numa.py | 305 +++++++ .../motrix_rl/fastsac/async_impl/panels.py | 186 +++++ .../src/motrix_rl/fastsac/async_impl/stats.py | 73 ++ .../motrix_rl/fastsac/async_impl/topology.py | 372 +++++++++ .../src/motrix_rl/fastsac/async_impl/train.py | 708 +++++++++++++--- .../fastsac/async_impl/transport/__init__.py | 2 + .../fastsac/async_impl/transport/common.py | 48 +- .../fastsac/async_impl/transport/handshake.py | 115 +++ .../motrix_rl/fastsac/async_impl/worker.py | 768 ++++++++++-------- motrix_rl/src/motrix_rl/fastsac/buffer.py | 83 +- motrix_rl/src/motrix_rl/fastsac/config.py | 12 + motrix_rl/src/motrix_rl/system_metrics.py | 5 + motrix_rl/tests/fastsac_async_mocks.py | 37 + motrix_rl/tests/test_fastsac_async_multi.py | 556 +++++++++++++ motrix_rl/tests/test_fastsac_boot_panel.py | 167 ++++ motrix_rl/tests/test_fastsac_collector.py | 2 - .../tests/test_fastsac_ddp_equivalence.py | 170 ++++ motrix_rl/tests/test_fastsac_ipc_ring.py | 110 ++- motrix_rl/tests/test_fastsac_learner.py | 24 +- .../test_fastsac_pipeline_equivalence.py | 185 +++++ motrix_rl/tests/test_rl_sim_backend.py | 41 +- .../fastsac-async-heterogeneous-trainer.md | 82 +- wiki/plan/fastsac-async-multi-learner.md | 17 + wiki/plan/index.md | 2 +- 29 files changed, 4068 insertions(+), 665 deletions(-) create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/numa.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/panels.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/stats.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/transport/handshake.py create mode 100644 motrix_rl/tests/fastsac_async_mocks.py create mode 100644 motrix_rl/tests/test_fastsac_async_multi.py create mode 100644 motrix_rl/tests/test_fastsac_boot_panel.py create mode 100644 motrix_rl/tests/test_fastsac_ddp_equivalence.py create mode 100644 motrix_rl/tests/test_fastsac_pipeline_equivalence.py create mode 100644 wiki/plan/fastsac-async-multi-learner.md diff --git a/configs/algo_base/motrix.fastsac.yaml b/configs/algo_base/motrix.fastsac.yaml index 68974902..5e3291d8 100644 --- a/configs/algo_base/motrix.fastsac.yaml +++ b/configs/algo_base/motrix.fastsac.yaml @@ -117,3 +117,14 @@ trainer: weight_ipc: "auto" # Minimum flattened parameter bytes for the CUDA-IPC path under "auto". weight_ipc_min_bytes: 16777216 + # Number of collector processes; num_envs is split evenly across them (must divide evenly). + num_collectors: 1 + # NUMA nodes are assigned automatically: collectors spread round-robin over + # the host's nodes, learners bind to their GPU-local node. + # CPUs assigned to each collector from its node/affinity set; null uses all of them. + cpus_per_collector: null + # DDP data-parallel learner replicas (one process per GPU, single node); + # requires num_collectors to divide evenly; 1 keeps the single-learner behavior. + num_learners: 1 + # One device per learner, e.g. [cuda:0, cuda:1]; null replicates the trainer device. + learner_devices: null diff --git a/motrix_rl/src/motrix_rl/console.py b/motrix_rl/src/motrix_rl/console.py index 9b6715ed..4bbf69b2 100644 --- a/motrix_rl/src/motrix_rl/console.py +++ b/motrix_rl/src/motrix_rl/console.py @@ -46,7 +46,7 @@ class TrainingPanelStats: iteration: int total_iterations: int - steps_per_second: float + steps_per_second: float | None elapsed_seconds: float mean_return: float mean_episode_length: float @@ -57,6 +57,9 @@ class TrainingPanelStats: learn_ms: float learn_percent: float warming: bool = False + # Rolling-window iteration rate; None = same "warming up" semantics as + # steps_per_second. Falls back to the cumulative average when absent. + iterations_per_second: float | None = None training_metrics: Mapping[str, Any] | None = None reward_terms: Mapping[str, Any] = field(default_factory=dict) env_metrics: Mapping[str, Any] = field(default_factory=dict) @@ -187,6 +190,14 @@ def _eta_seconds(stats: TrainingPanelStats) -> float | None: return (stats.total_iterations - stats.iteration) / rate +def _iter_rate_text(stats: TrainingPanelStats) -> str: + if stats.iterations_per_second is not None: + return f"{stats.iterations_per_second:,.0f} iter/s" + if stats.steps_per_second is None: + return "warming up" + return f"{stats.iteration / max(stats.elapsed_seconds, 1e-9):,.0f} iter/s" + + def format_training_panel(stats: TrainingPanelStats, *, title: str = "rl") -> str: """Render a plain-text RL training panel from backend-provided scalar stats.""" @@ -203,9 +214,10 @@ def si(n: float) -> str: pct = 100.0 * stats.iteration / max(stats.total_iterations, 1) eta = _eta_seconds(stats) eta_text = f" - eta ~{hms(eta)}" if eta is not None else "" + sps_text = f"{stats.steps_per_second:.0f} env-steps/s" if stats.steps_per_second is not None else "warming up" header = ( f" {title} - iter {stats.iteration}/{stats.total_iterations} ({pct:.1f}%) - " - f"{stats.steps_per_second:.0f} env-steps/s - {hms(stats.elapsed_seconds)}{eta_text}" + f"{sps_text} - {hms(stats.elapsed_seconds)}{eta_text}" ) lines = [ "-" * width, @@ -477,7 +489,9 @@ def _prototype_bar(fraction: float, *, width: int = 22, style: str = "cyan"): return result -def _cpu_spectrum_rows(per_core: Sequence[float], load_style, per_row: int = 48) -> list: +def _cpu_spectrum_rows( + per_core: Sequence[float], load_style, core_ids: Sequence[int] | None = None, per_row: int = 48 +) -> list: """Per-core CPU utilization spectrum rows for the System view. Style is selected by ``MOTRIX_PANEL_CPU_SPECTRUM`` (height | shade | @@ -505,7 +519,11 @@ def _cpu_spectrum_rows(per_core: Sequence[float], load_style, per_row: int = 48) label_width = 14 # "cores 192-199 " for start in range(0, len(per_core), per_row): chunk = per_core[start : start + per_row] - label = f"cores {start}-{start + len(chunk) - 1}".ljust(label_width) + if core_ids is not None: + ids = core_ids[start : start + per_row] + label = f"cores {ids[0]}-{ids[-1]}".ljust(label_width) + else: + label = f"cores {start}-{start + len(chunk) - 1}".ljust(label_width) if style in ("height", "shade"): # Leading space: idle cores render blank so the dim ▁ ceiling caps # are the ONLY ▁ on screen (otherwise an idle column would look @@ -557,7 +575,7 @@ def _prototype_system_page(stats: TrainingPanelStats, load_style, memory_style): if load is not None and load.model_name: cpu_parts.append(Text(load.model_name, style="dim")) if load is not None and load.per_core_percent: - cpu_parts.extend(_cpu_spectrum_rows(load.per_core_percent, load_style)) + cpu_parts.extend(_cpu_spectrum_rows(load.per_core_percent, load_style, load.per_core_ids)) else: # RAM is not repeated here — the always-visible System health card in # the summary row already carries it. @@ -752,8 +770,13 @@ def memory_style(memory: MemoryUsage | None) -> str: card( "Throughput", Group( - Text(f"{stats.steps_per_second:,.0f} env-steps/s", style="bold cyan"), - Text(f"{stats.iteration / max(stats.elapsed_seconds, 1e-9):,.0f} iter/s", style="white"), + Text( + f"{stats.steps_per_second:,.0f} env-steps/s" + if stats.steps_per_second is not None + else "warming up", + style="bold cyan", + ), + Text(_iter_rate_text(stats), style="white"), ), ), card( diff --git a/motrix_rl/src/motrix_rl/fastsac/agent.py b/motrix_rl/src/motrix_rl/fastsac/agent.py index 332f7ef1..3daa1f45 100644 --- a/motrix_rl/src/motrix_rl/fastsac/agent.py +++ b/motrix_rl/src/motrix_rl/fastsac/agent.py @@ -16,6 +16,7 @@ import time import torch +import torch.distributed as dist import torch.nn.functional as F from torch import nn, optim @@ -55,6 +56,7 @@ def __init__( action_scale: torch.Tensor | None = None, action_bias: torch.Tensor | None = None, writer=None, + world_size: int = 1, ): """Build the actor, twin distributional critics, optimizers and replay buffer. @@ -63,10 +65,15 @@ def __init__( sizes the replay buffer's per-env rings. ``action_scale`` / ``action_bias`` map the tanh-squashed policy output into the environment's action range (defaults to identity when ``None``). AMP autocast and ``torch.compile`` - are enabled per ``cfg`` but auto-disabled on CPU. + are enabled per ``cfg`` but auto-disabled on CPU. ``world_size > 1`` + (process group already initialized by the caller) averages gradients + across ranks manually after each backward and splits ``batch_size`` + across ranks; the all-reduce stays in the eager orchestrator so the + compiled halves remain CUDA-graph capturable (see _update_main). """ self.cfg = cfg self.device = device + self.world_size = world_size self.obs_dim = obs_dim self.critic_obs_dim = critic_obs_dim self.act_dim = act_dim @@ -154,17 +161,26 @@ def __init__( # Runtime callables default to the canonical modules. Checkpointing, # optimizer ownership and inter-process weight publication always use - # the canonical modules so torch.compile remains a runtime-only detail. + # the canonical modules so torch.compile remains a runtime-only + # detail. Multi-learner uses the canonical modules too and averages + # gradients manually after each backward (see _allreduce_module_grads): + # the update boundary calls custom methods (get_actions_and_log_probs, + # projection, get_value) rather than module.forward, which DDP wrappers + # neither proxy nor synchronize. self._actor_runtime = self.actor self._qnet_runtime = self.qnet self._qnet_target_runtime = self.qnet_target - self._update_main_runtime = self._update_main - self._update_pol_runtime = self._update_pol - - # Stage A compile boundary: compile the complete learner update rather - # than compiling individual network modules. This follows Holosoma's - # FastSAC structure and avoids nested eager/compiled boundaries around - # actor sampling, projection, loss, backward, and optimizers. + self._critic_backward_runtime = self._critic_backward + self._actor_backward_runtime = self._actor_backward + + # Stage A compile boundary: compile the PURE-COMPUTE halves of the + # update (forward + backward, no cross-process sync, no optimizer + # step) rather than individual network modules. The eager orchestrator + # methods (_update_main / _update_pol) own gradient averaging, clipping + # and optimizer steps: CUDA graphs cannot capture NCCL collectives, so + # keeping the all-reduce outside the compiled region lets DDP ranks + # (world_size > 1) use the same reduce-overhead graphs as the + # single-learner path instead of falling back to eager. if bool(cfg.compile) and device.type == "cuda": import torch._inductor.config as inductor_config @@ -177,8 +193,8 @@ def __init__( # 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") + self._critic_backward_runtime = torch.compile(self._critic_backward, mode="reduce-overhead") + self._actor_backward_runtime = torch.compile(self._actor_backward, mode="reduce-overhead") # --------------------------------------------------------------- normalize @staticmethod @@ -187,6 +203,27 @@ def _norm(normalizer: nn.Module, obs: torch.Tensor, update: bool) -> torch.Tenso return normalizer(obs, update=update) return obs + def _allreduce_module_grads(self, modules) -> None: + """Average the gradients of the given modules across DDP ranks. + + One fused all-reduce per call (all grads concatenated into a flat + buffer), replacing DDP's per-backward hook sync — the update boundary + calls custom module methods, which a DDP wrapper would neither proxy + nor synchronize. No-op at world_size 1. + """ + if self.world_size <= 1: + return + params = [p for m in modules for p in m.parameters() if p.grad is not None] + if not params: + return + flat = torch.cat([p.grad.reshape(-1) for p in params]) + dist.all_reduce(flat, op=dist.ReduceOp.AVG) + offset = 0 + for p in params: + n = p.grad.numel() + p.grad.copy_(flat[offset : offset + n].view_as(p.grad)) + offset += n + def _autocast(self): """torch.autocast context manager, or a no-op when AMP is disabled. @@ -200,7 +237,15 @@ def _autocast(self): return contextlib.nullcontext() # --------------------------------------------------------------- updates - def _update_main(self, b: dict): + def _critic_backward(self, b: dict): + """Pure-compute half of the critic update: forward + backward only. + + Everything here is capturable by a CUDA graph (no cross-process + collective, no optimizer mutation of parameters outside the graph's + static inputs). Returns owned scalar stats; ``next_logp_mean`` feeds + the eager alpha update — mean commutes with the constant + ``target_entropy``, so the scalar is exactly the alpha loss basis. + """ cfg = self.cfg rewards = b["rewards"] dones = b["dones"].bool() @@ -227,17 +272,25 @@ def _update_main(self, b: dict): self.q_optimizer.zero_grad(set_to_none=True) qf_loss.backward() + return ( + qf_loss.detach().float().clone(), + next_logp.detach().float().mean().clone(), + target_values.max().detach().float().clone(), + target_values.min().detach().float().clone(), + ) + + def _update_main(self, b: dict): + cfg = self.cfg + qf_loss, next_logp_mean, tv_max, tv_min = self._critic_backward_runtime(b) + self._allreduce_module_grads([self.qnet]) if cfg.max_grad_norm > 0: 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. + # Target-network soft update. Ordering constraints: after the target + # read at the top of _critic_backward 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()] @@ -246,19 +299,19 @@ def _update_main(self, b: dict): alpha_loss = torch.zeros((), device=self.device) if cfg.use_autotune: - alpha_loss = (-self.log_alpha.exp() * (next_logp.detach() + self.target_entropy)).mean() + alpha_loss = -self.log_alpha.exp() * (next_logp_mean + self.target_entropy) self.alpha_optimizer.zero_grad(set_to_none=True) alpha_loss.backward() + if self.world_size > 1 and self.log_alpha.grad is not None: + # log_alpha is a bare tensor, not inside a module; average its + # gradient manually to keep every rank's temperature identical. + dist.all_reduce(self.log_alpha.grad, op=dist.ReduceOp.AVG) self.alpha_optimizer.step() - return ( - qf_loss.detach().float(), - alpha_loss.detach().float(), - target_values.max().detach().float(), - target_values.min().detach().float(), - ) + return (qf_loss, alpha_loss.detach().float(), tv_max, tv_min) - def _update_pol(self, b: dict): + def _actor_backward(self, b: dict): + """Pure-compute half of the actor update: forward + backward only.""" with self._autocast(): actions, log_probs = self._actor_runtime.get_actions_and_log_probs(b["obs"]) q_outputs = self._qnet_runtime(b["critic_obs"], actions) @@ -268,10 +321,18 @@ def _update_pol(self, b: dict): self.actor_optimizer.zero_grad(set_to_none=True) actor_loss.backward() + return actor_loss.detach().float().clone(), (-log_probs.mean()).detach().float().clone() + + def _update_pol(self, b: dict): + actor_loss, neg_logp = self._actor_backward_runtime(b) + # The policy loss backpropagates into BOTH the actor and the critic (the + # critic's q_values feed the objective); averaging both keeps every + # rank's parameters identical, matching what DDP would sync. + self._allreduce_module_grads([self.actor, self.qnet]) if self.cfg.max_grad_norm > 0: torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.cfg.max_grad_norm) self.actor_optimizer.step() - return actor_loss.detach().float(), (-log_probs.mean()).detach().float() + return actor_loss, neg_logp def update(self, num_updates: int): """Run ``num_updates`` gradient steps, each on a fresh batch. @@ -296,7 +357,10 @@ def update(self, num_updates: int): cfg = self.cfg if num_updates <= 0: return None - batch_per_env = max(cfg.batch_size // self.num_envs, 1) + # Global batch split across DDP ranks (== batch_size when world_size 1); + # gradient averaging makes the update equivalent to the sync + # trainer's single global-batch step. + batch_per_env = max(cfg.batch_size // self.world_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")} update_started = time.perf_counter() @@ -332,13 +396,13 @@ def update(self, num_updates: int): # any output that outlives its own iteration goes through `_own`. torch.compiler.cudagraph_mark_step_begin() stage_started = time.perf_counter() - outputs = self._update_main_runtime(b) + outputs = self._update_main(b) timing_s["critic_alpha"] += time.perf_counter() - stage_started actor_pair = (last[3], last[4]) if (self.update_idx + i) % cfg.policy_frequency == 0: stage_started = time.perf_counter() - pol_outputs = self._update_pol_runtime(b) + pol_outputs = self._update_pol(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 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 48559c9a..ad9b4c4c 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py @@ -72,7 +72,6 @@ def __init__( env: FastSacEnvWrap, cfg: FastSacCfg, obs_dim: int, - critic_obs_dim: int, act_dim: int, action_scale: torch.Tensor, action_bias: torch.Tensor, @@ -80,6 +79,7 @@ def __init__( weights: WeightReceiver, control: Control, is_resume: bool = False, + collector_id: int = 0, ): self.env = env self.cfg = cfg @@ -94,6 +94,7 @@ def __init__( self.ring = ring self.weights = weights self.control = control + self.collector_id = collector_id self.is_resume = is_resume # Env-internal step profiling (Perf on the wrapped env, when present) # feeds the panel's env_step sub-stage tree. ``env.env`` is the @@ -227,7 +228,6 @@ def policy_lag(self) -> int: """How many published versions behind the collector's local policy is.""" return self.weights.lag - # ------------------------------------------------------------------ ring handoff # ------------------------------------------------------------------ step def step_once(self) -> bool: """Run one env-step batch and push it to the ring. @@ -247,7 +247,8 @@ def step_once(self) -> bool: self._wait_started = None t0 = now - warming = self.control.collector_steps < self._learning_starts + steps = self.control.collector_steps_at(self.collector_id) + warming = steps < self._learning_starts t_sample_actions = time.perf_counter() actions = self._sample_actions(warming) t_env = time.perf_counter() @@ -297,10 +298,10 @@ def step_once(self) -> bool: self.obs = next_obs self.critic_obs = next_critic_obs - self.control.inc_collector_steps() + self.control.inc_collector_steps(self.collector_id) t_sync = time.perf_counter() - if self.control.collector_steps % max(self.async_options.weight_poll_interval, 1) == 0: + if self.control.collector_steps_at(self.collector_id) % max(self.async_options.weight_poll_interval, 1) == 0: self.sync_weights(record_timing=True) t_done = time.perf_counter() @@ -324,6 +325,7 @@ def snapshot_stats(self) -> dict: rr, rl = self.recent_returns, self.recent_lengths term_means = {k: v / max(self.term_count, 1) for k, v in self.term_accum.items()} stats = { + "collector_id": self.collector_id, "return": (sum(rr) / len(rr)) if rr else float("nan"), "ep_len": (sum(rl) / len(rl)) if rl else float("nan"), "episodes": self.n_episodes, 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 5086b5cd..80ba0dc7 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py @@ -18,57 +18,200 @@ from __future__ import annotations import time +from dataclasses import dataclass import torch +import torch.distributed as dist from motrix_rl.fastsac.agent import FastSacAgent 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.buffer import EmpiricalNormalization from motrix_rl.fastsac.config import FastSacCfg +@dataclass +class CollectorEndpoint: + """The learner-side endpoints for one collector. + + Both halves belong to collector ``i`` of this rank: the ring the + collector's transitions arrive on, and the weight channel they are + published back to. Index alignment is the ownership relation — a + misordered pair silently cross-wires two collectors. + """ + + ring: SharedTransitionRing | IpcTransitionRing + weight_sender: WeightSender + + +class GenerationAssembler: + """Cross-collector generation assembly: shards in, complete generations out. + + Shards are indexed by their ABSOLUTE generation number (the ring slot's + ordinal since process start), so collectors delivering at different rates + land on the same timeline. A generation is complete once every ring's + shard has arrived; :meth:`pop_ready` pops complete generations in strict + ordinal order, and incomplete ones wait here indefinitely without + blocking the other rings (their own rings backpressure them instead). + """ + + def __init__(self, num_rings: int): + self.num_rings = num_rings + self.next_gen = 0 + self._slots: dict[int, list[tuple | None]] = {} + + def add(self, ring_id: int, gen: int, part: tuple) -> None: + """Place one ring's shard at its absolute generation slot.""" + slots = self._slots.setdefault(gen, [None] * self.num_rings) + slots[ring_id] = part + + def pop_ready(self) -> list[list[tuple]]: + """Pop consecutive complete generations starting at the cursor.""" + ready: list[list[tuple]] = [] + while (slots := self._slots.get(self.next_gen)) is not None and None not in slots: + ready.append(slots) + del self._slots[self.next_gen] + self.next_gen += 1 + return ready + + def stats(self) -> tuple[int, int, int, int]: + """``(next_gen, oldest pending gen, newest pending gen, pending count)``.""" + return ( + self.next_gen, + min(self._slots) if self._slots else -1, + max(self._slots) if self._slots else -1, + len(self._slots), + ) + + +class DrainStaging: + """Pinned-staging H2D pipeline for host-ring shards on a CUDA learner. + + Exists only when the learner runs on CUDA and owns at least one host + ring. Host-ring views are memcpy'd into pinned buffers and moved to the + GPU with one non-blocking H2D copy per field on a dedicated stream, so + ingest overlaps gradient updates. IPC shards need none of this — their + slots are already on the device (see _drain_generations' host pull-down). + + Lifecycle discipline: + + * :meth:`acquire` — wait before overwriting buffers that may still back + the previous drain's in-flight copies; + * :meth:`stage` — host views -> pinned -> async H2D ``rb.extend_batch``; + * :meth:`mark_in_flight` — record the event at the end of a drain, so + later sampling waits for the copies; + * :meth:`wait` — make the compute stream wait for in-flight copies + before sampling; + * :meth:`finish` — block until everything landed (end of ingesting). + """ + + def __init__(self, rb, chunk: int, device: torch.device): + self.rb = rb + self.stream = torch.cuda.Stream(device=device) + self.event = torch.cuda.Event() + self.pending = False + pin = lambda *shape: torch.empty(*shape, pin_memory=True) # noqa: E731 + self.buffers = ( + 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), + ) + + def acquire(self) -> None: + """Join in-flight copies before the buffers get overwritten.""" + self.stream.synchronize() + + def stage(self, views) -> None: + """Copy one host run into the pinned buffers and issue the H2D ingest.""" + k = views[0].shape[0] + for stage, view in zip(self.buffers, views): + stage[:k].copy_(view) # ring -> pinned (plain CPU memcpy) + with torch.cuda.stream(self.stream): + self.rb.extend_batch(*(stage[:k] for stage in self.buffers)) + self.pending = True + + def mark_in_flight(self) -> None: + """Record the event that later ``wait`` calls synchronize on.""" + if self.pending: + self.event.record(self.stream) + + def wait(self) -> None: + """Make the compute stream wait for in-flight copies, then clear.""" + if self.pending: + self.event.wait() + self.pending = False + + def finish(self) -> None: + """Block until every issued ingest copy has landed in the buffer.""" + self.stream.synchronize() + self.pending = False + + class Learner: + """Drains N independent SPSC rings (one per collector) and broadcasts weights. + + Each collector owns its own ring and its own weight-channel sender; the + single-writer invariants of both primitives are preserved unchanged. The + learner round-robins the rings (per-ring ingest bound so a fast collector's + full ring never starves the others) and publishes the current actor snapshot + to every collector in turn — each publish is independent and lock-free, so + one slow reader never blocks the others. + + Ring slots carry ``num_envs / num_collectors`` transitions each. The replay + buffer (and its n-step time adjacency per env) is built around full + ``num_envs`` batches, so shards are merged by generation: the k-th slot of + every collector assembles into the k-th full batch (env ids map to + contiguous shard blocks). A slow collector's generation stays pending on + the CPU slot views until complete — its own ring fills up and backpressures + only that collector; read cursors are committed only after the merged batch + reached the GPU, preserving the ring's no-clobber guarantee. + """ + def __init__( self, agent: FastSacAgent, cfg: FastSacCfg, - ring: SharedTransitionRing | IpcTransitionRing, - weights: WeightSender, + collectors: list[CollectorEndpoint], control: Control, + ddp_rank: int | None = None, ): self.agent = agent self.cfg = cfg self.async_options = cfg.trainer.async_options - self.ring = ring - self.weights = weights + self.collectors = collectors + self.rings = [collector.ring for collector in collectors] + self.weights = [collector.weight_sender for collector in collectors] self.control = control + self._pending = GenerationAssembler(num_rings=len(self.rings)) + # Ring transport per LEARNER RANK is all-or-nothing (topology), so + # each rank gets its own ingest path: a PURE-HOST rank on CUDA runs + # the staged async H2D pipeline (DrainStaging); pure-IPC and MIXED + # ranks lift host shards at assembly and consume everything D2D + # (see _drain_rings), so staging would be dead weight there. + has_host = any(not isinstance(ring, IpcTransitionRing) for ring in self.rings) + has_ipc = any(isinstance(ring, IpcTransitionRing) for ring in self.rings) + self._staging: DrainStaging | None = None + if agent.device.type == "cuda" and has_host and not has_ipc: + self._staging = DrainStaging(agent.rb, max(self.async_options.max_ingest_per_iter, 1), agent.device) + # This rank owns and publishes to its own collectors only: the slice + # must cover exactly this rank's share of the global collector count + # (the total lives in the shared control block). + per_learner = control.num_collectors // agent.world_size + if len(collectors) != per_learner: + raise ValueError( + f"{len(collectors)} collector endpoints must cover this rank's slice of " + f"{per_learner} collectors ({control.num_collectors} total)" + ) + # DDP replica id (multi-learner); None keeps the single-learner path. + self.ddp_rank = ddp_rank 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), - ) + # global-progress basis consumed by the last lockstep update call + # (multi-learner path only); the worker re-bases it on resume. + self._last_train_gstep = 0 # keep normalizers/actor in train mode: the learner is the update side. self.agent.set_train_mode() @@ -82,104 +225,226 @@ def update_idx(self) -> int: # ------------------------------------------------------------------ ingest def drain(self) -> int: - """Move up to ``max_ingest_per_iter`` ring slots into the replay buffer. + """Move up to ``max_ingest_per_iter`` slots per ring into the replay buffer. - Returns the number of slots ingested. Slots are consumed in contiguous - runs (``ring.read_span()``): + Returns the number of full ``num_envs`` batches ingested. Slots are + assembled into generations (one shard per collector) and complete + generations are merged, moved to the GPU and written with one + :meth:`extend_batch` per field per flush; a single ring is the + ``num_rings == 1`` special case (its shard alone is the generation, + no merge needed). Read cursors advance only after the GPU copy, so a + collector cannot clobber an in-flight slot; a full or slow ring only + blocks itself. + + Per-ring transport notes: * 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. + any subsequent sample by stream order). + * host ring on a pure-host CUDA rank: each shard 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). + * host ring in a MIXED rank: the shard is lifted to the ingest device + at assembly, so the merge never leaves the GPU. * 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. + following slot's stored observation, so no successor peek is needed + and a slot is ingested as soon as it is committed. """ - budget = max(self.async_options.max_ingest_per_iter, 1) + return self._drain_rings(max(self.async_options.max_ingest_per_iter, 1)) + + def _drain_rings(self, budget: int) -> int: + for ring_id, ring in enumerate(self.rings): + base = ring.read_idx + avail = min(budget, ring.size()) + if avail <= 0: + continue + count, views = ring.read_span() + # PURE-HOST rank (self._staging): keep the host views — the staged + # pipeline moves them with async H2D on the copy stream, + # overlapping the next gradient update. MIXED rank: lift host + # shards to the device at assembly so the merge is single-device + # D2D; IPC shards are already there. Host ring on CPU consumes + # directly into the buffer. + if self._staging is None and not isinstance(ring, IpcTransitionRing): + views = tuple(view.to(self.agent.device) for view in views) + count = min(count, avail) + for offset in range(count): + self._pending.add(ring_id, base + offset, tuple(v[offset] for v in views)) + return self._flush_complete_generations() + + def _stall_diag(self, tag: str) -> None: + import logging + import os + + if not os.environ.get("MOTRIX_STALL_DEBUG"): + return + next_gen, oldest, newest, count = self._pending.stats() + logging.getLogger(__name__).error( + "STALL[%s] next_gen=%d pending=[%s..%s] n_pending=%d rings=[%s]", + tag, + next_gen, + oldest, + newest, + count, + [(r.read_idx, r.write_idx) for r in self.rings], + ) + + def _flush_complete_generations(self) -> int: + generations = self._pending.pop_ready() + if not generations: + self._stall_diag("flush-empty") + return 0 + if self._staging is not None: + # staging may still back the previous drain's in-flight H2D copies + self._staging.acquire() + # env blocks concatenate in collector order -> env id mapping is + # stable across batches, so per-env trajectories stay contiguous in + # the buffer's time dimension. A single shard skips the cat. + runs = [ + parts[0] if len(parts) == 1 else tuple(torch.cat(field) for field in zip(*parts)) for parts in generations + ] + # runs arrive env-major per generation (n_env, dim); the buffer write + # path is slot-major (k, n_env, dim), so stack generations first. + stacked = tuple(torch.stack(column, dim=0) for column in zip(*runs)) + ingested = self._ingest_runs((stacked,)) + for ring in self.rings: + ring.commit_reads(ingested) 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 - 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) + self._staging.mark_in_flight() return ingested + def _ingest_runs(self, runs) -> int: + """Ingest slot-major device runs ``(k, n_env, dim)``; caller commits reads. + + Every shard was lifted to the ingest device before assembly, so each + ``extend_batch`` is a same-device strided copy. Returns the number of + RING SLOTS consumed — each run's leading dim IS the slot count; + returning ``len(runs)`` instead desynced the read cursor from + ``_next_gen`` by G-1 per flush (learner stalls permanently once the + gap exceeds the pending window). + """ + if self._staging is not None: + # pure-host rank: host run -> pinned buffers -> async H2D extend + self._staging.stage(runs[0]) + return runs[0][0].shape[0] + k = 0 + for views in runs: + self.agent.rb.extend_batch(*views) + k += views[0].shape[0] + return k + # ------------------------------------------------------------------ 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 + if self._staging is not None: + self._staging.finish() - def _ready(self) -> bool: - return self.control.collector_steps >= self._learning_starts and self.agent.rb.num_stored > 0 + def maybe_train(self, gstep: int) -> dict | None: + """Run ratio-governed updates when the global progress basis advanced. - def _num_updates_for(self, ingested: int) -> int: - """Decide how many gradient updates to run this iteration.""" - base = self.agent.cfg.num_updates - mode = self.async_options.utd_mode - if mode == "strict": - # exactly num_updates per ingested env-step batch -> matches sync UTD. - return ingested * base - # learner_bound: run a full batch of updates whenever ready. - # In the two-process path the learner loops continuously; here per-call. - return base - - def maybe_train(self, ingested: int) -> dict | None: - """Run ratio-governed updates. Returns last metrics dict or ``None``.""" - if not self._ready(): + ``gstep`` is the trainer's full-batch-equivalent progress + (``control.collector_steps // num_collectors``) — identical on every + DDP rank, and the single learner's own position when + ``num_collectors == 1``. One decision formula serves both topologies: + + * single learner: decides locally; an empty replay buffer merely + defers the update (the cumulative delta covers it next call); + * multi-learner: rank 0 decides, the count is broadcast, and every + rank waits for data rather than skipping — skipping would desync + the DDP collectives. + """ + is_src = self.ddp_rank in (None, 0) + n = 0 + if is_src: + warmup_met = self.control.collector_steps >= self._learning_starts * self.control.num_collectors + if gstep > self._last_train_gstep and warmup_met: + delta = gstep - self._last_train_gstep + base = self.agent.cfg.num_updates + n = base * delta if self.async_options.utd_mode == "strict" else base + if self.ddp_rank is not None: + count_tensor = torch.tensor([n], device=self.agent.device) + dist.broadcast(count_tensor, src=0) + n = int(count_tensor[0]) + # Global readiness does not guarantee THIS rank's shard has data + # yet; wait (bounded by the shared stop flag) rather than skipping + # the update — skipping would desynchronize the DDP collectives. + while self.agent.rb.num_stored == 0 and not self.control.stop: + self.drain() + time.sleep(self.async_options.idle_sleep_s) + if self.control.stop and self.agent.rb.num_stored == 0: + n = 0 + elif self.agent.rb.num_stored == 0: + # Single learner: an empty buffer merely DEFERS the update — the + # cumulative delta grows, so nothing is lost on the next call. + n = 0 + if n <= 0: 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 - # counter is the single source of truth for policy-frequency gating. - metrics = self.agent.update(n) - self._last_publish_ms = 0.0 - if metrics is not None: - started = time.perf_counter() - self.publish_if_due() - self._last_publish_ms = (time.perf_counter() - started) * 1000.0 + # Sampling reads replay-buffer slots the copy stream may still be + # filling (staged async H2D ingest); make the compute stream wait for + # the in-flight copies first, mirroring the pre-staging event wait. + if self._staging is not None: + self._staging.wait() + self._last_train_gstep = gstep + # Publish cadence per gradient step: chunk strict-mode ``n`` back to + # ``num_updates``-sized updates (measured 15x more publishes than one + # huge update call — collectors act on much fresher weights). + base = self.agent.cfg.num_updates + metrics = None + remaining = n + while remaining > 0 and not self.control.stop: + metrics = self.agent.update(min(base, remaining)) + remaining -= base + self._last_publish_ms = 0.0 + if metrics is not None: + started = time.perf_counter() + self.publish_if_due() + self._last_publish_ms = (time.perf_counter() - started) * 1000.0 return metrics # ------------------------------------------------------------------ publish + def _sync_normalizer_stats(self) -> None: + """All-reduce every normalizer's (count, sum, sumsq) across DDP ranks. + + Each rank's EmpiricalNormalization only sees its own collector shard, + so per-rank statistics diverge; collectors act with rank-0's published + stats while rank-1 trains on inputs normalized differently — the + averaged gradients then pull the networks toward two different input + scalings (observed as a large multi-learner convergence gap early in + training). Merging sufficient statistics (n, Sum, SumSq) with one + SUM all-reduce per normalizer makes every rank hold identical GLOBAL + stats at each weight publish. + + Sync cadence is a deliberate design point (survey of IsaacLab / + mjlab / rsl_rl / holosoma, see MotrixLab#75): rsl_rl syncs nothing + (per-rank stats diverge forever); holosoma embeds an all_reduce in + every update() — exact per-step identity, but it requires lockstep + stepping with equal batch shapes and doubles collectives in the hot + learner loop. We merge at the weight-publish cadence instead: drift + between publishes is bounded by ``weight_publish_interval`` steps of + rank-local batches layered on fresh global stats, which captures + essentially all of the convergence benefit at 1/Nth the collectives. + If tighter consistency is ever needed, raising the sync frequency is + a call-site change — per-update merging stays structurally possible + because the local accumulators make the merge exact and idempotent + at any cadence. + """ + for norm in (self.agent.obs_normalizer, self.agent.critic_obs_normalizer): + if not isinstance(norm, EmpiricalNormalization): + continue # nn.Identity when obs_normalization is off + if not norm.local_enabled: + norm.seed_local_accumulators() + flat = norm.local_sufficient_stats_flat() + dist.all_reduce(flat, op=dist.ReduceOp.SUM) + norm.apply_global_sufficient_stats(flat) + def publish_weights(self) -> None: - self.weights.publish(self.agent.actor, self.agent.obs_normalizer) + """Broadcast the current actor snapshot to every collector's snapshot.""" + if self.agent.world_size > 1: + self._sync_normalizer_stats() + for weights in self.weights: + weights.publish(self.agent.actor, self.agent.obs_normalizer) def publish_if_due(self) -> None: if self.agent.update_idx % max(self.async_options.weight_publish_interval, 1) == 0: diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/numa.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/numa.py new file mode 100644 index 00000000..d4f8f49f --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/numa.py @@ -0,0 +1,305 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Best-effort NUMA binding for multi-collector async FastSAC training. + +Each collector on a multi-NUMA-node server should run its env, staging buffers +and pinned host allocations with node-local memory, and the scheduler should not +migrate it across nodes. WHICH node each worker gets is decided by +``async_impl/topology.py`` (learners on their GPU's PCIe-local node, collectors +on their owning learner's node); this module applies the decided placement +(the ``numactl --cpunodebind= --membind=`` equivalent) from inside the +worker process: + +* CPU affinity via ``os.sched_setaffinity`` (portable stdlib); +* memory policy via libnuma's ``set_membind`` (the same call ``numactl`` uses), + loaded with :mod:`ctypes` — future allocations of the calling process come + from the bound node. + +Everything is best-effort: when libnuma is unavailable or the node is unknown +the binding logs a warning and continues with OS default placement, so single +NUMA machines and containers keep working unchanged. + +Binding must happen at process start, before the env and any staging buffer is +allocated — ``set_membind`` only affects future allocations. +""" + +from __future__ import annotations + +import contextlib +import ctypes +import functools +import logging +import os +import re +import subprocess +from pathlib import Path +from typing import Any + +_NODE_SYSFS = Path("/sys/devices/system/node") + +logger = logging.getLogger(__name__) + +_RANGE_RE = re.compile(r"^(\d+)-(\d+)$") + + +def parse_cpulist(text: str) -> list[int]: + """Parse a sysfs ``cpulist`` (e.g. ``"0-3,8,10-11"``) into CPU ids.""" + cpus: list[int] = [] + for part in text.strip().split(","): + if not part: + continue + match = _RANGE_RE.match(part) + if match: + cpus.extend(range(int(match.group(1)), int(match.group(2)) + 1)) + else: + cpus.append(int(part)) + return cpus + + +def available_numa_nodes() -> list[int]: + """NUMA node ids visible in sysfs; empty when the host is not NUMA-aware.""" + if not _NODE_SYSFS.is_dir(): + return [] + return sorted(int(p.name.removeprefix("node")) for p in _NODE_SYSFS.iterdir() if p.name.startswith("node")) + + +def numa_node_cpus(node: int) -> list[int]: + """CPU ids of a NUMA node; raises ``ValueError`` for unknown nodes.""" + cpulist = _NODE_SYSFS / f"node{node}" / "cpulist" + if not cpulist.is_file(): + raise ValueError(f"NUMA node {node} does not exist ({cpulist} missing)") + return parse_cpulist(cpulist.read_text()) + + +def select_collector_cpus( + base: list[int], + collector_id: int, + cpus_per_collector: int | None, +) -> list[int]: + """Pick this collector's CPU slice from ``base`` (its node's or the process's CPUs). + + ``collector_id`` indexes WITHIN ``base`` — callers holding a per-node base + pass the node-local ordinal, not the global collector id. Without + ``cpus_per_collector`` all collectors share the full ``base`` set and the + OS load-balances them; with it, the set is split into contiguous chunks + so collectors never compete for the same cores. + """ + if not base: + raise ValueError("empty CPU set for collector binding") + if cpus_per_collector is None: + return list(base) + if cpus_per_collector <= 0: + raise ValueError(f"cpus_per_collector must be positive, got {cpus_per_collector}") + start = collector_id * cpus_per_collector + end = min(start + cpus_per_collector, len(base)) + if start >= end: + raise ValueError( + f"cpus_per_collector={cpus_per_collector} leaves no CPUs for collector {collector_id} " + f"(base set has {len(base)} CPUs)" + ) + return base[start:end] + + +def apply_cpu_affinity(cpus: list[int], role: str) -> None: + """Pin the current process to ``cpus``, intersected with allowed CPUs.""" + allowed = os.sched_getaffinity(0) + selected = sorted(set(cpus) & allowed) + if not selected: + raise ValueError( + f"{role}: none of the requested CPUs {sorted(cpus)} are allowed (process affinity is {sorted(allowed)})" + ) + os.sched_setaffinity(0, selected) + if set(cpus) - allowed: + logger.warning("%s: dropped non-allowed CPUs %s from binding", role, sorted(set(cpus) - allowed)) + + +def _load_libnuma(): + try: + return ctypes.CDLL("libnuma.so.1", use_errno=True) + except OSError: + return None + + +def _membind(libnuma, bitmask_ptr, role: str, what: str, node: int | None = None) -> bool: + if libnuma.numa_set_membind(ctypes.c_void_p(bitmask_ptr)) != 0: + logger.warning("%s: numa_set_membind(%s) failed; %s not bound", role, node, what) + return False + return True + + +def _node_bitmask(node: int): + libnuma = _load_libnuma() + if libnuma is None: + return None, None + libnuma.numa_allocate_nodemask.restype = ctypes.c_void_p + mask = libnuma.numa_allocate_nodemask() + if not mask: + return None, None + libnuma.numa_bitmask_setbit(ctypes.c_void_p(mask), ctypes.c_uint(node)) + return libnuma, mask + + +def _free_bitmask(libnuma, mask) -> None: + libnuma.numa_bitmask_free(ctypes.c_void_p(mask)) + + +def apply_memory_policy(node: int, role: str) -> None: + """Bind future allocations of this process to ``node`` (libnuma ``numa_set_membind``).""" + libnuma, mask = _node_bitmask(node) + if libnuma is None: + logger.warning("%s: libnuma unavailable; memory policy left to the OS (node %d not bound)", role, node) + return + try: + _membind(libnuma, mask, role, "memory policy", node) + finally: + _free_bitmask(libnuma, mask) + + +@contextlib.contextmanager +def spawn_placement(node: int | None, role: str): + """Place a spawn-started child on ``node`` by pre-placing its parent. + + A ``spawn`` child re-imports torch / the simulator / glibc arenas before + the worker entry function runs, so a bind executed inside the child only + affects allocations made after those imports — simulator thread pools and + malloc arenas created at import time land on a random node and physics + stepping pays cross-node access forever. Affinity and memory policy both + survive fork+exec, so briefly switching the *parent* onto ``node`` around + ``Process.start()`` makes the child's import-time allocations node-local; + the worker's own :func:`apply_binding` call then re-affirms the same + binding (a no-op refinement). + """ + if node is None: + yield + return + saved_affinity = os.sched_getaffinity(0) + try: + cpus = sorted(set(numa_node_cpus(node)) & saved_affinity) + if cpus: + os.sched_setaffinity(0, cpus) + libnuma, mask = _node_bitmask(node) + if libnuma is not None and mask: + try: + _membind(libnuma, mask, role, "spawn placement", node) + finally: + _free_bitmask(libnuma, mask) + yield + finally: + os.sched_setaffinity(0, saved_affinity) + libnuma = _load_libnuma() + if libnuma is not None: + # restore the default "any node" policy for the parent + try: + all_nodes = ctypes.c_void_p.in_dll(libnuma, "numa_all_nodes_ptr") + libnuma.numa_set_membind(all_nodes) + except (ValueError, OSError, AttributeError): + pass + + +def apply_binding(role: str, node: int | None, cpus: list[int]) -> None: + """Apply a topology-decided binding: CPU affinity plus node memory policy. + + ``cpus`` comes pre-computed from the topology resolution (the worker + never decides a binding itself). Both steps are best-effort; with no + CPUs and no node this is a no-op. + """ + if cpus: + apply_cpu_affinity(cpus, role) + if node is not None: + apply_memory_policy(node, role) + + +_nvml_state: tuple[Any, list[Any]] | tuple[()] | None = None # None: untried; (): unavailable + + +def _nvml() -> tuple[Any, list[Any]] | None: + """Lazily initialize NVML and return ``(pynvml, pci_bus_ids)``, or ``None``. + + Same in-process NVML session pattern as ``system_metrics``: microsecond + queries instead of an ``nvidia-smi`` subprocess spawn, and still no CUDA + context. NVML maps a device index to its PCI bus id, from which the sysfs + ``numa_node`` file gives the hosting NUMA node. + """ + global _nvml_state + if _nvml_state is None: + try: + import pynvml + + pynvml.nvmlInit() + bus_ids = [ + pynvml.nvmlDeviceGetPciInfo(pynvml.nvmlDeviceGetHandleByIndex(index)).busId + for index in range(pynvml.nvmlDeviceGetCount()) + ] + _nvml_state = (pynvml, bus_ids) + except Exception: # ImportError (pynvml missing) or NVML init failure (no driver/GPU) + _nvml_state = () + return _nvml_state or None + + +def gpu_numa_node(device_index: int) -> int | None: + """NUMA node hosting a CUDA device; None when unknown. + + Primary source: NVML (pynvml, same session pattern as ``system_metrics``) + for the index -> PCI bus id mapping and sysfs ``numa_node`` for the node + lookup — no CUDA context is created, so the pre-spawn parent can call this + freely. Fallback (containers often do not expose the PCI ``numa_node`` + sysfs attribute): the GPU's ``CPU Affinity`` column from + ``nvidia-smi topo -m``, matched against each node's ``cpulist``. A sysfs + value of -1 (unknown / single-node host) maps to None (no binding). + """ + session = _nvml() + if session is None or device_index >= len(session[1]): + return None + try: + node = int((Path("/sys/bus/pci/devices") / session[1][device_index].lower() / "numa_node").read_text().strip()) + if node >= 0: + return node + except (OSError, ValueError): + pass + return _gpu_node_from_topology(device_index) + + +@functools.lru_cache(maxsize=1) +def _gpu_cpu_affinities() -> dict[int, frozenset[int]]: + """GPU index -> CPU affinity set, parsed from ``nvidia-smi topo -m``.""" + try: + proc = subprocess.run(["nvidia-smi", "topo", "-m"], capture_output=True, text=True, timeout=10, check=True) + except (OSError, subprocess.SubprocessError): + return {} + affinities: dict[int, frozenset[int]] = {} + for line in proc.stdout.splitlines(): + cells = line.split() + if not cells or not cells[0].lstrip("\x1b[4m").startswith("GPU"): + continue + try: + index = int(cells[0].lstrip("\x1b[4m").removeprefix("GPU")) + except ValueError: + continue + # The header and data rows do not share a column layout (ANSI escapes, + # padded diagonal cells, differing column counts), so positional + # indexing is unreliable: take the token that IS a multi-cpu list. + for cell in cells[1:]: + if not re.fullmatch(r"[0-9]+(?:-[0-9]+)?(?:,[0-9]+(?:-[0-9]+)?)+|[0-9]+-[0-9]+", cell): + continue + cpus = parse_cpulist(cell) + if len(cpus) > 1: + affinities[index] = frozenset(cpus) + break + return affinities + + +def _gpu_node_from_topology(device_index: int) -> int | None: + """Node whose CPU list overlaps the GPU's affinity the most.""" + affinity = _gpu_cpu_affinities().get(device_index) + if not affinity: + return None + best_node, best_overlap = None, 0 + for node in available_numa_nodes(): + try: + overlap = len(set(numa_node_cpus(node)) & affinity) + except ValueError: + continue + if overlap > best_overlap: + best_node, best_overlap = node, overlap + return best_node diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/panels.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/panels.py new file mode 100644 index 00000000..fa6597b3 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/panels.py @@ -0,0 +1,186 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Rich panel rendering for the async FastSAC trainer's parent process. + +The boot panel (worker startup table + worker-log tail) runs on the +alternate screen via ``rich.live.Live`` while workers boot at different +speeds. Everything here is pure rendering; the boot loop, readiness +tracking and error handling stay in ``train.py``. The post-boot training +panel is shared across frameworks and lives in ``motrix_rl.console``. + +Frame geometry is the core contract: the panel is a vertical ``Layout`` +where the worker table takes a fixed number of lines and the log region +expands to every remaining terminal line. Log lines are either cropped or +folded into full-width continuations, and the DISPLAY line count is capped +so the frame height never changes between refreshes — a Live redraw whose +frame grows, wraps or overflows the terminal cannot erase its previous +frame and tears. +""" + +from __future__ import annotations + +import math +import os +from collections.abc import Sequence +from pathlib import Path + +from rich.console import Console, Group +from rich.layout import Layout +from rich.panel import Panel +from rich.table import Table +from rich.text import Text + +_LOG_TAIL_PLACEHOLDER = "(waiting for worker logs…)" + + +def worker_log_tail( + log_dir: Path, + log_names: Sequence[str], + max_lines: int = 4, + line_width: int = 120, + per_file: int = 2, + wrap: bool = False, +) -> Text: + """Render the boot panel's worker-log region. + + Reads the last bytes of each worker log file and returns at most + ``max_lines`` display lines (the last ``per_file`` source lines per file, + newest last). Each source line becomes exactly one display line cropped + to ``line_width``, or — with ``wrap=True`` — as many full-width + continuation lines as it needs. Either way the panel's geometry stays + constant across refreshes (the display-line count is always capped at + ``max_lines``), which is what keeps the Live redraw tear-free. + + Lines are grouped per worker in ``log_names`` order (collectors first, + learners last), untagged — the surrounding panel/table names the worker. + Missing files (worker hasn't created its log yet) are skipped silently. + """ + display: list[str] = [] + for name in log_names: + path = Path(log_dir) / name + try: + with path.open("rb") as f: + f.seek(0, os.SEEK_END) + size = f.tell() + f.seek(max(0, size - 8192)) + chunk = f.read().decode(errors="replace") + except OSError: + continue + tail = [ln for ln in chunk.splitlines() if ln.strip()][-per_file:] + for ln in tail: + if wrap: + display.extend(ln[j : j + line_width] for j in range(0, len(ln), line_width)) + else: + display.append(ln[:line_width]) + return Text("\n".join(display[-max_lines:]) or _LOG_TAIL_PLACEHOLDER, style="dim") + + +class BootPanel: + """Worker-startup panel: worker table below a grid of per-worker log cells.""" + + def __init__(self, title: str, log_dir: Path, log_names: Sequence[str], workers: Sequence[tuple[str, int]]) -> None: + self._title = title + self._log_dir = Path(log_dir) + self._log_names = list(log_names) + self._workers = list(workers) + self.console = Console() + # Outer frame mirrors the training panel: task name on top, cyan + # border. Its border + padding cost 2 lines and 4 columns of the + # terminal, shrinking the inner regions accordingly. + self._frame_lines = 2 + self._frame_cols = 4 + # Table height = header + 3 rule rows + one row per worker, + # plus the always-reserved gate line. + # The log grid VERTICALLY EXPANDS to every remaining terminal line. + self._table_lines = len(self._workers) + 4 + 1 + # Near-square grid: 1-2 workers stay in one row, more pack into + # ceil(sqrt(n)) columns so cells keep a readable width. + n = max(1, len(self._log_names)) + self._grid_cols = math.ceil(math.sqrt(n)) + self._grid_rows = math.ceil(n / self._grid_cols) + # Cell inner size: the log region split by the grid shape, minus each + # cell Panel's border and padding columns/rows. + region_lines = max(self._grid_rows * 3, self.console.height - self._frame_lines - self._table_lines) + self.log_tail_lines = max(1, region_lines // self._grid_rows - 2) + inner_width = self.console.width - self._frame_cols + self.log_line_width = max(20, inner_width // self._grid_cols - 4) + # Without a TTY (piped/redirected stdout) Live cannot redraw in place + # and every refresh appends a new frame; the caller renders one + # static frame instead (auto_refresh=False, no updates). + self.interactive = self.console.is_terminal + + def _worker_panel(self, name: str) -> Panel: + """One worker's log cell: fixed-size Panel padded to full height.""" + worker = name.removesuffix(".log") + text = worker_log_tail( + self._log_dir, + [name], + self.log_tail_lines, + self.log_line_width, + per_file=self.log_tail_lines, + wrap=True, + ) + # Pad to a fixed line count so the panel (and thus the whole frame) + # keeps a constant height between refreshes. + body = text.plain.splitlines() + body += [""] * (self.log_tail_lines - len(body)) + return Panel( + Text("\n".join(body), style="dim"), + title=worker, + border_style="dim", + expand=True, + height=self.log_tail_lines + 2, + ) + + def render(self, ready: set[tuple[str, int]], gate: str | None = None, starting: bool = False) -> Panel: + """Render one boot/handoff frame. + + The whole view is wrapped in the same outer frame the training panel + uses (task name on top, cyan border). ``ready`` drives the boot phase + (booting…/ready). After the barrier releases, the caller sets + ``starting`` (every worker flips to a "starting" status) and passes + ``gate``, the one-line panel-data readiness summary; the same panel + then serves as the handoff view until the training panel's quiescence + gate opens, so the terminal never switches views mid-boot. + """ + inner_width = self.console.width - self._frame_cols + table = Table(expand=True, width=inner_width) + table.add_column("worker") + table.add_column("status") + for role, idx in sorted(self._workers): + name = f"{role}[{idx}]" + if starting: + table.add_row(name, "[yellow]starting[/yellow]") + elif (role, idx) in ready: + table.add_row(name, "[green]ready[/green]") + else: + table.add_row(name, "[dim]booting…[/dim]") + # Worker-log grid: rows of equal-height cells, each cell one worker. + logs = Layout(name="logs") + rows = [] + for r in range(self._grid_rows): + cells = self._log_names[r * self._grid_cols : (r + 1) * self._grid_cols] + row = Layout(name=f"log-row{r}") + row.split_row(*(Layout(name=name.removesuffix(".log"), ratio=1) for name in cells)) + rows.append(row) + logs.split_column(*rows) + for name in self._log_names: + logs[name.removesuffix(".log")].update(self._worker_panel(name)) + # The gate line is always reserved (empty during the boot phase) so + # the frame geometry is identical across the boot→handoff switch. + bottom = Group(table, Text(gate if gate else "", style="dim")) + layout = Layout() + layout.split_column( + Layout(name="logs"), + Layout(name="table", size=self._table_lines), + ) + layout["logs"].update(logs) + layout["table"].update(bottom) + return Panel( + layout, + title=self._title, + border_style="cyan", + padding=(0, 1), + height=self.console.height, + ) diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/stats.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/stats.py new file mode 100644 index 00000000..b7ecb607 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/stats.py @@ -0,0 +1,73 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Pure rollout/timing statistics aggregation for the async trainer. + +Parent-process-side helpers: collectors and learners emit per-process +snapshots; these functions merge them into the panel/TensorBoard-facing +views. No process, transport, or simulator knowledge — everything here is +a pure function over dicts and floats. +""" + +from __future__ import annotations + +import math +from typing import Any + + +def timing_mean(values: list[float]) -> float: + """Mean of a non-empty timing sample list (ms).""" + return sum(values) / len(values) + + +def nest_timing_path(tree: dict[str, Any], parts: tuple[str, ...], value: float) -> None: + """Insert one dotted timing path into a nested mapping. + + A stage's scalar total and its sub-stage paths may arrive in either order + (the collector emits parents before children); when both exist the scalar + becomes the node's ``total`` alongside its children. + """ + head, rest = parts[0], parts[1:] + node = tree.get(head) + if not rest: + if isinstance(node, dict): + node["total"] = value + else: + tree[head] = value + else: + if not isinstance(node, dict): + node = {"total": node} if node is not None else {} + tree[head] = node + nest_timing_path(node, rest, value) + + +def aggregate_collector_stats(per_collector: dict[int, dict | None]) -> dict: + """Merge per-collector rollout snapshots into one panel/TB-facing stats dict. + + Return/episode-length/reward/metric values are averaged over the collectors + that reported since the last log window; episode counts are summed; + ``policy_lag`` is the worst (max) staleness; timing values are averaged. + """ + snapshots = [stats for stats in per_collector.values() if stats] + + def _nanmean(key: str) -> float: + values = [stats[key] for stats in snapshots if not math.isnan(stats.get(key, float("nan")))] + return (sum(values) / len(values)) if values else float("nan") + + def _mean_dicts(key: str) -> dict[str, float]: + keys = {k for stats in snapshots for k in stats.get(key, {})} + return { + k: sum(stats[key][k] for stats in snapshots if k in stats.get(key, {})) + / sum(1 for stats in snapshots if k in stats.get(key, {})) + for k in keys + } + + return { + "return": _nanmean("return"), + "ep_len": _nanmean("ep_len"), + "episodes": sum(stats.get("episodes", 0) for stats in snapshots), + "reward_terms": _mean_dicts("reward_terms"), + "env_metrics": _mean_dicts("env_metrics"), + "policy_lag": max((stats.get("policy_lag", 0) for stats in snapshots), default=0), + "timing_ms": _mean_dicts("timing_ms"), + } diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py new file mode 100644 index 00000000..06576387 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py @@ -0,0 +1,372 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Worker topology decisions for the async FastSAC trainer. + +Pure decision functions over ``(opts, devices, counts)`` — how envs shard +across collectors, which NUMA node each worker is placed on, and which +transport each ring uses. The policy is co-location, informed by two +measurements on dual-socket hosts: + +1. a collector must never be placed on a different NUMA node than the learner + that drains its ring — the transition ring, pinned staging and weight + snapshots would all go cross-node; +2. a learner binds to its GPU's PCIe-local node, and each collector follows + its OWNING learner (collector ``i`` belongs to learner + ``i // (num_collectors // num_learners)``), so the pair shares one node's + memory domain. Restricting a single collector to one node does NOT hurt: + first-touch page locality inside the physics working set outweighs the + halved aggregate bandwidth (measured +58% on a 2-socket/2-GPU host). + +Everything degrades to "no binding" (all ``None``) on single-node hosts, +non-NUMA kernels/containers and CPU learners — the OS default placement is +already optimal there. + +The binding itself (affinity + libnuma memory policy, spawn-time +pre-placement) lives in :mod:`motrix_rl.fastsac.async_impl.numa`, the ring +runtime objects in :mod:`motrix_rl.fastsac.async_impl.transport`; this module +only computes the decisions. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass + +import torch + +from motrix_rl.fastsac.async_impl import numa +from motrix_rl.fastsac.config import FastSacAsyncOptionsCfg + + +@dataclass(frozen=True) +class LearnerInfo: + """Everything one learner rank needs to know about itself.""" + + rank: int + device: torch.device + numa_node: int | None + cpus: list[int] + + +@dataclass(frozen=True) +class CollectorInfo: + """Everything one collector process needs to know about itself.""" + + collector_id: int + owner_rank: int + num_envs: int + device: torch.device + numa_node: int | None + cpus: list[int] + ring_ipc: bool + weight_ipc: bool + + +@dataclass(frozen=True) +class TrainerTopology: + """The complete compute layout of one async trainer run. + + One structure describing how the distributed training network is wired: + per-worker self descriptions (devices, NUMA placement, CPU bindings, + transports, env shards) plus the collector→learner ownership relation. + ``resolve_trainer_topology`` is the single computation API producing it. + + Invariants (by construction): + + * ``collectors[i].numa_node == learners[collectors[i].owner_rank].numa_node`` + — a collector/learner pair never straddles a NUMA boundary; + * ``collectors[i].num_envs`` envs travel with collector ``i``; + * ``sum(c.num_envs for c in collectors) == num_envs``. + """ + + learners: list[LearnerInfo] + collectors: list[CollectorInfo] + + @property + def num_learners(self) -> int: + return len(self.learners) + + @property + def num_collectors(self) -> int: + return len(self.collectors) + + @property + def per_learner(self) -> int: + return self.num_collectors // self.num_learners + + @property + def env_shards(self) -> list[int]: + return [collector.num_envs for collector in self.collectors] + + def collector_owner(self, collector_id: int) -> int: + """The learner rank that owns (drains the rings of) this collector.""" + return self.collectors[collector_id].owner_rank + + def ring_slice_for_rank(self, rings: list, rank: int) -> list: + """The rings this learner rank drains: contiguous collector block ``rank``. + + Each ring still has exactly one consumer, so the SPSC contract is + unchanged. + """ + return rings[rank * self.per_learner : (rank + 1) * self.per_learner] + + +def _learner_node(device: torch.device, multi_node: bool) -> int | None: + """GPU-local node of one learner rank; ``None`` for CPU / unknown hosts.""" + if not multi_node: + return None + if device.type != "cuda": + return None + index = device.index if device.index is not None else 0 + return numa.gpu_numa_node(index) + + +def resolve_learner_devices( + learner_device_specs: list[str] | None, + num_learners: int, + default_device: torch.device, +) -> list[torch.device]: + """One indexed CUDA device per learner rank. + + ``None`` replicates the trainer device: for ``num_learners > 1`` an + unindexed ``cuda`` expands to consecutive indexes starting at the trainer + device's; anything else (an indexed single device or CPU) cannot back + multiple ranks and is rejected. A single learner with ``None`` resolves + in-process (empty list — the worker keeps its own resolution). + """ + if num_learners == 1 and learner_device_specs is None: + return [] # single-learner path resolves its device in-process + if learner_device_specs is None: + base = default_device + if base.type != "cuda": + raise ValueError(f"num_learners={num_learners} requires CUDA learner devices, got {base}") + start = base.index if base.index is not None else 0 + specs = [f"cuda:{start + i}" for i in range(num_learners)] + else: + if len(learner_device_specs) != num_learners: + raise ValueError( + f"learner_devices must have exactly num_learners={num_learners} entries, " + f"got {len(learner_device_specs)}" + ) + specs = learner_device_specs + devices = [torch.device(spec) for spec in specs] + indexes = [d.index for d in devices] + if any(d.type != "cuda" or d.index is None for d in devices): + raise ValueError(f"learner_devices must be explicit indexed CUDA devices, got {specs}") + if len(set(indexes)) != len(indexes): + raise ValueError(f"learner_devices reference a device more than once: {specs}") + available = torch.cuda.device_count() + for d in devices: + if d.index >= available: + raise ValueError(f"learner device {d} does not exist (only {available} CUDA device(s) available)") + return devices + + +def split_num_envs(num_envs: int, num_collectors: int) -> list[int]: + """Shard ``num_envs`` evenly across collectors (requires exact divisibility). + + Topological complement of the collector→learner ownership encoded in + :class:`TrainerTopology`: shard ``i`` (and so its envs) always travels + with collector ``i``, which never straddles a NUMA boundary from its + owning learner. + """ + if num_collectors < 1: + raise ValueError(f"num_collectors must be >= 1, got {num_collectors}") + if num_envs % num_collectors != 0: + raise ValueError( + f"num_envs={num_envs} must divide evenly across num_collectors={num_collectors} " + f"({num_envs} % {num_collectors} != 0)" + ) + per_collector = num_envs // num_collectors + return [per_collector] * num_collectors + + +def use_ipc_weight_channel( + opts: FastSacAsyncOptionsCfg, + learner_device: torch.device, + collector_device: torch.device, + actor_param_numel: int, +) -> bool: + """Whether the weight channel for one collector should use CUDA-IPC. + + Like :func:`use_ipc_transition_ring`, plus a size threshold: IPC device + slots only pay off when the actor parameters reach the configured + ``weight_ipc_min_bytes``; host shared memory otherwise. + """ + if not use_ipc_transition_ring(opts, learner_device, collector_device): + return False + return actor_param_numel * 4 >= opts.weight_ipc_min_bytes + + +def resolve_trainer_topology( + num_envs: int, + num_collectors: int, + num_learners: int, + learner_devices: list[torch.device], + collector_devices: list[torch.device], + default_device: torch.device, + async_options: FastSacAsyncOptionsCfg, + actor_param_numel: int, + cpus_per_collector: int | None = None, +) -> TrainerTopology: + """Single computation API: derive the full trainer topology in one pass. + + Combines env sharding (:func:`split_num_envs`), NUMA placement, CPU + bindings, and the transport decisions for transition rings + (:func:`ring_transport_is_ipc`) and weight channels + (:func:`use_ipc_weight_channel`) into one :class:`TrainerTopology`. + ``actor_param_numel`` is the parent-computed actor parameter count that + sizes the weight-transport threshold; ``cpus_per_collector`` optionally + chunks each binding base into per-collector slices. ``learner_devices`` + carries one indexed device per rank (empty for a single learner, whose + device comes from ``default_device`` — an index-less ``cuda`` means + device 0). ``collector_devices[i]`` is the inference device of collector + ``i``. + """ + if num_learners < 1 or num_collectors < 1: + raise ValueError(f"invalid worker counts: {num_collectors=} {num_learners=}") + if num_collectors % num_learners != 0: + raise ValueError(f"num_collectors={num_collectors} must divide evenly across num_learners={num_learners}") + if len(collector_devices) != num_collectors: + raise ValueError(f"expected {num_collectors} collector devices, got {len(collector_devices)}") + + env_shards = split_num_envs(num_envs, num_collectors) + multi_node = len(numa.available_numa_nodes()) >= 2 + per_learner = num_collectors // num_learners + + # Per-worker CPU bindings, decided here so workers only apply them. Base + # is the NUMA node's CPUs when the worker is node-bound, else the + # process's own affinity mask (children inherit it through spawn). An + # unreadable node degrades to no binding, matching the numa module's + # best-effort contract. + def _base_cpus(node: int | None) -> list[int]: + if node is None: + return sorted(os.sched_getaffinity(0)) + try: + return numa.numa_node_cpus(node) + except ValueError: + return [] + + def _chunk(base: list[int], local_id: int) -> list[int]: + if not base: + return [] + return numa.select_collector_cpus(base, local_id, cpus_per_collector) + + # resolve_learner_devices' contract: learner_devices is empty iff + # num_learners == 1, and the single learner runs on default_device. + def _rank_device(rank: int) -> torch.device: + return learner_devices[rank] if learner_devices else default_device + + ring_ipc = ring_transport_is_ipc( + async_options, learner_devices, collector_devices, num_collectors, num_learners, default_device + ) + + learners: list[LearnerInfo] = [] + for rank in range(num_learners): + device = _rank_device(rank) + node = _learner_node(device, multi_node) + learners.append(LearnerInfo(rank=rank, device=device, numa_node=node, cpus=_base_cpus(node))) + + collectors: list[CollectorInfo] = [] + for i, (num_envs_i, collector_dev) in enumerate(zip(env_shards, collector_devices)): + owner = i // per_learner + owner_node = learners[owner].numa_node + collectors.append( + CollectorInfo( + collector_id=i, + owner_rank=owner, + num_envs=num_envs_i, + device=collector_dev, + numa_node=owner_node, + # chunk by the NODE-LOCAL ordinal: the base list is the + # owner node's own CPUs, so a global index would offset + # rank>=1 collectors past their node's list (asymmetric + # bindings, or "leaves no CPUs" on small nodes). + cpus=_chunk(_base_cpus(owner_node), i % per_learner), + ring_ipc=ring_ipc[i], + weight_ipc=use_ipc_weight_channel(async_options, _rank_device(owner), collector_dev, actor_param_numel), + ) + ) + return TrainerTopology(learners=learners, collectors=collectors) + + +def same_cuda_device(learner_device: torch.device, collector_device: torch.device) -> bool: + """Whether the learner and the collector's inference device share one GPU. + + An index-less ``cuda`` means the default current device (index 0 — nothing + in the trainer ever calls ``torch.cuda.set_device``), so it is resolved + with 0 rather than treated as a wildcard matching any explicit index: + ``learner=cuda`` (effectively cuda:0) with ``collector_inference_device: + cuda:1`` is a cross-GPU setup and must NOT enable the device transports. + Pure device arithmetic — no CUDA context is created, so the pre-spawn + parent can call it as safely as the workers. + """ + if learner_device.type != "cuda" or collector_device.type != "cuda": + return False + learner_index = learner_device.index if learner_device.index is not None else 0 + collector_index = collector_device.index if collector_device.index is not None else 0 + return learner_index == collector_index + + +def use_ipc_transition_ring( + opts: FastSacAsyncOptionsCfg, learner_device: torch.device, collector_device: torch.device +) -> bool: + """Whether the transition ring should use CUDA-IPC device slots. + + Requires learner and collector inference on the same GPU (see + :func:`same_cuda_device`); otherwise the host shared-memory ring is used. + Purely device-object arithmetic — no CUDA context is created here, so the + parent can call it safely. + """ + mode = opts.transition_ipc + # YAML 1.1 parses unquoted ``on``/``off`` scalars as booleans; accept that + # form so ``transition_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.transition_ipc must be auto, on or off, got {mode!r}") + same_gpu = same_cuda_device(learner_device, collector_device) + if mode == "off": + return False + if mode == "on" and not same_gpu: + reason = ( + "collector inference device is not CUDA" + if collector_device.type != "cuda" + else f"learner device {learner_device} and collector device {collector_device} are different GPUs" + if learner_device.type == "cuda" + else "learner device is not CUDA" + ) + logging.getLogger(__name__).warning( + "async_options.transition_ipc=on requires learner and collector inference on the same GPU, " + "but %s; falling back to the host shared-memory transition ring", + reason, + ) + return same_gpu + + +def ring_transport_is_ipc( + opts: FastSacAsyncOptionsCfg, + learner_devices: list[torch.device], + collector_devices: list[torch.device], + num_collectors: int, + num_learners: int, + default_owner: torch.device, +) -> list[bool]: + """Per-collector CUDA-IPC ring decision, all-or-nothing per learner rank. + + A rank's rings share one transport so its drain path stays homogeneous + (the generation merge would choke on mixed host/device views). Collector + i belongs to rank i // (num_collectors // num_learners) and its ring goes + IPC only when EVERY collector of that rank infers on the rank's GPU. + """ + per_learner = num_collectors // num_learners + decisions: list[bool] = [] + for rank in range(num_learners): + owner = learner_devices[rank] if num_learners > 1 else default_owner + rank_devs = collector_devices[rank * per_learner : (rank + 1) * per_learner] + rank_ipc = all(use_ipc_transition_ring(opts, owner, dev) for dev in rank_devs) + decisions.extend([rank_ipc] * per_learner) + return decisions 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 59e8989a..408550f8 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py @@ -16,8 +16,10 @@ from __future__ import annotations +import os import random import time +from pathlib import Path from queue import Empty import numpy as np @@ -26,19 +28,33 @@ from motrix_env_core import registry as env_registry from motrix_env_core.renderer import RenderConfig +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 resolve_collector_inference_device +from motrix_rl.fastsac.async_impl.numa import spawn_placement +from motrix_rl.fastsac.async_impl.panels import BootPanel +from motrix_rl.fastsac.async_impl.stats import aggregate_collector_stats, nest_timing_path, timing_mean +from motrix_rl.fastsac.async_impl.topology import resolve_learner_devices, resolve_trainer_topology from motrix_rl.fastsac.async_impl.transport import Control, RingCursors, SharedTransitionRing +from motrix_rl.fastsac.async_impl.transport.handshake import StartupHandshake from motrix_rl.fastsac.async_impl.transport.weight_channel import WeightChannelShared from motrix_rl.fastsac.async_impl.worker import ( + actor_param_numel, build_env, + inherit_log_stdio, run_collector_process, run_learner_process, - use_ipc_transition_ring, ) from motrix_rl.fastsac.config import FastSacCfg from motrix_rl.fastsac.wrap import FastSacEnvWrap from motrix_rl.frameworks import TrainerBase, TrainerContext +from motrix_rl.system_metrics import ( + CpuLoadSampler, + GpuMemoryUsageSampler, + GpuUtilizationSampler, + MemoryUsageSampler, + sample_gpu_devices, +) torch.set_float32_matmul_precision("high") @@ -131,18 +147,81 @@ def train(self) -> None: learner_device = self._device() collector_device = resolve_collector_inference_device(async_options.collector_inference_device) - # shared-memory primitives allocated in the parent, inherited by children. - # The transition ring's transport is decided here: host shared-memory - # fields, or (same-GPU collector/learner) bare cursors — the learner - # then allocates the CUDA-IPC device slots and ships them to the - # collector through the one-shot slot_queue handshake. + num_collectors = async_options.num_collectors + num_learners = async_options.num_learners + if num_learners < 1: + raise ValueError(f"num_learners must be >= 1, got {num_learners}") + if num_learners > 1 and num_collectors % num_learners != 0: + raise ValueError(f"num_collectors={num_collectors} must divide evenly across num_learners={num_learners}") + if cfg.agent.batch_size % num_learners != 0: + raise ValueError( + f"agent.batch_size={cfg.agent.batch_size} must divide evenly across num_learners={num_learners}" + ) + # NUMA placement (see async_impl/topology.py): each learner binds to + # its GPU's PCIe-local node and every collector follows its owning + # learner, so a collector/learner pair never straddles a NUMA node. + # Falls back to the OS default placement on single-node hosts. + learner_devices = resolve_learner_devices(async_options.learner_devices, num_learners, self._device()) + cpus_per_collector = async_options.cpus_per_collector num_envs = self._context.num_envs - if use_ipc_transition_ring(async_options, learner_device, collector_device): - ring: SharedTransitionRing | RingCursors = RingCursors() + # Generic "cuda" collector inference resolves to the owning learner's + # GPU (num_collectors // num_learners collectors per learner); explicit + # specs pass through. The single-learner path keeps the in-process + # resolution (collector_device=None) byte-identical. + collector_device_specs = None + if num_learners > 1: + per_learner = num_collectors // num_learners + spec = async_options.collector_inference_device + collector_device_specs = [ + spec + if not (spec == "cuda" and learner_devices[i // per_learner].type == "cuda") + else f"cuda:{learner_devices[i // per_learner].index}" + for i in range(num_collectors) + ] + # Shared-memory primitives allocated in the parent, inherited by children. + # One SPSC ring + one weight channel per collector: every shared quantity + # keeps exactly one producer and one consumer, so the lock-free + # single-writer invariants are unchanged by the collector count. + # + # The ring transport is decided per LEARNER RANK (all-or-nothing over + # its collectors, so the rank's drain stays transport-homogeneous): + # when every collector of the rank infers on the rank's GPU, its rings + # are bare cursors and each learner allocates the CUDA-IPC device + # slots, shipping them to its collectors through the one-shot + # ring-slot queue; otherwise the rank's rings are host shared memory. + # Collectors are co-located with their owning learner's GPU by + # default (collector_inference_device="cuda" resolves per owner), so + # the IPC path is the default whenever both sides share that GPU. + if collector_device_specs is not None: + collector_devices = [torch.device(spec) for spec in collector_device_specs] else: - ring = SharedTransitionRing(async_options.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim) - weights = WeightChannelShared(obs_dim=obs_dim) - control = Control() + collector_devices = [collector_device] * num_collectors + # One resolution pass derives the whole compute layout: env shards, + # NUMA placement, and per-collector transports (rings + weights). + param_numel = actor_param_numel(cfg, dims, action_scale, action_bias) + topology = resolve_trainer_topology( + num_envs, + num_collectors, + num_learners, + learner_devices, + collector_devices, + self._device(), + async_options, + param_numel, + cpus_per_collector=cpus_per_collector, + ) + numa_nodes = [collector.numa_node for collector in topology.collectors] + learner_numa_nodes = [learner.numa_node for learner in topology.learners] + env_shards = topology.env_shards + ring_ipc = [collector.ring_ipc for collector in topology.collectors] + rings: list[SharedTransitionRing | RingCursors] = [ + RingCursors() + if ipc + else SharedTransitionRing(async_options.ring_capacity, shard, obs_dim, critic_obs_dim, act_dim) + for ipc, shard in zip(ring_ipc, env_shards) + ] + weights = [WeightChannelShared(obs_dim=obs_dim) for _ in range(num_collectors)] + control = Control(num_collectors) resume_step = 0 is_resume = False @@ -152,24 +231,31 @@ def train(self) -> None: ckpt = torch.load(self._resume_from, map_location="cpu", weights_only=False) resume_step = int(ckpt.get("global_step", 0)) is_resume = resume_step > 0 - control.collector_steps = resume_step + control.resume_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) + stats_queues = [ctx.Queue(maxsize=8) for _ in range(num_collectors)] 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) + # Per-window worker payloads for the parent-rendered panel: rollout + # snapshots from every collector (stats_queues below) plus one learner + # payload per rank. All workers are peers; the parent is the recorder. + panel_queue = ctx.Queue(maxsize=8 * max(num_collectors + num_learners, 1)) + # Startup handshake (parent-arbitrated, see transport/handshake.py): + # learners ship slot tensors through the per-collector one-shot + # queues; workers report readiness, the parent renders a boot panel + # and releases everyone at once. + handshake = StartupHandshake(ctx, num_collectors) reported_errors: set[tuple[str, str]] = set() seed = self._context.seed + def _record_error_trace(process_name: str, traceback_text: str) -> Path: + error_dir = self._context.run_dir / "async_errors" + error_dir.mkdir(parents=True, exist_ok=True) + error_path = error_dir / f"{process_name}_error.log" + error_path.write_text(traceback_text) + return error_path + def _drain_child_errors() -> list[tuple[str, str]]: errors = [] while True: @@ -183,107 +269,527 @@ def _drain_child_errors() -> list[tuple[str, str]]: reported_errors.add(key) errors.append(key) - error_dir = self._context.run_dir / "async_errors" - error_dir.mkdir(parents=True, exist_ok=True) - error_path = error_dir / f"{process_name}_error.log" - error_path.write_text(traceback_text) + error_path = _record_error_trace(process_name, traceback_text) print(f"[motrix.fastsac async] {process_name} traceback written to {error_path}") print(traceback_text.rstrip()) return errors print( - f"[motrix.fastsac async] two-process training '{self._env_name}' learner={learner_device} " - f"collector_env=cpu collector_inference={collector_device} num_envs={num_envs} iters={num_iterations} " + f"[motrix.fastsac async] collector/learner training '{self._env_name}' learner={learner_device} " + f"learner_replicas={num_learners} " + + (f"learner_devices={[str(d) for d in learner_devices]} " if num_learners > 1 else "") + + f"collector_env=cpu collector_inference={collector_device} num_collectors={num_collectors} " + f"numa_nodes={numa_nodes} learner_numa_nodes={learner_numa_nodes} " + f"num_envs={num_envs} iters={num_iterations} " f"from={resume_step} utd_mode={async_options.utd_mode}" ) - p_learner = ctx.Process( - target=run_learner_process, - args=( - cfg, - num_envs, - dims, - action_scale, - action_bias, - ring, - weights, - control, - stats_queue, - error_queue, - num_iterations, - logging_interval, - save_interval, - str(self._context.run_dir), - self._env_name, - str(self._context.checkpoint_dir), - self._context.checkpoint_format, - self._resume_from, - seed, - slot_queue, - ), - name="fastsac-async-learner", + rendezvous_file = None + if num_learners > 1: + # Resolve to an absolute path: init_method="file://" is + # URL-parsed (host/path), which mangles a relative run_dir into a + # bogus absolute path and blocks forever inside the file store. + rendezvous_path = (self._context.run_dir / "ddp_rendezvous").resolve() + rendezvous_path.parent.mkdir(parents=True, exist_ok=True) + rendezvous_path.unlink(missing_ok=True) # the file store requires a fresh file + rendezvous_file = str(rendezvous_path) + per_learner = num_collectors // num_learners + p_learners = [ + ctx.Process( + target=run_learner_process, + args=( + cfg, + num_envs, + dims, + action_scale, + action_bias, + topology.ring_slice_for_rank(rings, rank), + # every rank publishes to its OWN collectors; rank 0 is + # additionally the logger / checkpointer. + weights[rank * per_learner : (rank + 1) * per_learner], + control, + error_queue, + num_iterations, + logging_interval, + save_interval, + str(self._context.run_dir), + str(self._context.checkpoint_dir), + self._context.checkpoint_format, + self._resume_from, + seed, + handshake, + ), + kwargs={ + "rank": rank, + "panel_queue": panel_queue, + "num_learners": num_learners, + "rendezvous_file": rendezvous_file, + "learner_device": str(learner_devices[rank]) if learner_devices else None, + "all_rings": rings, + "weight_ipc": [ + c.weight_ipc for c in topology.collectors[rank * per_learner : (rank + 1) * per_learner] + ], + "learner_cpus": topology.learners[rank].cpus if learner_numa_nodes else None, + "learner_numa_node": learner_numa_nodes[rank] if learner_numa_nodes else None, + }, + name="fastsac-async-learner" if num_learners == 1 else f"fastsac-async-learner-{rank}", + ) + for rank in range(num_learners) + ] + p_collectors = [ + ctx.Process( + target=run_collector_process, + kwargs={ + "env_spec": self._env_spec, + "cfg": cfg, + "num_envs": env_shards[i], + "dims": dims, + "action_scale": action_scale, + "action_bias": action_bias, + "ring": rings[i], + "weights": weights[i], + "control": control, + "stats_queue": stats_queues[i], + "error_queue": error_queue, + "num_iterations": num_iterations, + "logging_interval": logging_interval, + "is_resume": is_resume, + "seed": None if seed is None else seed + i, + "collector_id": i, + "numa_node": numa_nodes[i], + "cpus": topology.collectors[i].cpus, + "collector_device": collector_device_specs[i] if collector_device_specs is not None else None, + "run_dir": str(self._context.run_dir), + "handshake": handshake, + }, + name=f"fastsac-async-collector-{i}", + ) + for i in range(num_collectors) + ] + + # Start each child while the parent is pre-placed on the child's NUMA + # node: spawn children re-import torch/simulator before their entry + # function runs, and affinity + memory policy survive fork+exec, so + # the child's import-time allocations are node-local from the start + # (the worker re-binds itself afterwards as a no-op refinement). + log_dir = self._context.run_dir / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + learner_log_name = (lambda rank: f"learner{rank}.log") if num_learners > 1 else (lambda rank: "learner.log") + for rank, p in enumerate(p_learners): + with ( + spawn_placement(learner_numa_nodes[rank] if learner_numa_nodes else None, f"learner-spawn[{rank}]"), + inherit_log_stdio(log_dir / learner_log_name(rank)), + ): + p.start() + for i, p in enumerate(p_collectors): + with ( + spawn_placement(numa_nodes[i], f"collector-spawn[{i}]"), + inherit_log_stdio(log_dir / f"collector{i}.log"), + ): + p.start() + + # -------------------------------------------------- startup panel + barrier + # Workers boot at different speeds (env/numba compile vs CUDA init and + # the DDP rendezvous). The parent tracks readiness on a live table and + # releases everyone at once, so stepping starts in lockstep. A worker + # crash or the boot timeout aborts via the shared stop flag. + expected = {("collector", i) for i in range(num_collectors)} | {("learner", r) for r in range(num_learners)} + ready: set[tuple[str, int]] = set() + boot_deadline = time.time() + max(float(os.environ.get("MOTRIX_BOOT_TIMEOUT_S", "600")), 30.0) + from rich.live import Live + + log_names = [f"collector{i}.log" for i in range(num_collectors)] + log_names += [learner_log_name(rank) for rank in range(num_learners)] + boot_panel = BootPanel( + title=f"{self._env_name}/motrix.fastsac", + log_dir=log_dir, + log_names=log_names, + workers=sorted(expected), ) - p_collector = ctx.Process( - target=run_collector_process, - args=( - self._env_spec, - cfg, - num_envs, - dims, - action_scale, - action_bias, - ring, - weights, - control, - stats_queue, - error_queue, - num_iterations, - logging_interval, - is_resume, - seed, - slot_queue, - ), - name="fastsac-async-collector", + interactive = boot_panel.interactive + + # Panel-data state must exist BEFORE the boot barrier: the handoff + # loop below drains the same queues the training panel will use, so + # the transition between panels is a data-source switch, not a wait. + per_collector_stats: dict[int, dict | None] = {i: None for i in range(num_collectors)} + learner_payloads: dict[int, dict | None] = {r: None for r in range(num_learners)} + + def _drain_worker_stats() -> None: + for queue in stats_queues: + try: + while True: + snapshot = queue.get_nowait() + per_collector_stats[snapshot["collector_id"]] = snapshot + except Empty: + pass + try: + while True: + payload = panel_queue.get_nowait() + learner_payloads[payload["rank"]] = payload + except Empty: + pass + + def _panel_data_ready() -> bool: + """The training panel's quiescence gate (see _render_panel).""" + return ( + all(v is not None for v in per_collector_stats.values()) + and any(v is not None for v in learner_payloads.values()) + and control.collector_steps > 0 + ) + + def _handoff_gate_line() -> str: + """One-line panel-data readiness summary for the handoff view.""" + stats_marks = " ".join( + f"c{i} {'✓' if per_collector_stats[i] is not None else '…'}" for i in range(num_collectors) + ) + payload_marks = " ".join( + f"l{r} {'✓' if learner_payloads[r] is not None else '…'}" for r in range(num_learners) + ) + return f"panel data — stats: {stats_marks} | payloads: {payload_marks} | steps: {control.collector_steps}" + + try: + # Bare prints inside Live would tear its frame; collect abort + # reasons here and print them after Live has exited. + abort_reason: str | None = None + # On a real terminal the panel runs on the alternate screen + # (htop-style): every refresh redraws the whole screen, so no + # frame can ever survive to tear — even when the frame fills the + # full terminal height, where cursor-up erasure is unreliable. + # On pipes/files Live cannot redraw in place at all; the + # per-refresh prints are skipped and stop() emits one final frame. + with Live( + boot_panel.render(ready), + refresh_per_second=4, + auto_refresh=interactive, + screen=interactive, + ) as boot_live: + while len(ready) < len(expected): + fresh = handshake.drain_ready() + if fresh: + ready |= fresh + if interactive: + boot_live.update(boot_panel.render(ready)) # log tail advances even without new readiness + try: + process_name, traceback_text = error_queue.get_nowait() + except Empty: + process_name = None + if process_name is not None: + control.set_stop() + _record_error_trace(process_name, traceback_text) + abort_reason = ( + f"worker {process_name} failed during startup " + f"(traceback: async_errors/{process_name}_error.log)" + ) + break + if time.time() > boot_deadline: + control.set_stop() + missing = sorted(expected - ready) + abort_reason = f"worker startup timed out after 600s; missing: {missing}" + break + time.sleep(0.25) + # ------------------------------------------------ handoff phase + # Release the barrier NOW: the handoff gate below waits for + # collectors to step (shared counters), and workers held at + # wait_for_start cannot step — waiting for the gate before + # releasing is a parent/child deadlock. + if abort_reason is None: + handshake.release() + # The barrier released every worker, but the training panel's + # queue-fed stats have not converged yet. Keep the SAME Live + # (and alt screen) running with a shared-counter progress + # view — every field is well-defined immediately — until the + # quiescence gate opens, then fall through to the training + # panel whose first frame is therefore fully aggregated. + if abort_reason is None: + while not _panel_data_ready(): + _drain_worker_stats() + try: + process_name, traceback_text = error_queue.get_nowait() + except Empty: + process_name = None + if process_name is not None: + control.set_stop() + abort_reason = f"worker {process_name} failed after startup" + break + if interactive: + boot_live.update(boot_panel.render(ready, gate=_handoff_gate_line(), starting=True)) + time.sleep(0.25) + # Live has exited — plain prints are safe on the terminal again. + if abort_reason is not None: + print(f"[motrix.fastsac async] {abort_reason}; aborting.") + raise RuntimeError(abort_reason) + print(f"[motrix.fastsac async] all {len(expected)} workers ready — training starts") + finally: + handshake.release() # idempotent; frees workers on an aborted boot + # -------------------------------------------------- parent-owned panel + # The parent is the only process that sees every worker (collectors' + # rollout snapshots + every learner rank's payload) AND the whole + # machine: it is unbound, so the system panel shows all cores, and + # workers stay peers with no "primary renderer" role. + console, live = None, None # opened lazily once real data flows (see below) + panel_open_attempted = False + try: + from torch.utils.tensorboard import SummaryWriter + + tb_writer = SummaryWriter(log_dir=str(self._context.run_dir)) + except Exception: + tb_writer = None + panel_cpu_sampler = CpuLoadSampler() + panel_gpu_sampler = GpuUtilizationSampler() + panel_memory_sampler = MemoryUsageSampler() + panel_gpu_memory_sampler = GpuMemoryUsageSampler() + panel_start_time = time.time() + panel_anchored = False + # Rolling (time, step) samples for the DISPLAYED env-steps/s — a + # time-boxed rate decoupled from the TensorBoard log-window + # bookkeeping below, which never resets the panel's number. + panel_rate_samples: list[tuple[float, int]] = [] + panel_last_render = 0.0 + panel_last_log_time = panel_start_time + panel_last_log_step = resume_step + panel_last_updates = 0 + panel_last_weight_version = 0 + panel_next_log = ( + ((resume_step // max(logging_interval, 1)) + 1) * logging_interval if logging_interval > 0 else 0 ) - p_learner.start() - p_collector.start() + def _render_panel(force: bool = False) -> None: + nonlocal panel_anchored, panel_start_time, panel_last_render, console, live, panel_open_attempted + nonlocal panel_rate_samples + nonlocal panel_last_log_time, panel_last_log_step, panel_last_updates, panel_next_log + nonlocal panel_last_weight_version + # ALWAYS drain first: the quiescence gate below waits for learner + # payloads that can only arrive through this drain — draining + # after the gate would deadlock the workers on a full queue. + _drain_worker_stats() + if console is None and not panel_open_attempted: + # Quiescence gate (see _panel_data_ready): stay off the + # terminal until every collector has reported at least one + # snapshot, a learner payload exists and collection started. + if not _panel_data_ready(): + return + panel_open_attempted = True + console, live = open_training_live() + if console is not None: + console.clear() + panel_last_render = 0.0 # render immediately after the gate opens + # non-TTY (nohup/redirect): open_training_live returns Nones — + # fall through to emit_training_panel's plain-frame path and + # never retry the open on every tick. + if not force and time.time() - panel_last_render < 1.0: + return + panel_last_render = time.time() + step = control.collector_steps // num_collectors + if step > resume_step and not panel_anchored: + panel_start_time = time.time() + panel_last_log_time = panel_start_time + panel_anchored = True + last_stats = aggregate_collector_stats(per_collector_stats) + payloads = [p for p in learner_payloads.values() if p] + primary = payloads[0] if payloads else None + now = time.time() + # Windowed rate. The window is meaningless right after the anchor + # (Δt ≈ 0 → absurd rate) or before any step landed (Δstep = 0); + # once it has accumulated real elapsed time a PARTIAL first + # window is already a honest rate, so report it instead of + # waiting a full logging interval on "warming up". + # Rolling rate over the last ~10s of render samples: monotone + # updates, no reset when a TB log point fires, honest from the + # first ~5s of stepping (below that there is not enough span). + panel_rate_samples.append((now, step)) + while len(panel_rate_samples) >= 2 and panel_rate_samples[1][0] <= now - 10.0: + panel_rate_samples.pop(0) + t0, s0 = panel_rate_samples[0] + t1, s1 = panel_rate_samples[-1] + # Short-window init: as soon as steps flow, compute over the + # available span (floor 1s against div-by-tiny) so the panel + # shows a real — if rough — rate immediately; None (warming up) + # only means "no steps flowed in the window" (pre-start or a + # genuine throughput stall). + if s1 > s0: + span = max(t1 - t0, 1.0) + sps = (s1 - s0) * num_envs / span + iters = (s1 - s0) / span + else: + sps = None + iters = None + warming = step < cfg.agent.learning_starts + # learn_ms is 0.0 (not None) in the immediate step-0 warmup + # payload — filter on None, not truthiness, or the list empties + # and timing_mean divides by zero. + learn_ms = timing_mean([p["learn_ms"] for p in payloads if p["learn_ms"] is not None]) if payloads else 0.0 + learn_pct = timing_mean([p["learn_pct"] for p in payloads]) if payloads else 0.0 + updates = primary["updates"] if primary else panel_last_updates + utd = updates / max(step, 1) + collector_timing_ms = last_stats.get("timing_ms", {}) + collector_timing_detail_ms = {k: v for k, v in collector_timing_ms.items() if k != "collect"} + collector_items: dict = {} + for key, value in collector_timing_detail_ms.items(): + nest_timing_path(collector_items, tuple(key.split(".")), value) + timing_groups: dict = {"collector": collector_items} + for payload in payloads: + rank_key = "learner" if num_learners == 1 else f"learner[{payload['rank']}]" + if payload["learner_timing"]: + timing_groups[rank_key] = payload["learner_timing"] + stats = TrainingPanelStats( + iteration=step, + total_iterations=num_iterations, + steps_per_second=sps, + iterations_per_second=iters, + elapsed_seconds=now - panel_start_time, + mean_return=last_stats["return"], + mean_episode_length=last_stats["ep_len"], + episodes=last_stats["episodes"], + buffer_size=primary["buffer_size"] if primary else 0, + buffer_capacity=primary["buffer_capacity"] if primary else 0, + collect_ms=collector_timing_ms.get("collect", 0.0), + learn_ms=learn_ms, + learn_percent=learn_pct, + warming=warming, + training_metrics=primary["metrics"] if primary else None, + reward_terms=last_stats["reward_terms"], + env_metrics=last_stats["env_metrics"], + timing_groups=timing_groups, + diagnostics={"UTD": utd}, + cpu_load=panel_cpu_sampler.sample(), + gpu_utilization_percent=panel_gpu_sampler.sample(), + memory_usage=panel_memory_sampler.sample(), + gpu_memory_usage=panel_gpu_memory_sampler.sample(), + gpu_devices=sample_gpu_devices(panel_gpu_sampler, panel_gpu_memory_sampler), + checkpoint_path=primary["checkpoint_path"] if primary else None, + ) + emit_training_panel(live, stats, title=f"{self._env_name}/motrix.fastsac") + if tb_writer is not None and step >= panel_next_log: + tb_writer.add_scalar("rollout/mean_return", last_stats["return"], step) + tb_writer.add_scalar("rollout/mean_ep_len", last_stats["ep_len"], step) + if sps is not None: + tb_writer.add_scalar("perf/env_steps_per_s", sps, step) + tb_writer.add_scalar( + "perf/updates_per_s", (updates - panel_last_updates) / max(now - panel_last_log_time, 1e-6), step + ) + tb_writer.add_scalar("async/policy_lag", last_stats["policy_lag"], step) + if primary and primary.get("weight_version") is not None: + tb_writer.add_scalar( + "async/weight_publishes_per_iter", + (primary["weight_version"] - panel_last_weight_version) / max(step - panel_last_log_step, 1), + step, + ) + panel_last_weight_version = primary["weight_version"] + if primary: + fills = primary.get("ring_fill") or [] + tb_writer.add_scalar("async/ring_fill", sum(f for f in fills if f is not None), step) + if primary.get("weight_version") is not None: + tb_writer.add_scalar("async/weight_version", primary["weight_version"], step) + tb_writer.add_scalar("async/utd", utd, step) + if num_collectors > 1: + for i in range(num_collectors): + stats_i = per_collector_stats.get(i) + if stats_i: + tb_writer.add_scalar(f"async/policy_lag_collector{i}", stats_i["policy_lag"], step) + if primary: + fill = (primary.get("ring_fill") or [None] * num_collectors)[i] + if fill is not None: + tb_writer.add_scalar(f"async/ring_fill_collector{i}", fill, step) + tb_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(): + tb_writer.add_scalar(f"perf/collector_{k}_ms", v, step) + tb_writer.add_scalar("perf/learn_ms_total", learn_ms, step) + tb_writer.add_scalar("perf/learn_pct", learn_pct, step) + for k, v in last_stats["env_metrics"].items(): + tb_writer.add_scalar(f"metrics/{k}", v, step) + for k, v in last_stats["reward_terms"].items(): + tb_writer.add_scalar(f"reward/{k}", v, step) + if primary and primary["metrics"] is not None: + for k, v in primary["metrics"].items(): + tb_writer.add_scalar(f"train/{k}", v, step) + panel_last_log_time, panel_last_log_step, panel_last_updates = now, step, updates + panel_next_log += logging_interval + try: - # monitor: exit when the learner finishes; abort both if either crashes. + # monitor: exit when every learner has finished; abort all if any worker crashes. while True: - if not p_learner.is_alive(): - if p_learner.exitcode not in (0, None): - print(f"[motrix.fastsac async] learner crashed (exit {p_learner.exitcode}); stopping.") - _drain_child_errors() - break - if not p_collector.is_alive() and p_collector.exitcode not in (0, None): - print(f"[motrix.fastsac async] collector crashed (exit {p_collector.exitcode}); stopping.") + crashed = [p for p in (*p_learners, *p_collectors) if not p.is_alive() and p.exitcode not in (0, None)] + learners_done = [p for p in p_learners if not p.is_alive()] + if crashed: + for p in crashed: + print(f"[motrix.fastsac async] {p.name} crashed (exit {p.exitcode}); stopping.") _drain_child_errors() break + if learners_done: + _drain_child_errors() + if len(learners_done) == len(p_learners): + break + _render_panel() time.sleep(0.5) finally: + # Ctrl+C reaches the parent as KeyboardInterrupt and every worker via + # its SIGINT handler (stop-flag unwind). Shutdown is an escalation + # ladder with per-step exception isolation: one stuck child must + # never skip the cleanup of the others (that is how orphans holding + # GPU memory were produced). control.set_stop() - p_collector.join(timeout=30) - p_learner.join(timeout=30) - for p in (p_collector, p_learner): - if p.is_alive(): - print(f"[motrix.fastsac async] force-terminating {p.name}") - p.terminate() - p.join(timeout=10) - _drain_child_errors() - # drain the queue so the feeder thread can shut down cleanly. + + def _shutdown(processes, grace_s): + for p in processes: + try: + p.join(timeout=grace_s) + if not p.is_alive(): + continue + print(f"[motrix.fastsac async] terminating {p.name}") + p.terminate() # SIGTERM: unwinds queue feeder threads + p.join(timeout=3.0) + if p.is_alive(): + print(f"[motrix.fastsac async] killing {p.name}") + p.kill() # SIGKILL: e.g. blocked inside a NCCL collective + p.join(timeout=3.0) + except Exception as exc: # noqa: BLE001 - isolation by design + print(f"[motrix.fastsac async] shutdown of {p.name} failed: {exc}") + + _shutdown(p_collectors, grace_s=8.0) + _shutdown(p_learners, grace_s=8.0) try: - while True: - stats_queue.get_nowait() + _drain_child_errors() + except Exception: + pass + try: + _render_panel(force=True) + except Exception: + pass + if live is not None: + try: + live.stop() + except Exception: + pass + if tb_writer is not None: + try: + tb_writer.close() + except Exception: + pass + try: + panel_queue.close() + except Exception: + pass + handshake.close() + for stats_queue in stats_queues: + try: + while True: + stats_queue.get_nowait() + except Exception: + pass + stats_queue.close() + try: + error_queue.close() except Exception: pass - stats_queue.close() - error_queue.close() - if p_learner.exitcode not in (0, None): - raise RuntimeError(f"motrix.fastsac async learner process failed with exit code {p_learner.exitcode}") - if p_collector.exitcode not in (0, None): - raise RuntimeError(f"motrix.fastsac async collector process failed with exit code {p_collector.exitcode}") + for p in p_learners: + if p.exitcode not in (0, None): + raise RuntimeError( + f"motrix.fastsac async learner process '{p.name}' failed with exit code {p.exitcode}" + ) + for i, p in enumerate(p_collectors): + if p.exitcode not in (0, None): + raise RuntimeError(f"motrix.fastsac async collector {i} process failed with exit code {p.exitcode}") # ------------------------------------------------------------------ play def play(self, policy: str) -> None: diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py index d30f8bc8..8f134f6d 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py @@ -27,6 +27,7 @@ flatten_params, load_flat_params, ) +from motrix_rl.fastsac.async_impl.transport.handshake import StartupHandshake from motrix_rl.fastsac.async_impl.transport.ipc_ring import IpcTransitionRing from motrix_rl.fastsac.async_impl.transport.ring import RingCursors, SharedTransitionRing from motrix_rl.fastsac.async_impl.transport.weight_channel import ( @@ -42,6 +43,7 @@ __all__ = [ "Control", + "StartupHandshake", "GpuIpcWeightReceiver", "GpuIpcWeightSender", "HostWeightReceiver", diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py index afcf4d35..3b8f540e 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py @@ -58,12 +58,23 @@ def _shared(shape, dtype) -> torch.Tensor: # ---------------------------------------------------------------- control block class Control: - """A handful of shared scalar controls / counters.""" + """A handful of shared scalar controls / counters. + + ``collector_steps`` is a per-collector counter array: each collector + increments only its own entry (single writer), and the aggregate property + sums them. The learner uses the aggregate as the training-progress basis; + each collector compares its own entry against ``num_iterations`` / + ``learning_starts`` (one entry == one env-step batch of that collector's + env shard). With ``num_collectors=1`` this is byte-identical to the previous + single-counter behavior. + """ - def __init__(self): + def __init__(self, num_collectors: int = 1): + self.num_collectors = num_collectors 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 + # one env-step-batch counter per collector, single-writer each + self._collector_steps = [_shared((1,), torch.int64) for _ in range(num_collectors)] @property def stop(self) -> bool: @@ -82,14 +93,29 @@ def global_step(self, v: int) -> None: @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 + """Aggregate env-step batches produced by all collectors.""" + return sum(int(counter[0]) for counter in self._collector_steps) + + def resume_collector_steps(self, per_collector_v: int) -> None: + """Restart every collector's own counter at ``per_collector_v``. + + Resume entry point: ``per_collector_v`` is the checkpointed training + step in full ``num_envs``-batch equivalents (== the checkpoint's + ``global_step``); the aggregate becomes ``num_collectors * + per_collector_v``, matching the invariant that the learner's progress + basis is the aggregate while each collector terminates against its own + entry. Intentionally a named method, not a setter: the property reads + as the aggregate while this writes per-collector values — same units + under both directions would be a trap. + """ + for counter in self._collector_steps: + counter[0] = per_collector_v + + def collector_steps_at(self, collector_id: int) -> int: + return int(self._collector_steps[collector_id][0]) + + def inc_collector_steps(self, collector_id: int) -> None: + self._collector_steps[collector_id][0] += 1 # ---------------------------------------------------------------- flat-param helpers diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/handshake.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/handshake.py new file mode 100644 index 00000000..a883d2c5 --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/handshake.py @@ -0,0 +1,115 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Startup rendezvous for the async trainer: slot delivery + readiness barrier. + +One structured object owns every piece of the boot handshake that used to be +scattered across per-worker queues: + +* one-shot slot delivery — learners ship CUDA-IPC handle bundles (weight + slots, transition-ring device slots) to the collectors they feed, through + per-collector one-shot queues; +* the readiness barrier — every worker reports readiness and then holds on + the start event until the parent releases everyone at once, so stepping + begins in lockstep (warmup counters, weight generations and generation + merges align from step one). + +Created in the parent with the spawn context and inherited by the children +through the process args; the parent is the only ``release()`` caller. +With ``barrier=False`` the readiness half is inert (direct callers that +drive a single collector by hand). +""" + +from __future__ import annotations + +from multiprocessing.queues import Queue + + +class StartupHandshake: + """Parent-allocated boot rendezvous shared by every trainer worker.""" + + SLOT_TIMEOUT_S = 60.0 + + def __init__(self, ctx, num_collectors: int, barrier: bool = True): + self.barrier = barrier + self.ready_queue = ctx.Queue() + self.start_event = ctx.Event() + # One-shot slot queue per collector (maxsize 1: exactly one message + # per direction is ever sent). Ring-slot messages carry ``None`` for + # the host shared-memory ring. + self.slot_queues: list[Queue] = [ctx.Queue(maxsize=1) for _ in range(num_collectors)] + self.ring_slot_queues: list[Queue] = [ctx.Queue(maxsize=1) for _ in range(num_collectors)] + + # ------------------------------------------------------- learner side + def ship_weight_slots(self, collector_id: int, params) -> None: + """Ship the learner-built weight-slot pair for one collector.""" + self.slot_queues[collector_id].put(params) + + def ship_ring_slots(self, collector_id: int, ring_slots) -> None: + """Ship the CUDA-IPC transition-ring slots (``None`` for host rings).""" + self.ring_slot_queues[collector_id].put(ring_slots) + + # ------------------------------------------------------ collector side + def await_weight_slots(self, collector_id: int): + """Block until the publishing learner ships this collector's weight slots.""" + return self._await(self.slot_queues[collector_id], "weight-slot") + + def await_ring_slots(self, collector_id: int): + """Block until the draining learner ships this collector's ring slots. + + The message is ``None`` when the collector's ring is host shared + memory (no device slots to map). + """ + return self._await(self.ring_slot_queues[collector_id], "ring-slot") + + def _await(self, queue: Queue, kind: str): + import queue as queue_mod + + try: + return queue.get(timeout=self.SLOT_TIMEOUT_S) + except queue_mod.Empty as exc: + raise RuntimeError( + f"timed out after {self.SLOT_TIMEOUT_S:.0f}s waiting for the learner to ship the {kind} tensors " + "(learner startup — agent build / checkpoint load / CUDA warmup — " + "likely failed; check the learner process's error queue/log)" + ) from exc + + # ------------------------------------------------------------- barrier + def report_ready(self, role: str, worker_id: int) -> None: + """Report this worker as booted (no-op when the barrier is disabled).""" + if self.barrier: + self.ready_queue.put((role, worker_id)) + + def wait_for_start(self, control) -> None: + """Hold until the parent releases every worker (no-op without barrier).""" + if not self.barrier: + return + while not self.start_event.is_set() and not control.stop: + self.start_event.wait(timeout=0.2) + + def drain_ready(self) -> set: + """Pop every readiness report delivered so far (parent only).""" + import queue as queue_mod + + ready: set = set() + try: + while True: + ready.add(self.ready_queue.get_nowait()) + except queue_mod.Empty: + pass + return ready + + def release(self) -> None: + """Release the barrier; also used to free workers on an aborted boot.""" + self.start_event.set() + + def close(self) -> None: + """Drain and close the one-shot queues so feeder threads can shut down.""" + for one_shot in (*self.slot_queues, *self.ring_slot_queues): + try: + while True: + one_shot.get_nowait() + except Exception: + pass + one_shot.close() + self.ready_queue.close() 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 86a13e63..eb68910c 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -21,76 +21,49 @@ import sys import time import traceback +from contextlib import contextmanager from multiprocessing.queues import Queue from pathlib import Path -from queue import Empty -from typing import Any +from queue import Empty, Full +from typing import TYPE_CHECKING, Any import numpy as np import torch +if TYPE_CHECKING: + from torch.utils.tensorboard import SummaryWriter + from motrix_env_core.array.env import ArrayEnv from motrix_env_core.registry import EnvBuildSpec from motrix_env_core.renderer import RenderConfig from motrix_env_motrixsim.torch_env import TorchEnv 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, resolve_collector_inference_device -from motrix_rl.fastsac.async_impl.learner import Learner +from motrix_rl.fastsac.async_impl.collector import Collector +from motrix_rl.fastsac.async_impl.learner import CollectorEndpoint, Learner +from motrix_rl.fastsac.async_impl.numa import apply_binding +from motrix_rl.fastsac.async_impl.stats import timing_mean from motrix_rl.fastsac.async_impl.transport import ( Control, IpcTransitionRing, RingCursors, SharedTransitionRing, ) +from motrix_rl.fastsac.async_impl.transport.handshake import StartupHandshake from motrix_rl.fastsac.async_impl.transport.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 from motrix_rl.fastsac.wrap_torch import FastSacTorchEnvWrap -from motrix_rl.system_metrics import ( - CpuLoadSampler, - GpuMemoryUsageSampler, - GpuUtilizationSampler, - MemoryUsageSampler, - sample_gpu_devices, -) - - -def _timing_mean(values: list[float]) -> float: - """Mean of a non-empty timing sample list (ms).""" - return sum(values) / len(values) - - -def _nest_timing_path(tree: dict[str, Any], parts: tuple[str, ...], value: float) -> None: - """Insert one dotted timing path into a nested mapping. - A stage's scalar total and its sub-stage paths may arrive in either order - (the collector emits parents before children); when both exist the scalar - becomes the node's ``total`` alongside its children. - """ - head, rest = parts[0], parts[1:] - node = tree.get(head) - if not rest: - if isinstance(node, dict): - node["total"] = value - else: - tree[head] = value - else: - if not isinstance(node, dict): - node = {"total": node} if node is not None else {} - tree[head] = node - _nest_timing_path(node, rest, value) +logger = logging.getLogger(__name__) -# ------------------------------------------------------------------ builders def set_seed(seed: int | None) -> None: if seed is None: return @@ -142,7 +115,23 @@ def actor_param_numel(cfg: FastSacCfg, dims, action_scale, action_bias) -> int: return sum(p.numel() for p in actor.parameters()) -def build_agent(cfg: FastSacCfg, dims, num_envs, device, action_scale, action_bias, writer=None) -> FastSacAgent: +def build_agent( + cfg: FastSacCfg, + dims: tuple[int, int, int], + num_envs: int, + device: torch.device, + action_scale: torch.Tensor | None, + action_bias: torch.Tensor | None, + writer: SummaryWriter | None = None, + world_size: int = 1, +) -> FastSacAgent: + """Build the learner's ``FastSacAgent``. + + ``dims`` is ``(obs_dim, critic_obs_dim, act_dim)``. ``action_scale`` / + ``action_bias`` map the tanh-squashed policy output to the env action + range (``None`` = identity). ``writer`` is the parent-created + TensorBoard writer; ``world_size`` is the DDP learner-rank count. + """ obs_dim, critic_obs_dim, act_dim = dims return FastSacAgent( obs_dim=obs_dim, @@ -154,62 +143,10 @@ def build_agent(cfg: FastSacCfg, dims, num_envs, device, action_scale, action_bi action_scale=action_scale, action_bias=action_bias, writer=writer, + world_size=world_size, ) -def same_cuda_device(learner_device: torch.device, collector_device: torch.device) -> bool: - """Whether the learner and the collector's inference device share one GPU. - - An index-less ``cuda`` means the default current device (index 0 — nothing - in the trainer ever calls ``torch.cuda.set_device``), so it is resolved - with 0 rather than treated as a wildcard matching any explicit index: - ``learner=cuda`` (effectively cuda:0) with ``collector_inference_device: - cuda:1`` is a cross-GPU setup and must NOT enable the device transports. - Pure device arithmetic — no CUDA context is created, so the pre-spawn - parent can call it as safely as the workers. - """ - if learner_device.type != "cuda" or collector_device.type != "cuda": - return False - learner_index = learner_device.index if learner_device.index is not None else 0 - collector_index = collector_device.index if collector_device.index is not None else 0 - return learner_index == collector_index - - -def use_ipc_transition_ring(opts, learner_device: torch.device, collector_device: torch.device) -> bool: - """Whether the transition ring should use CUDA-IPC device slots. - - Requires learner and collector inference on the same GPU (see - :func:`same_cuda_device`); otherwise the host shared-memory ring is used. - Purely device-object arithmetic — no CUDA context is created here, so the - parent can call it safely. - """ - mode = opts.transition_ipc - # YAML 1.1 parses unquoted ``on``/``off`` scalars as booleans; accept that - # form so ``transition_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.transition_ipc must be auto, on or off, got {mode!r}") - same_gpu = same_cuda_device(learner_device, collector_device) - if mode == "off": - return False - if mode == "on" and not same_gpu: - reason = ( - "collector inference device is not CUDA" - if collector_device.type != "cuda" - else f"learner device {learner_device} and collector device {collector_device} are different GPUs" - if learner_device.type == "cuda" - else "learner device is not CUDA" - ) - logging.getLogger(__name__).warning( - "async_options.transition_ipc=on requires learner and collector inference on the same GPU, " - "but %s; falling back to the host shared-memory transition ring", - reason, - ) - return same_gpu - - # ------------------------------------------------------------------ collector process def _available_cpu_ids() -> set[int]: """CPU ids this process may run on (Linux affinity mask, Windows process mask).""" @@ -302,62 +239,74 @@ def _pin_worker_cpus(cpus: set[int]) -> None: 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. +def _install_graceful_sigint(control: Control) -> None: + """Turn SIGINT (terminal Ctrl+C hits the whole process group) into the + shared stop flag so the worker unwinds at its next loop safe point instead + of dying mid-CUDA/queue operation with a traceback. - ``basicConfig`` only applies when the root logger has no handlers yet; when - a handler is already configured, raising the root level is enough for the - startup INFO records to be emitted through the existing setup. + A worker blocked inside a NCCL collective never returns to the interpreter, + so the handler cannot fire there — the parent's escalation ladder + (join -> terminate -> kill, see train.py) is the backstop for that case. """ - root = logging.getLogger() - if root.handlers: - root.setLevel(logging.INFO) - return - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + import signal + def _handler(signum, frame): # noqa: ARG001 + control.set_stop() -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). + try: + signal.signal(signal.SIGINT, _handler) + except ValueError: # not on the main thread (defensive; workers are processes) + pass + + +@contextmanager +def inherit_log_stdio(log_file: Path): + """Parent-side spawn guard: point fd 1/2 at the worker's log file while starting it. + + Spawn children inherit the parent's terminal fds and re-import + torch/simulator modules BEFORE their entry function runs + (``_configure_process_logging``), so import-time prints (library banners, + profiler messages, warnings) would land raw on the shared terminal and + jitter the parent's live panel. Holding the log file on fd 1/2 across + ``Process.start()`` makes the child inherit it from its first + instruction; the parent's own fds are restored immediately afterwards. """ - 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 = same_cuda_device(device, collector_device) - 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) + log_fd = os.open(str(log_file), os.O_WRONLY | os.O_CREAT | os.O_APPEND) + saved_stdout, saved_stderr = os.dup(1), os.dup(2) + try: + os.dup2(log_fd, 1) + os.dup2(log_fd, 2) + yield + finally: + os.dup2(saved_stdout, 1) + os.dup2(saved_stderr, 2) + os.close(saved_stdout) + os.close(saved_stderr) + os.close(log_fd) + + +def _configure_process_logging(log_file: Path) -> None: + """Route worker logs (e.g. manager env startup) to a file, not the terminal. + + The parent renders the live panel on the shared terminal; raw worker log + bytes would tear the rich Live frames apart. Startup records stay + inspectable under the run's ``logs/`` directory; errors additionally + travel the error-queue channel as before. + """ + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + filename=str(log_file), + filemode="a", + force=True, + ) + # Hard hat: dup the log file onto stdout/stderr so C-level output + # (third-party import banners, library warnings) can never reach the + # shared terminal and tear the parent's Rich panel. + log_fd = os.open(str(log_file), os.O_WRONLY | os.O_CREAT | os.O_APPEND) + os.dup2(log_fd, 1) + os.dup2(log_fd, 2) + os.close(log_fd) def run_collector_process( @@ -376,40 +325,54 @@ def run_collector_process( logging_interval: int, is_resume: bool, seed: int | None, - slot_queue: Queue, + run_dir: str, + handshake: StartupHandshake, + collector_id: int = 0, + numa_node: int | None = None, + cpus: list[int] | None = None, + collector_device: str | None = None, ) -> None: + role = f"collector[{collector_id}]" + _install_graceful_sigint(control) try: - _configure_process_logging() + log_dir = Path(run_dir) / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + _configure_process_logging(log_dir / f"collector{collector_id}.log") + # Bind before any env / staging allocation: the memory policy only + # affects future pages, so this must be the first thing the worker does. + # The binding itself was decided by the parent topology pass. + apply_binding(role, numa_node, cpus or []) set_seed(seed) opts = cfg.trainer.async_options + # Multi-learner: the parent resolves the generic "cuda" spec to the + # owning learner's GPU; an explicit spec passes through unchanged. + if collector_device is not None: + opts.collector_inference_device = collector_device _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") + env_started = time.perf_counter() env = build_env(env_spec, num_envs, device, seed=seed) - # Handshake: build the endpoints from the slot tensors the learner - # shipped (weight slots, plus the CUDA-IPC transition-ring slots when - # that transport was selected), before the collector is wired up. - try: - weight_slots, ring_slots = slot_queue.get(timeout=60.0) - except Empty as exc: - raise RuntimeError( - "timed out waiting for the learner to ship the 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 + logger.info( + "collector startup: env build (scene + manager kernels) finished in %.3fs", + time.perf_counter() - env_started, + ) + # Handshake: build the endpoints from the slot tensors the learners + # shipped before the collector is wired up. Weight slots come from + # the publishing learner (rank 0), CUDA-IPC transition-ring slots + # from the learner that drains this collector's ring (its owner — + # always rank 0 in the single-learner topology). + weight_slots = handshake.await_weight_slots(collector_id) + ring_slots = handshake.await_ring_slots(collector_id) weight_rx = weight_receiver_for(weights, weight_slots) if isinstance(ring, RingCursors): if ring_slots is None: raise RuntimeError("the learner shipped no transition-ring slots for the IPC ring handshake") - ring = IpcTransitionRing( - ring, ring_slots, async_options.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim - ) + ring = IpcTransitionRing(ring, ring_slots, opts.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim) collector = Collector( env, cfg, obs_dim, - critic_obs_dim, act_dim, action_scale, action_bias, @@ -417,16 +380,43 @@ def run_collector_process( weight_rx, control, is_resume=is_resume, + collector_id=collector_id, ) collector.reset() collector.sync_weights() + warmup_started = time.perf_counter() collector.warmup_inference() + logger.info( + "collector startup: inference warmup (torch.compile + CUDA graphs) finished in %.3fs", + time.perf_counter() - warmup_started, + ) - while not control.stop and control.collector_steps < num_iterations: + # Startup barrier: report readiness and hold until every worker (all + # collectors AND learners) has booted, so stepping starts in lockstep + # — warmup counters, weight generations and generation merge align + # from step one instead of converging mid-run. + handshake.report_ready("collector", collector_id) + handshake.wait_for_start(control) + if control.stop: + return + + # First snapshot once a few steps have landed: the parent's panel + # quiescence gate opens early (not after a full logging window of + # silence) while the handoff progress view gets to show real steps, + # and the snapshot already carries a little real data. The + # per-interval send below then replaces it. + first_stats_steps = 10 + first_stats_sent = False + + while not control.stop and control.collector_steps_at(collector_id) < num_iterations: if not collector.step_once(): - time.sleep(async_options.idle_sleep_s) # ring full -> backpressure + time.sleep(opts.idle_sleep_s) # ring full -> backpressure continue - if collector.control.collector_steps % max(logging_interval, 1) == 0: + steps = control.collector_steps_at(collector_id) + if not first_stats_sent and steps >= first_stats_steps: + first_stats_sent = True + stats_queue.put(collector.snapshot_stats()) + elif steps % max(logging_interval, 1) == 0: # replace any stale snapshot so the learner always sees the latest. try: while True: @@ -435,7 +425,10 @@ def run_collector_process( pass stats_queue.put(collector.snapshot_stats()) except BaseException: - error_queue.put(("collector", traceback.format_exc())) + if not isinstance(sys.exc_info()[1], (KeyboardInterrupt, SystemExit)): + # Ctrl+C is an operator action, not a defect: the stop flag is + # already set by the SIGINT handler; don't spam async_errors. + error_queue.put((role, traceback.format_exc())) raise finally: control.set_stop() # signal the learner if the collector exits for any reason @@ -448,64 +441,164 @@ def run_learner_process( dims: tuple[int, int, int], action_scale: torch.Tensor, action_bias: torch.Tensor, - ring: SharedTransitionRing | RingCursors, - weights: WeightChannelShared, + rings: list[SharedTransitionRing | RingCursors], + weights: list[WeightChannelShared], control: Control, - stats_queue: Queue, - error_queue: Queue, + error_queue, num_iterations: int, logging_interval: int, save_interval: int, run_dir: str, - env_name: str, checkpoint_dir: str, checkpoint_format: str, resume_from: str | None, seed: int | None, - slot_queue: Queue, + handshake: StartupHandshake, + all_rings: list[SharedTransitionRing], + weight_ipc: list[bool], + panel_queue: Queue, + learner_cpus: list[int] | None = None, + learner_numa_node: int | None = None, + rank: int = 0, + num_learners: int = 1, + rendezvous_file: str | None = None, + learner_device: str | None = None, ) -> None: - _configure_process_logging() - console, live = open_training_live() + """One learner process; ``rank 0`` additionally owns stats aggregation, + logging and checkpointing. + + Multi-learner (DDP) specifics: ``rings`` is this rank's slice (the parent + partitions by collector ownership) while ``all_rings`` is the full list + (rank-0 logging only); ``weights`` is this rank's slice of the weight + channels — every learner publishes to its OWN collectors only; update + counts are derived from the global progress basis (see + ``Learner.maybe_train``). + """ + is_primary = num_learners == 1 or rank == 0 + if num_learners > 1: + import torch.distributed as dist + else: + dist = None + _install_graceful_sigint(control) + _learner_log_dir = Path(run_dir) / "logs" + _learner_log_dir.mkdir(parents=True, exist_ok=True) + _configure_process_logging(_learner_log_dir / (f"learner{rank}.log" if num_learners > 1 else "learner.log")) try: - set_seed(seed) + apply_binding(f"learner[{rank}]" if num_learners > 1 else "learner", learner_numa_node, learner_cpus) + set_seed(None if seed is None else seed + rank) 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 - try: - from torch.utils.tensorboard import SummaryWriter - - writer = SummaryWriter(log_dir=run_dir) - except Exception: - writer = None - - agent = build_agent(cfg, dims, num_envs, device, action_scale, action_bias, writer=writer) + num_collectors = len(all_rings) + device = ( + torch.device(learner_device) + if learner_device is not None + else torch.device(cfg.device or ("cuda" if torch.cuda.is_available() else "cpu")) + ) + if num_learners > 1: + if device.type != "cuda" or device.index is None: + raise ValueError(f"multi-learner requires explicit indexed CUDA learner devices, got {device}") + torch.cuda.set_device(device) + # Bounded collective timeout: with the default 30-minute timeout a + # rank waiting on a crashed/blocked peer hangs until the parent's + # force-kill instead of unwinding through its own error path. + import datetime as _dt + + dist.init_process_group( + "nccl", + init_method=f"file://{rendezvous_file}", + rank=rank, + world_size=num_learners, + timeout=_dt.timedelta(seconds=90), + ) + # Replay buffer sized by THIS rank's env shard: each learner drains only + # its collectors' rings, so every ingested batch carries + # num_envs // num_learners envs (env-level sharding keeps each env's + # full trajectory — and n-step adjacency — inside one rank). The batch + # math still divides the GLOBAL batch_size by world_size, so per-rank + # sample rows == batch_size / num_learners, matching DDP averaging. + agent_started = time.perf_counter() + agent = build_agent( + cfg, + dims, + num_envs // num_learners, + device, + action_scale, + action_bias, + world_size=num_learners, + ) if resume_from: ckpt = torch.load(resume_from, map_location=device, weights_only=False) agent.load_state_dict(ckpt, load_optimizers=True) + logger.info( + "learner startup: agent build%s finished in %.3fs", + " + checkpoint load" if resume_from else "", + time.perf_counter() - agent_started, + ) - # 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. The - # IPC transition-ring slots ship in the same one-shot message: this - # process is the owner (it allocated the device tensor and must keep it - # alive), the collector maps it through the queue's CUDA-IPC reducers. - weight_tx = _build_weight_sender(weights, cfg, dims, action_scale, action_bias, device) - if isinstance(ring, RingCursors): - if device.type != "cuda": - raise RuntimeError("the IPC transition ring was selected but the learner device is not CUDA") - obs_dim, critic_obs_dim, act_dim = dims - feat = obs_dim + critic_obs_dim + act_dim + 3 - ring_slots = torch.zeros(async_options.ring_capacity, num_envs, feat, dtype=torch.float32, device=device) - ring = IpcTransitionRing( - ring, ring_slots, async_options.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim + resume_step = control.collector_steps // num_collectors # full-batch equivalents + per_collector_envs = num_envs // num_collectors + per_learner = num_collectors // num_learners + # Build one sender endpoint per collector OWNED by this rank (collector + # i belongs to rank i // per_learner) and ship the slot tensors BEFORE + # the first publish so each collector (blocking on its handshake queue) + # builds the matching receiver from step one. The IPC transition-ring + # slots ship in the same one-shot message: this process is the owner + # (it allocated the device tensor and must keep it alive), the + # collector maps it through the queue's CUDA-IPC reducers. + # Transport per channel was decided by the parent topology pass + # (TrainerTopology.weight_ipc, indexed per collector); this process + # only constructs the endpoints. One throwaway CPU actor counts the + # parameters that size both transports' buffers. + param_numel = actor_param_numel(cfg, dims, action_scale, action_bias) + weight_txs = [] + for j, shared in enumerate(weights): + channel = rank * per_learner + j + # CUDA-IPC device slots when the topology resolved IPC for this + # channel (the collector infers on this rank's GPU), host + # shared-memory slots otherwise; this rank is the exporter (its + # CUDA context must keep the slot tensors alive). + tx = ( + GpuIpcWeightSender(shared, param_numel, device) + if weight_ipc[j] + else HostWeightSender(shared, param_numel) ) - slot_queue.put((weight_tx.params, ring_slots)) - else: - slot_queue.put((weight_tx.params, None)) - learner = Learner(agent, cfg, ring, weight_tx, control) - learner.publish_weights() # give the collector an initial policy before it warms up + handshake.ship_weight_slots(channel, tx.params) + weight_txs.append(tx) + # Wrap RingCursors into IpcTransitionRing (allocating the device slots + # and shipping them to the collector) BEFORE building the endpoints — + # the endpoints must hold the FINAL ring objects the drain path uses. + for j, ring in enumerate(rings): + ring_slots = None + if isinstance(ring, RingCursors): + if device.type != "cuda": + raise RuntimeError("the IPC transition ring was selected but the learner device is not CUDA") + obs_dim, critic_obs_dim, act_dim = dims + feat = obs_dim + critic_obs_dim + act_dim + 3 + ring_slots = torch.zeros( + opts.ring_capacity, per_collector_envs, feat, dtype=torch.float32, device=device + ) + rings[j] = IpcTransitionRing( + ring, + ring_slots, + opts.ring_capacity, + per_collector_envs, + obs_dim, + critic_obs_dim, + act_dim, + ) + handshake.ship_ring_slots(rank * per_learner + j, ring_slots) + + endpoints = [CollectorEndpoint(ring=rings[j], weight_sender=weight_txs[j]) for j in range(len(weight_txs))] + + learner = Learner( + agent, + cfg, + endpoints, + control, + ddp_rank=rank if num_learners > 1 else None, + ) + learner._last_train_gstep = resume_step # lockstep basis continues from the checkpoint + learner.publish_weights() # give every collector an initial policy before it warms up # Elapsed/ETA anchor. The collector's env build (scene compile, numba # JIT) runs concurrently with the learner build and can outlast it by @@ -515,20 +608,16 @@ def run_learner_process( start_time = time.time() start_anchored = False last_log_time = start_time - resume_step = control.collector_steps last_log_step = resume_step - last_update_idx = 0 - last_stats = { - "return": float("nan"), - "ep_len": float("nan"), - "episodes": 0, - "reward_terms": {}, - "env_metrics": {}, - "policy_lag": 0, - "timing_ms": {}, - } last_metrics = None - next_log = ((resume_step // logging_interval) + 1) * logging_interval if logging_interval > 0 else 0 + # First payload goes out immediately (step 0) so the parent's panel + # quiescence gate opens as soon as collectors take their first steps, + # instead of after a silent full log window; resume keeps window + # alignment. + if logging_interval > 0 and resume_step: + next_log = ((resume_step // logging_interval) + 1) * logging_interval + else: + next_log = 0 next_save = ((resume_step // save_interval) + 1) * save_interval if save_interval > 0 else 0 t_learn_win = 0.0 # wall-clock spent in learner train calls this log window learner_train_samples_ms: list[float] = [] @@ -536,21 +625,32 @@ def run_learner_process( learner_breakdown_samples_ms: dict[str, list[float]] = {} learner_ring_wait_samples_ms: list[float] = [] learner_gate_wait_samples_ms: list[float] = [] - cpu_sampler = CpuLoadSampler() - gpu_sampler = GpuUtilizationSampler() - memory_sampler = MemoryUsageSampler() - gpu_memory_sampler = GpuMemoryUsageSampler() last_checkpoint_path: str | None = None - - def _drain_stats(): - nonlocal last_stats - try: - while True: - last_stats = stats_queue.get_nowait() - except Empty: - pass - - while not control.stop and control.collector_steps < num_iterations: + # The parent process renders the panel and writes TensorBoard: this + # worker ships one compact payload per log window (all learner ranks + # send — workers are peers, the parent is the recorder). + + # Startup barrier: report readiness and hold until every worker booted + # (see run_collector_process). Weight slots were already shipped and + # the initial weights published, so collectors blocking on their + # handshakes complete as soon as every learner reaches this point. + handshake.report_ready("learner", rank) + handshake.wait_for_start(control) + if control.stop: + return + + # Each collector runs until its own counter reaches num_iterations; the + # aggregate (sum) therefore reaches num_iterations * num_collectors. + total_collector_steps = num_iterations * num_collectors + while not control.stop and control.collector_steps < total_collector_steps: + # Progress basis: full num_envs-batch equivalents — the aggregate + # counter counts per-collector batches of per_collector_envs + # transitions each, so dividing by num_collectors recovers the + # sync-trainer "iteration collects num_envs transitions" unit + # (identical to the raw counter when num_collectors == 1). All + # ranks read the same shared value. + step = control.collector_steps // num_collectors + control.global_step = step t_drain = time.perf_counter() ingested = learner.drain() if ingested: @@ -560,7 +660,10 @@ def _drain_stats(): start_anchored = True learner_drain_samples_ms.append((time.perf_counter() - t_drain) * 1000.0) t_l = time.perf_counter() - metrics = learner.maybe_train(ingested) + # One decision path for both topologies: updates are gated on the + # global progress basis (identical on every DDP rank; the single + # learner's own position when num_collectors == 1). + metrics = learner.maybe_train(step) if metrics is not None: last_metrics = metrics elapsed_learn_s = time.perf_counter() - t_l @@ -575,7 +678,7 @@ def _drain_stats(): else: # warmup or starved: avoid a hot spin. t_idle = time.perf_counter() - time.sleep(async_options.idle_sleep_s) + time.sleep(opts.idle_sleep_s) idle_ms = (time.perf_counter() - t_idle) * 1000.0 if ingested == 0: # No ring slot was available: learner is waiting for collector data. @@ -583,124 +686,65 @@ def _drain_stats(): else: # Data arrived, but replay/batch readiness still gated training. learner_gate_wait_samples_ms.append(idle_ms) - step = control.collector_steps - control.global_step = step - if step >= next_log: - _drain_stats() + # One compact payload per log window to the parent (panel + + # TensorBoard live there; see train.py). All ranks send. now = time.time() sps = (step - last_log_step) * num_envs / max(now - last_log_time, 1e-6) warming = step < agent.cfg.learning_starts metrics_log = {k: float(v) for k, v in last_metrics.items()} if (last_metrics and not warming) else None - updates = learner.update_idx - utd = updates / max(step, 1) - # Per-process timing (collector and learner run concurrently, so - # these do NOT sum to 100% like the sync panel): - # timing_ms[collect] — collector's avg ms per env-step batch (from queue) - # timing_ms[wait] — avg ring-backpressure wait per batch - # learn_ms — learner's avg ms per train call (one UTD - # execution, which may run several gradient - # updates); idle waits are not included - # learn_pct — fraction of learner wall-clock spent updating vs - # idle/starved (≈100% when GPU-bound, lower if the - # collector can't keep the buffer fed) - # Window means only: live-panel percentiles are noise at these - # sample counts (benchmarks own the tail statistics). learn_pct = 100.0 * t_learn_win / max(now - last_log_time, 1e-9) - collector_timing_ms = last_stats.get("timing_ms", {}) - collector_timing_detail_ms = { - key: value for key, value in collector_timing_ms.items() if key != "collect" - } - # Panel tree is per-process; the headline collect/learn means - # live on TrainingPanelStats. Every timing key is either a flat - # stage name or a dotted path (``sync.wait_writer``, - # ``env_step.physics``); nesting is rebuilt with one rule, and - # a stage's own total folds into its node. (The collector - # already reports at most one sub-stage level per stage.) - collector_items: dict[str, Any] = {} - for key, value in collector_timing_detail_ms.items(): - _nest_timing_path(collector_items, tuple(key.split(".")), value) - timing_groups = {"collector": collector_items} learner_items: dict[str, Any] = {} - drain_ms = _timing_mean(learner_drain_samples_ms) if learner_drain_samples_ms else 0.0 - ring_wait_ms = _timing_mean(learner_ring_wait_samples_ms) if learner_ring_wait_samples_ms else 0.0 - gate_wait_ms = _timing_mean(learner_gate_wait_samples_ms) if learner_gate_wait_samples_ms else 0.0 if learner_drain_samples_ms: - learner_items["drain"] = drain_ms + learner_items["drain"] = timing_mean(learner_drain_samples_ms) if learner_ring_wait_samples_ms: - learner_items["ring wait"] = ring_wait_ms + learner_items["ring wait"] = timing_mean(learner_ring_wait_samples_ms) if learner_gate_wait_samples_ms: - learner_items["gate wait"] = gate_wait_ms - update_items = {key: _timing_mean(values) for key, values in learner_breakdown_samples_ms.items()} + learner_items["gate wait"] = timing_mean(learner_gate_wait_samples_ms) + update_items = {key: timing_mean(values) for key, values in learner_breakdown_samples_ms.items()} if update_items: # publish is a child stage of the learner update in the # panel, so include it in the displayed update total too. if "publish" in update_items and "total" in update_items: update_items["total"] += update_items["publish"] learner_items["update"] = update_items - if learner_items: - timing_groups["learner"] = learner_items - learn_ms = _timing_mean(learner_train_samples_ms) if learner_train_samples_ms else 0.0 - stats = TrainingPanelStats( - iteration=step, - total_iterations=num_iterations, - steps_per_second=sps, - elapsed_seconds=now - start_time, - mean_return=last_stats["return"], - mean_episode_length=last_stats["ep_len"], - episodes=last_stats["episodes"], - buffer_size=agent.rb.num_stored * num_envs, - buffer_capacity=agent.rb.buffer_size * num_envs, - collect_ms=collector_timing_ms.get("collect", 0.0), - learn_ms=learn_ms, - learn_percent=learn_pct, - warming=warming, - training_metrics=metrics_log, - reward_terms=last_stats["reward_terms"], - env_metrics=last_stats["env_metrics"], - timing_groups=timing_groups, - diagnostics={"UTD": utd}, - cpu_load=cpu_sampler.sample(), - gpu_utilization_percent=gpu_sampler.sample(), - memory_usage=memory_sampler.sample(), - gpu_memory_usage=gpu_memory_sampler.sample(), - gpu_devices=sample_gpu_devices(gpu_sampler, gpu_memory_sampler), - checkpoint_path=last_checkpoint_path, - ) - emit_training_panel(live, stats, title=f"{env_name}/motrix.fastsac") - if writer is not None: - writer.add_scalar("rollout/mean_return", last_stats["return"], step) - writer.add_scalar("rollout/mean_ep_len", last_stats["ep_len"], step) - writer.add_scalar("perf/env_steps_per_s", sps, step) - writer.add_scalar( - "perf/updates_per_s", (updates - last_update_idx) / max(now - last_log_time, 1e-6), step - ) - 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", 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(): - writer.add_scalar(f"perf/collector_{k}_ms", v, step) - writer.add_scalar("perf/learn_ms_total", learn_ms, step) - writer.add_scalar("perf/learn_pct", learn_pct, step) - for k, v in last_stats["env_metrics"].items(): - writer.add_scalar(f"metrics/{k}", v, step) - for k, v in last_stats["reward_terms"].items(): - writer.add_scalar(f"reward/{k}", v, step) - if metrics_log is not None: - for k, v in metrics_log.items(): - writer.add_scalar(f"train/{k}", v, step) - last_log_time, last_log_step, last_update_idx = now, step, updates + payload = { + "rank": rank, + "step": step, + "sps": sps, + "updates": learner.update_idx, + "metrics": metrics_log, + "warming": warming, + "learn_ms": timing_mean(learner_train_samples_ms) if learner_train_samples_ms else 0.0, + "learn_pct": learn_pct, + "learner_timing": learner_items, + "buffer_size": agent.rb.num_stored * agent.num_envs, + "buffer_capacity": agent.rb.buffer_size * agent.num_envs, + # bare RingCursors (IPC transport) carry no size; the + # wrapped IpcTransitionRing lives in this worker. + "ring_fill": [None if isinstance(ring, RingCursors) else ring.size() for ring in all_rings], + "weight_version": max((tx.version for tx in weight_txs), default=None), + "checkpoint_path": last_checkpoint_path, + } + # Never block the training loop on a slow panel consumer: + # drop the window's payload if the queue is full (the next + # window's numbers supersede it anyway). + try: + panel_queue.put_nowait(payload) + except Full: + pass + last_log_time, last_log_step = now, step t_learn_win = 0.0 learner_train_samples_ms = [] learner_drain_samples_ms = [] learner_breakdown_samples_ms = {} learner_ring_wait_samples_ms = [] learner_gate_wait_samples_ms = [] - next_log += logging_interval + # Recompute the window edge (not +=) so the immediate first + # payload cannot shift every later window off the interval. + next_log = ((step // logging_interval) + 1) * logging_interval if logging_interval > 0 else 0 - if save_interval > 0 and step >= next_save and step > 0: + if is_primary and save_interval > 0 and step >= next_save and step > 0: agent.global_step = step learner.wait_ingest() # checkpoint reads the rb tensors path = Path(checkpoint_dir) / f"model_{step:07d}.pt" @@ -712,39 +756,43 @@ def _drain_stats(): checkpoints.TRAINING_STATE, checkpoint_format=checkpoint_format, ) - if console is not None: - last_checkpoint_path = str(path) - else: - print(f"[motrix.fastsac async] saved checkpoint {path}") + last_checkpoint_path = str(path) + logger.info("saved checkpoint %s", path) next_save += save_interval - # final checkpoint (identical structure to sync fastsac) - agent.global_step = control.collector_steps - learner.wait_ingest() # checkpoint reads the rb tensors - ckpt_path = checkpoints.final_checkpoint_path(checkpoint_format, Path(run_dir)) - ckpt_path.parent.mkdir(parents=True, exist_ok=True) - torch.save(agent.state_dict(), ckpt_path) - checkpoints.record_checkpoint_artifact( - Path(run_dir), - checkpoints.LATEST_TRAINING_STATE, - ckpt_path, - checkpoints.TRAINING_STATE, - checkpoint_format=checkpoint_format, - ) - checkpoints.record_checkpoint_artifact( - Path(run_dir), - checkpoints.BEST_POLICY, - ckpt_path, - checkpoints.POLICY, - checkpoint_format=checkpoint_format, - ) - (console.print if console else print)(f"[motrix.fastsac async] saved checkpoint to {ckpt_path}") - if writer is not None: - writer.close() + if num_learners > 1: + # All ranks have consumed the identical global step count, so the + # collectives inside update() are matched; this barrier aligns + # teardown so no rank destroys the process group while another is + # still inside a collective. + dist.barrier() + if is_primary: + # final checkpoint (identical structure to sync fastsac) + agent.global_step = control.collector_steps // num_collectors + learner.wait_ingest() # checkpoint reads the rb tensors + ckpt_path = checkpoints.final_checkpoint_path(checkpoint_format, Path(run_dir)) + ckpt_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(agent.state_dict(), ckpt_path) + checkpoints.record_checkpoint_artifact( + Path(run_dir), + checkpoints.LATEST_TRAINING_STATE, + ckpt_path, + checkpoints.TRAINING_STATE, + checkpoint_format=checkpoint_format, + ) + checkpoints.record_checkpoint_artifact( + Path(run_dir), + checkpoints.BEST_POLICY, + ckpt_path, + checkpoints.POLICY, + checkpoint_format=checkpoint_format, + ) + logger.info("saved checkpoint to %s", ckpt_path) except BaseException: - error_queue.put(("learner", traceback.format_exc())) + if not isinstance(sys.exc_info()[1], (KeyboardInterrupt, SystemExit)): + error_queue.put((f"learner[{rank}]" if num_learners > 1 else "learner", traceback.format_exc())) raise finally: - if live is not None: - live.stop() - control.set_stop() # tell the collector to exit + if num_learners > 1: + dist.destroy_process_group() + control.set_stop() # tell the collectors to exit diff --git a/motrix_rl/src/motrix_rl/fastsac/buffer.py b/motrix_rl/src/motrix_rl/fastsac/buffer.py index 58eaa4de..f7179b77 100644 --- a/motrix_rl/src/motrix_rl/fastsac/buffer.py +++ b/motrix_rl/src/motrix_rl/fastsac/buffer.py @@ -221,12 +221,27 @@ def sample(self, batch_size: int) -> dict: class EmpiricalNormalization(nn.Module): - """Normalize mean and variance of values based on empirical values.""" + """Normalize mean and variance of values based on empirical values. + + Also maintains optional LOCAL accumulators (``_local_*``) for the + multi-learner cross-rank merge: the public ``_mean``/``_var``/``count`` + may be overwritten by a merge (they then hold GLOBAL stats), so they + must never feed the next merge — that would double-count all history on + every sync (count doubles per merge until float32 overflows). The + accumulators stay disabled (``local_enabled=False``) until a learner + turns them on; single-learner runs never pay the mirrored-update cost. + Local accumulators are plain attributes, deliberately NOT registered + buffers — they are per-process bookkeeping, not checkpoint state. + """ def __init__(self, shape, device, eps=1e-2, until=None): super().__init__() self.eps = eps self.until = until + self.local_enabled = False + self._local_mean = torch.zeros(shape, dtype=torch.float64, device=device) + self._local_var = torch.ones(shape, dtype=torch.float64, device=device) + self._local_count = 0 self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0).to(device)) self.register_buffer("_var", torch.ones(shape).unsqueeze(0).to(device)) self.register_buffer("_std", torch.ones(shape).unsqueeze(0).to(device)) @@ -257,3 +272,69 @@ def update(self, x: torch.Tensor) -> None: self._var.copy_(big_m2 / new_count) self._std.copy_(self._var.sqrt()) self.count.copy_(new_count) + # Mirror into the LOCAL accumulators: the per-rank cumulative stats + # that cross-rank merges consume. The public _mean/_var/_count may be + # overwritten by a merge (they then hold GLOBAL stats), so they must + # never feed the next merge — that would double-count all history on + # every sync (count doubles per merge until float32 overflows). + if self.local_enabled: + b_mean = batch_mean.squeeze(0) + b_var = batch_var.squeeze(0) + l_count = self._local_count + batch_size + delta = b_mean - self._local_mean + self._local_mean += delta * (batch_size / l_count) + delta2 = b_mean - self._local_mean + l_m2 = ( + self._local_var * self._local_count + + b_var * batch_size + + delta2.pow(2) * (self._local_count * batch_size / l_count) + ) + self._local_var.copy_(l_m2 / l_count) + self._local_count = l_count + + def seed_local_accumulators(self) -> None: + """Enable local accumulation, seeded from the current public stats. + + Used on the first cross-rank merge (including after resume): seeding + from the loaded global stats means the first merge does not discard + pre-resume history (the shared history is counted once per rank, an + acceptable one-off bias that washes out as new data arrives). + """ + self.local_enabled = True + self._local_mean = self._mean.detach().clone().squeeze(0).double() + self._local_var = self._var.detach().clone().squeeze(0).double() + self._local_count = int(self.count) + + def local_sufficient_stats_flat(self) -> torch.Tensor: + """Pack the LOCAL accumulators as ``[count, sum, sumsq]`` (float64). + + Single-tensor layout so a cross-rank merge needs exactly one + collective. ``sumsq`` is reconstructed from var via + ``E[x^2] = var + mean^2`` — SUM is additive while mean/var are not. + """ + count = torch.tensor(float(self._local_count), dtype=torch.float64, device=self._mean.device) + total = count * self._local_mean + sumsq = self._local_var * count + total * total / count.clamp_min(1.0) + return torch.cat([count.reshape(1), total, sumsq]) + + def apply_global_sufficient_stats(self, flat: torch.Tensor) -> None: + """Restore merged ``[count, sum, sumsq]`` into the public buffers. + + Computes the global mean/var in float64 (``sumsq`` is a sum of + same-scale squares; float32 would catastrophically cancel), clamps + the sample count to ``until`` to match the update freeze, and writes + ``_mean``/``_var``/``_std``/``count``. The local accumulators are + untouched: they keep tracking rank-local samples only. + """ + d = (flat.numel() - 1) // 2 + n, s, q = flat[0], flat[1 : 1 + d], flat[1 + d :] + if n <= 0: + return + mean = (s / n).float() + var = ((q - s * s / n) / n).clamp_min_(0.0).float() + if self.until is not None: + n = torch.minimum(n, torch.tensor(float(self.until), dtype=n.dtype, device=n.device)) + self._mean.copy_(mean.unsqueeze(0)) + self._var.copy_(var.unsqueeze(0)) + self._std.copy_(var.sqrt().unsqueeze(0)) + self.count.copy_(n.long()) diff --git a/motrix_rl/src/motrix_rl/fastsac/config.py b/motrix_rl/src/motrix_rl/fastsac/config.py index e5dfcd35..9ec0e672 100644 --- a/motrix_rl/src/motrix_rl/fastsac/config.py +++ b/motrix_rl/src/motrix_rl/fastsac/config.py @@ -96,6 +96,18 @@ class FastSacAsyncOptionsCfg: # host transfer they avoid. weight_ipc: str = "auto" weight_ipc_min_bytes: int = 16 * 1024 * 1024 + # Number of collector processes; num_envs is split evenly across them. + num_collectors: int = MISSING + # Restrict each collector to this many CPUs taken from its NUMA node / + # affinity set (nodes are assigned automatically); None uses all of them. + cpus_per_collector: int | None = MISSING + # DDP data-parallel learner replicas, one process per GPU (single node). + # Requires num_collectors % num_learners == 0; 1 keeps the single-learner + # behavior byte-identical (no process group, no DDP wrap). + num_learners: int = MISSING + # One device per learner (len == num_learners); None replicates `device`. + # Learners bind to their GPU's NUMA node automatically. + learner_devices: list[str] | None = MISSING @dataclass diff --git a/motrix_rl/src/motrix_rl/system_metrics.py b/motrix_rl/src/motrix_rl/system_metrics.py index b840b4d9..c1d17c8e 100644 --- a/motrix_rl/src/motrix_rl/system_metrics.py +++ b/motrix_rl/src/motrix_rl/system_metrics.py @@ -38,6 +38,10 @@ class CpuLoad: iowait_percent: float steal_percent: float per_core_percent: tuple[float, ...] | None = None + # The logical CPU ids aligned 1:1 with ``per_core_percent`` (sorted). Under + # NUMA binding the sampler only sees the process's affinity, so positional + # labels would misreport which physical cores are shown. + per_core_ids: tuple[int, ...] | None = None # Static "model name" from /proc/cpuinfo (Linux); None where unavailable. model_name: str | None = None @@ -101,6 +105,7 @@ def sample(self) -> CpuLoad | None: iowait_percent=100.0 * iowait / total, steal_percent=100.0 * steal / total, per_core_percent=per_core, + per_core_ids=tuple(sorted(common_ids)), model_name=self._model_name, ) diff --git a/motrix_rl/tests/fastsac_async_mocks.py b/motrix_rl/tests/fastsac_async_mocks.py new file mode 100644 index 00000000..12a9d0c1 --- /dev/null +++ b/motrix_rl/tests/fastsac_async_mocks.py @@ -0,0 +1,37 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Shared white-box mocks for the async FastSAC learner's process-free tests. + +``make_mock_learner`` builds a ``Learner`` without running ``__init__`` (no +agent build, no CUDA): the drain path only needs the rings, a capturing +replay-buffer stub and the generation assembler. Used by +test_fastsac_async_multi.py and test_fastsac_pipeline_equivalence.py. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from motrix_rl.fastsac.async_impl.learner import GenerationAssembler, Learner +from motrix_rl.fastsac.async_impl.transport import Control + + +def make_mock_learner(rings: list, extends: list | None) -> Learner: + """A drain-capable Learner over ``rings``; ``extends`` captures extend_batch columns.""" + learner = Learner.__new__(Learner) + learner.agent = SimpleNamespace( + device=torch.device("cpu"), + rb=SimpleNamespace(extend_batch=lambda *c: extends.append([x.clone() for x in c])), + cfg=SimpleNamespace(learning_starts=0), + ) + learner.async_options = SimpleNamespace(max_ingest_per_iter=8) + learner.control = Control(num_collectors=len(rings)) + learner.rings = rings + learner.weights = [] + learner._pending = GenerationAssembler(num_rings=len(rings)) + learner._staging = None + learner._pending_copy = False + return learner diff --git a/motrix_rl/tests/test_fastsac_async_multi.py b/motrix_rl/tests/test_fastsac_async_multi.py new file mode 100644 index 00000000..046a871d --- /dev/null +++ b/motrix_rl/tests/test_fastsac_async_multi.py @@ -0,0 +1,556 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 + +"""Multi-collector topology unit tests: control aggregation, env sharding, +strict-UTD accounting across sharded slots, stats aggregation and NUMA +helper selection. Process-level behavior is covered by end-to-end training.""" + +import math +from types import SimpleNamespace + +import pytest +import torch +from fastsac_async_mocks import make_mock_learner + +from motrix_rl.fastsac.async_impl.learner import Learner +from motrix_rl.fastsac.async_impl.numa import parse_cpulist, select_collector_cpus +from motrix_rl.fastsac.async_impl.stats import aggregate_collector_stats +from motrix_rl.fastsac.async_impl.topology import ( + CollectorInfo, + LearnerInfo, + TrainerTopology, + resolve_learner_devices, + resolve_trainer_topology, + ring_transport_is_ipc, + split_num_envs, +) +from motrix_rl.fastsac.async_impl.train import Trainer +from motrix_rl.fastsac.async_impl.transport import Control, SharedTransitionRing + + +def test_control_aggregates_per_collector_steps() -> None: + control = Control(num_collectors=3) + control.inc_collector_steps(0) + control.inc_collector_steps(1) + control.inc_collector_steps(1) + control.inc_collector_steps(2) + + assert control.collector_steps == 4 + assert control.collector_steps_at(0) == 1 + assert control.collector_steps_at(1) == 2 + assert control.collector_steps_at(2) == 1 + + +def test_control_resume_restarts_each_collector_from_checkpointed_iteration() -> None: + control = Control(num_collectors=2) + control.resume_collector_steps(5) + + # v is in full-batch equivalents: each collector's own counter restarts at v + assert (control.collector_steps_at(0), control.collector_steps_at(1)) == (5, 5) + assert control.collector_steps == 10 + + +def test_single_collector_control_matches_previous_single_counter() -> None: + control = Control() + control.resume_collector_steps(7) + control.inc_collector_steps(0) + + assert control.collector_steps == 8 + + +def test_split_num_envs_requires_even_division() -> None: + assert split_num_envs(2048, 2) == [1024, 1024] + assert split_num_envs(2048, 1) == [2048] + with pytest.raises(ValueError, match="divide evenly"): + split_num_envs(2048, 3) + with pytest.raises(ValueError, match="num_collectors must be >= 1"): + split_num_envs(2048, 0) + + +def _strict_learner(num_updates: int) -> Learner: + learner = Learner.__new__(Learner) + learner.agent = SimpleNamespace(cfg=SimpleNamespace(num_updates=num_updates)) + learner.async_options = SimpleNamespace(utd_mode="strict") + learner.control = SimpleNamespace(collector_steps=0, num_collectors=1) + return learner + + +def _push_batch(ring, tag: float) -> None: + n = ring.num_envs + obs = torch.full((n, ring.obs.shape[-1]), tag) + actions = torch.zeros(n, ring.actions.shape[-1]) + rewards = torch.zeros(n) + dones = torch.zeros(n, dtype=torch.long) + pushed = ring.push(obs, obs, actions, rewards, dones, dones) + assert pushed + + +def test_drain_merges_shard_generations_into_full_batches() -> None: + rings = [SharedTransitionRing(4, 4, 3, 3, 2) for _ in range(2)] + extends: list = [] + learner = make_mock_learner(rings, extends) + + # only collector 0 produced generation 0: nothing is ingestible yet + _push_batch(rings[0], 1.0) + assert learner.drain() == 0 + assert extends == [] + + # collector 1 catches up: generation 0 merges into one full 8-env batch + _push_batch(rings[1], 2.0) + assert learner.drain() == 1 + assert len(extends) == 1 + merged_obs = extends[0][0] + # batched runs are slot-major: (count, num_envs, dim) + assert merged_obs.shape == (1, 8, 3) + # env blocks concatenate in collector order + assert merged_obs[0, :4].eq(1.0).all() + assert merged_obs[0, 4:].eq(2.0).all() + # both read cursors advanced past the consumed generation + assert all(ring.read_idx == 1 for ring in rings) + + +def test_drain_holds_generation_until_all_rings_deliver() -> None: + rings = [SharedTransitionRing(4, 4, 3, 3, 2) for _ in range(2)] + extends: list = [] + learner = make_mock_learner(rings, extends) + + _push_batch(rings[0], 1.0) + _push_batch(rings[0], 1.5) + _push_batch(rings[1], 2.0) + assert learner.drain() == 1 + # generation 1 waits for collector 1's second slot + assert learner._pending.next_gen == 1 + _push_batch(rings[1], 2.5) + assert learner.drain() == 1 + assert all(ring.read_idx == 2 for ring in rings) + + +def test_single_ring_drain_batches_multiple_slots() -> None: + ring = SharedTransitionRing(4, 4, 3, 3, 2) + extends: list = [] + learner = make_mock_learner([ring], extends) + + for tag in (1.0, 2.0, 3.0): + _push_batch(ring, tag) + assert learner.drain() == 3 + # one extend_batch with all three batches stacked along the slot axis + assert len(extends) == 1 + assert extends[0][0].shape == (3, 4, 3) + assert extends[0][0][0].eq(1.0).all() + assert extends[0][0][2].eq(3.0).all() + assert ring.read_idx == 3 + + +def _snapshot(collector_id: int, ret: float, episodes: int, lag: int, collect_ms: float) -> dict: + return { + "collector_id": collector_id, + "return": ret, + "ep_len": 100.0, + "episodes": episodes, + "reward_terms": {"alive": 0.5}, + "env_metrics": {}, + "policy_lag": lag, + "timing_ms": {"collect": collect_ms}, + } + + +def test_aggregate_collector_stats_merges_snapshots() -> None: + per_collector = {0: _snapshot(0, 10.0, 4, 1, 8.0), 1: _snapshot(1, 20.0, 6, 3, 12.0)} + stats = aggregate_collector_stats(per_collector) + + assert stats["return"] == pytest.approx(15.0) + assert stats["episodes"] == 10 + assert stats["policy_lag"] == 3 + assert stats["timing_ms"]["collect"] == pytest.approx(10.0) + + +def test_aggregate_collector_stats_skips_missing_collectors() -> None: + stats = aggregate_collector_stats({0: _snapshot(0, 10.0, 4, 1, 8.0), 1: None}) + + assert stats["return"] == pytest.approx(10.0) + assert stats["episodes"] == 4 + + +def test_aggregate_collector_stats_defaults_without_snapshots() -> None: + stats = aggregate_collector_stats({0: None, 1: None}) + + assert math.isnan(stats["return"]) + assert stats["episodes"] == 0 + assert stats["policy_lag"] == 0 + + +def test_extend_batch_matches_repeated_extend_including_wraparound() -> None: + """extend_batch writes exactly what successive extend calls would, including + a run that wraps the circular buffer end — n-step time adjacency depends on it.""" + from motrix_rl.fastsac.buffer import SimpleReplayBuffer + + def make(): + return SimpleReplayBuffer(n_env=4, buffer_size=5, n_obs=3, n_act=2, n_critic_obs=3, device="cpu") + + def batch(tag: float, count: int): + obs = torch.stack([torch.full((4, 3), tag) for _ in range(count)], dim=0) + actions = torch.zeros(count, 4, 2) + rewards = torch.tensor([[float(tag + j)] * 4 for j in range(count)]) + flags = torch.zeros(count, 4, dtype=torch.long) + return obs, obs.clone(), actions, rewards, flags, flags.clone() + + stepwise, batched = make(), make() + tags = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] # runs past buffer_size -> wraps + i = 0 + while i < len(tags): + count = 2 if i + 1 < len(tags) else 1 + b = batch(0.0, count) # obs tag overwritten per column below + obs = torch.stack([torch.full((4, 3), tags[i + j]) for j in range(count)], dim=0) + rewards = torch.tensor([[float(i + j)] * 4 for j in range(count)]) + batched.extend_batch(obs, obs.clone(), b[2], rewards, b[4], b[5]) + for j in range(count): + single = batch(tags[i + j], 1) + single = (single[0], single[1], single[2], torch.full((4, 1), float(i + j)), single[4], single[5]) + stepwise.extend(*(f[:, 0] for f in single)) + i += count + + for name in ("observations", "actions", "rewards", "dones"): + torch.testing.assert_close(getattr(stepwise, name), getattr(batched, name)) + assert batched.ptr == stepwise.ptr + assert batched.num_stored == 5 + + +def test_parse_cpulist() -> None: + assert parse_cpulist("0-3,8,10-11") == [0, 1, 2, 3, 8, 10, 11] + assert parse_cpulist(" 5 ") == [5] + assert parse_cpulist("") == [] + + +def test_select_collector_cpus_chunks_and_shares() -> None: + base = [0, 1, 2, 3, 4, 5] + + # without cpus_per_collector every collector shares the full set + assert select_collector_cpus(base, 0, None) == base + assert select_collector_cpus(base, 1, None) == base + + assert select_collector_cpus(base, 0, 3) == [0, 1, 2] + assert select_collector_cpus(base, 1, 3) == [3, 4, 5] + # over-subscription is a config error, not a silent empty binding + with pytest.raises(ValueError, match="no CPUs"): + select_collector_cpus(base, 2, 4) + + with pytest.raises(ValueError, match="positive"): + select_collector_cpus(base, 0, 0) + with pytest.raises(ValueError, match="empty"): + select_collector_cpus([], 0, None) + + +def test_topology_pairs_collectors_with_owning_learner(monkeypatch) -> None: + from motrix_rl.fastsac.async_impl import numa + + monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0, 1]) + monkeypatch.setattr(numa, "gpu_numa_node", lambda index: {0: 0, 1: 1}.get(index)) + cuda = lambda n: torch.device("cuda", n) # noqa: E731 + opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) + + # 2 learners x 2 collectors: each collector sits on its owner's GPU-local + # node — the pair never straddles a NUMA boundary + topo = resolve_trainer_topology( + 8, + 4, + 2, + [cuda(0), cuda(1)], + [cuda(0), cuda(0), cuda(1), cuda(1)], + cuda(0), + opts, + actor_param_numel=0, + ) + assert [learner.numa_node for learner in topo.learners] == [0, 1] + assert [collector.numa_node for collector in topo.collectors] == [0, 0, 1, 1] + assert topo.env_shards == [2, 2, 2, 2] + assert topo.collector_owner(3) == 1 + assert [collector.ring_ipc for collector in topo.collectors] == [True, True, True, True] + + # multi-collector single learner: every collector follows the one learner + topo = resolve_trainer_topology(8, 4, 1, [], [cuda(1)] * 4, cuda(1), opts, actor_param_numel=0) + assert [learner.numa_node for learner in topo.learners] == [1] + assert [collector.numa_node for collector in topo.collectors] == [1, 1, 1, 1] + + +def test_topology_chunks_cpus_by_node_local_ordinal(monkeypatch) -> None: + """Each node's collectors chunk that node's CPU list from offset 0. + + A global collector index against a per-node base would offset rank>=1 + collectors past their own node's CPUs (asymmetric bindings, or a + "leaves no CPUs" abort once the offset exceeds the node's list). + """ + from motrix_rl.fastsac.async_impl import numa + + node_cpus = {0: list(range(8)), 1: list(range(100, 108))} + monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0, 1]) + monkeypatch.setattr(numa, "gpu_numa_node", lambda index: {0: 0, 1: 1}.get(index)) + monkeypatch.setattr(numa, "numa_node_cpus", lambda node: node_cpus[node]) + cuda = lambda n: torch.device("cuda", n) # noqa: E731 + opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) + + topo = resolve_trainer_topology( + 8, + 4, + 2, + [cuda(0), cuda(1)], + [cuda(0), cuda(0), cuda(1), cuda(1)], + cuda(0), + opts, + actor_param_numel=0, + cpus_per_collector=4, + ) + + # node 0's collectors take node 0's CPUs from offset 0; node 1's + # collectors take node 1's CPUs from offset 0 (NOT from offset 2*4) + assert [c.cpus for c in topo.collectors] == [ + [0, 1, 2, 3], + [4, 5, 6, 7], + [100, 101, 102, 103], + [104, 105, 106, 107], + ] + + +def test_topology_unbound_without_gpu_locality(monkeypatch) -> None: + from motrix_rl.fastsac.async_impl import numa + + cuda = torch.device("cuda", 0) + cpu = torch.device("cpu") + opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) + + # single-node host: no binding anywhere + monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0]) + topo = resolve_trainer_topology(2, 2, 1, [], [cuda, cpu], cuda, opts, actor_param_numel=0) + assert [learner.numa_node for learner in topo.learners] == [None] + assert [collector.numa_node for collector in topo.collectors] == [None, None] + + # multi-node host, CPU learner: no binding anywhere + monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0, 1]) + topo = resolve_trainer_topology(2, 2, 1, [], [cpu, cpu], cpu, opts, actor_param_numel=0) + assert [learner.numa_node for learner in topo.learners] == [None] + assert [collector.numa_node for collector in topo.collectors] == [None, None] + + # multi-node host, unknown GPU locality: no binding anywhere + monkeypatch.setattr(numa, "gpu_numa_node", lambda index: None) + topo = resolve_trainer_topology(2, 2, 1, [], [cuda, cuda], cuda, opts, actor_param_numel=0) + assert [learner.numa_node for learner in topo.learners] == [None] + assert [collector.numa_node for collector in topo.collectors] == [None, None] + + # index-less cuda means device 0 + monkeypatch.setattr(numa, "gpu_numa_node", lambda index: {0: 0}.get(index)) + topo = resolve_trainer_topology(2, 2, 1, [], [cuda, cuda], torch.device("cuda"), opts, actor_param_numel=0) + assert [learner.numa_node for learner in topo.learners] == [0] + + +def test_topology_rejects_bad_counts() -> None: + opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) + with pytest.raises(ValueError, match="divide evenly"): + resolve_trainer_topology(3, 3, 2, [], [], torch.device("cpu"), opts, actor_param_numel=0) + with pytest.raises(ValueError, match="invalid worker counts"): + resolve_trainer_topology(4, 0, 1, [], [], torch.device("cpu"), opts, actor_param_numel=0) + with pytest.raises(ValueError, match="collector devices"): + resolve_trainer_topology(4, 2, 1, [], [torch.device("cpu")], torch.device("cpu"), opts, actor_param_numel=0) + + +def test_ring_slice_partitions_collectors_by_ownership() -> None: + rings = list(range(6)) + topo = TrainerTopology( + learners=[LearnerInfo(rank=r, device=torch.device("cpu"), numa_node=None, cpus=[]) for r in range(3)], + collectors=[CollectorInfo(i, i // 2, 1, torch.device("cpu"), None, [], False, False) for i in range(6)], + ) + + # collector i belongs to learner i // k, so each rank drains a contiguous block + assert topo.ring_slice_for_rank(rings, 0) == [0, 1] + assert topo.ring_slice_for_rank(rings, 1) == [2, 3] + assert topo.ring_slice_for_rank(rings, 2) == [4, 5] + # 1:1 topology: one ring per learner + two = TrainerTopology( + learners=[ + LearnerInfo(rank=0, device=torch.device("cpu"), numa_node=None, cpus=[]), + LearnerInfo(rank=1, device=torch.device("cpu"), numa_node=None, cpus=[]), + ], + collectors=[CollectorInfo(i, i, 1, torch.device("cpu"), None, [], False, False) for i in range(2)], + ) + assert two.ring_slice_for_rank(rings, 1) == [1] + # single learner owns everything + one = TrainerTopology( + learners=[LearnerInfo(rank=0, device=torch.device("cpu"), numa_node=None, cpus=[])], + collectors=[CollectorInfo(i, 0, 1, torch.device("cpu"), None, [], False, False) for i in range(6)], + ) + assert one.ring_slice_for_rank(rings, 0) == rings + + +def _lockstep_learner(num_updates: int, utd_mode: str, last_gstep: int = 5) -> tuple[Learner, list]: + learner = Learner.__new__(Learner) + learner.agent = SimpleNamespace( + cfg=SimpleNamespace(num_updates=num_updates), + rb=SimpleNamespace(num_stored=1), + device=torch.device("cpu"), + ) + learner.async_options = SimpleNamespace(utd_mode=utd_mode, idle_sleep_s=0.0) + learner.control = SimpleNamespace(collector_steps=100, num_collectors=2, stop=False) + learner.ddp_rank = 0 + learner._learning_starts = 10 + learner._staging = None + learner._last_train_gstep = last_gstep + learner._last_publish_ms = 0.0 + learner.weights = [] + learner.agent.rb.num_stored = 10 + calls: list = [] + learner.agent.update = lambda n: calls.append(n) or {} + learner.publish_if_due = lambda: None + return learner, calls + + +def test_lockstep_strict_due_scales_with_global_step_delta(monkeypatch) -> None: + import torch.distributed as dist + + monkeypatch.setattr(dist, "broadcast", lambda tensor, src=0: tensor) + learner, calls = _lockstep_learner(num_updates=4, utd_mode="strict") + + assert learner.maybe_train(8) is not None # delta 3 + assert sum(calls) == 12 and all(c == 4 for c in calls) + assert learner.maybe_train(8) is None # no new global batches -> no update + assert sum(calls) == 12 + assert learner.maybe_train(9) is not None + assert sum(calls) == 16 + + +def test_lockstep_learner_bound_runs_base_on_any_progress(monkeypatch) -> None: + import torch.distributed as dist + + monkeypatch.setattr(dist, "broadcast", lambda tensor, src=0: tensor) + learner, calls = _lockstep_learner(num_updates=4, utd_mode="learner_bound") + + assert learner.maybe_train(6) is not None + assert sum(calls) == 4 + assert learner.maybe_train(6) is None + assert sum(calls) == 4 + + +def test_lockstep_gated_before_learning_starts(monkeypatch) -> None: + import torch.distributed as dist + + monkeypatch.setattr(dist, "broadcast", lambda tensor, src=0: tensor) + learner, calls = _lockstep_learner(num_updates=4, utd_mode="strict") + learner.control.collector_steps = 19 # < learning_starts(10) * num_collectors(2) + + assert learner.maybe_train(8) is None + assert calls == [] + + +def _ipc_opts(mode: str = "auto") -> SimpleNamespace: + return SimpleNamespace(transition_ipc=mode) + + +def test_ring_transport_ipc_colocated_collectors() -> None: + """Collectors co-located with their owning learner's GPU get IPC rings per rank.""" + cuda = lambda n: torch.device("cuda", n) # noqa: E731 + + # multi-learner: each rank's collectors infer on its GPU -> all IPC + flags = ring_transport_is_ipc(_ipc_opts(), [cuda(0), cuda(1)], [cuda(0), cuda(0), cuda(1), cuda(1)], 4, 2, cuda(0)) + assert flags == [True, True, True, True] + + # a cross-GPU collector drags its whole rank back to host rings + flags = ring_transport_is_ipc(_ipc_opts(), [cuda(0), cuda(1)], [cuda(0), cuda(1), cuda(1), cuda(1)], 4, 2, cuda(0)) + assert flags == [False, False, True, True] + + # single learner, multi collector, all on its GPU -> IPC + flags = ring_transport_is_ipc(_ipc_opts(), [], [cuda(0)] * 4, 4, 1, cuda(0)) + assert flags == [True] * 4 + + # cpu collector inference -> host rings + flags = ring_transport_is_ipc(_ipc_opts(), [], [torch.device("cpu")] * 2, 2, 1, cuda(0)) + assert flags == [False, False] + + +def test_normalizer_stat_sync_merges_shards(monkeypatch) -> None: + """_sync_normalizer_stats merges per-rank (n, Sum, SumSq) into global stats.""" + import torch.distributed as dist + + from motrix_rl.fastsac.buffer import EmpiricalNormalization + + def make_norm(seed: int, data: torch.Tensor) -> EmpiricalNormalization: + norm = EmpiricalNormalization(data.shape[1:], torch.device("cpu")) + torch.manual_seed(seed) + norm.update(data) + return norm + + a = make_norm(0, torch.randn(1000, 5) * 2.0 + 1.0) + b = make_norm(1, torch.randn(300, 5) * 0.5 - 2.0) + + # the other rank's flat sufficient statistics (seeded from its publics) + b.seed_local_accumulators() + other = b.local_sufficient_stats_flat() + + def fake_all_reduce(flat, op=None): # noqa: ARG001 + flat += other + + a_count, a_mean, a_var = a.count.item(), a._mean.clone(), a._var.clone() + + monkeypatch.setattr(dist, "all_reduce", fake_all_reduce) + learner = Learner.__new__(Learner) + learner.agent = SimpleNamespace(world_size=2, obs_normalizer=a, critic_obs_normalizer=None) + learner._sync_normalizer_stats() + + # global empirical stats over both shards (Chan merge of sufficient stats) + n = a_count + b.count.item() + global_mean = (a_count * a_mean.squeeze(0) + b.count * b._mean.squeeze(0)) / n + torch.testing.assert_close(a._mean.squeeze(0), global_mean) + global_var = ( + a_count * a_var.squeeze(0) + + b.count * b._var.squeeze(0) + + a_count * b.count / n * (a_mean.squeeze(0) - b._mean.squeeze(0)) ** 2 + ) / n + torch.testing.assert_close(a._var.squeeze(0), global_var) + assert a.count == n + + +def test_normalizer_stat_sync_zero_count_is_noop(monkeypatch) -> None: + """A never-updated normalizer (count=0) must contribute zeros, not NaN.""" + import torch.distributed as dist + + from motrix_rl.fastsac.buffer import EmpiricalNormalization + + fresh = EmpiricalNormalization(5, torch.device("cpu")) # count=0 + other = EmpiricalNormalization(5, torch.device("cpu")) + other.update(torch.randn(300, 5)) + + other.seed_local_accumulators() + other_flat = other.local_sufficient_stats_flat() + monkeypatch.setattr(dist, "all_reduce", lambda flat, op=None: flat.__iadd__(other_flat)) + learner = Learner.__new__(Learner) + learner.agent = SimpleNamespace(world_size=2, obs_normalizer=fresh, critic_obs_normalizer=None) + learner._sync_normalizer_stats() + + assert torch.isfinite(fresh._mean).all() and torch.isfinite(fresh._std).all() + torch.testing.assert_close(fresh._mean.squeeze(0), other._mean.squeeze(0)) + torch.testing.assert_close(fresh._var.squeeze(0), other._var.squeeze(0)) + assert fresh.count == 300 + + +def _trainer_with_device(device: str | None) -> Trainer: + trainer = Trainer.__new__(Trainer) + trainer._rlcfg = SimpleNamespace(device=device) + return trainer + + +def test_resolve_learner_devices_expands_generic_cuda(monkeypatch) -> None: + monkeypatch.setattr("torch.cuda.device_count", lambda: 4) + + devices = resolve_learner_devices(None, 2, torch.device("cuda")) + assert [str(d) for d in devices] == ["cuda:0", "cuda:1"] + + devices = resolve_learner_devices(["cuda:3", "cuda:1"], 2, torch.device("cuda")) + assert [str(d) for d in devices] == ["cuda:3", "cuda:1"] + + +def test_resolve_learner_devices_rejects_bad_configs(monkeypatch) -> None: + monkeypatch.setattr("torch.cuda.device_count", lambda: 2) + + with pytest.raises(ValueError, match="exactly num_learners"): + resolve_learner_devices(["cuda:0"], 2, torch.device("cuda")) + with pytest.raises(ValueError, match="more than once"): + resolve_learner_devices(["cuda:0", "cuda:0"], 2, torch.device("cuda")) + with pytest.raises(ValueError, match="does not exist"): + resolve_learner_devices(["cuda:0", "cuda:5"], 2, torch.device("cuda")) + with pytest.raises(ValueError, match="requires CUDA"): + resolve_learner_devices(None, 2, torch.device("cpu")) diff --git a/motrix_rl/tests/test_fastsac_boot_panel.py b/motrix_rl/tests/test_fastsac_boot_panel.py new file mode 100644 index 00000000..60f1fc47 --- /dev/null +++ b/motrix_rl/tests/test_fastsac_boot_panel.py @@ -0,0 +1,167 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 + +"""Boot panel worker-log tail tests: geometry stability is the contract. + +The async trainer renders its startup panel with rich Live erase-and-redraw; +the log region must therefore keep a constant line count and line width (no +wrapping) and tolerate workers that have not created their log file yet. +""" + +from pathlib import Path + +from motrix_rl.fastsac.async_impl.panels import BootPanel, worker_log_tail + + +def test_tail_shows_placeholder_before_any_log_exists(tmp_path: Path) -> None: + result = worker_log_tail(tmp_path, ["collector0.log", "learner.log"]) + + assert result.plain == "(waiting for worker logs…)" + + +def test_tail_skips_missing_files_and_keeps_last_two_per_file(tmp_path: Path) -> None: + (tmp_path / "collector0.log").write_text("old line\nline1\nline2\n") + # learner.log deliberately absent — learner not spawned its log yet + + result = worker_log_tail(tmp_path, ["collector0.log", "learner.log"]) + + assert result.plain == "line1\nline2" + + +def test_tail_orders_files_in_given_order_newest_last(tmp_path: Path) -> None: + (tmp_path / "collector0.log").write_text("c-line\n") + (tmp_path / "learner.log").write_text("l-line\n") + + result = worker_log_tail(tmp_path, ["collector0.log", "learner.log"]) + + assert result.plain == "c-line\nl-line" + + +def test_tail_caps_total_lines_to_max_lines(tmp_path: Path) -> None: + (tmp_path / "collector0.log").write_text("a\nb\n") + (tmp_path / "learner.log").write_text("c\nd\n") + + result = worker_log_tail(tmp_path, ["collector0.log", "learner.log"], max_lines=2) + + # keeps the newest tail: the learner's two lines, collector0 dropped + assert result.plain == "c\nd" + + +def test_tail_truncates_long_lines_to_fixed_width(tmp_path: Path) -> None: + long_line = "x" * 500 + (tmp_path / "learner.log").write_text(long_line + "\n") + + result = worker_log_tail(tmp_path, ["learner.log"], line_width=80) + + line = result.plain.splitlines()[0] + assert len(line) == 80 + + +def test_tail_reads_only_the_end_of_large_files(tmp_path: Path) -> None: + filler = "\n".join(f"filler-{i}" for i in range(5000)) + (tmp_path / "learner.log").write_text(f"{filler}\nfinal-line\n") + + result = worker_log_tail(tmp_path, ["learner.log"]) + + assert "filler-0" not in result.plain + assert result.plain.endswith("final-line") + + +def test_tail_blank_lines_are_ignored(tmp_path: Path) -> None: + (tmp_path / "learner.log").write_text("\n\n \nreal-line\n\n") + + result = worker_log_tail(tmp_path, ["learner.log"]) + + assert result.plain == "real-line" + + +def test_tail_per_file_controls_lines_read_from_each_file(tmp_path: Path) -> None: + (tmp_path / "learner.log").write_text("l1\nl2\nl3\n") + + result = worker_log_tail(tmp_path, ["learner.log"], max_lines=3, per_file=3) + + assert result.plain == "l1\nl2\nl3" + + +def test_tail_splits_evenly_across_files_with_per_file(tmp_path: Path) -> None: + (tmp_path / "collector0.log").write_text("c1\nc2\nc3\n") + (tmp_path / "learner.log").write_text("l1\nl2\nl3\n") + + result = worker_log_tail(tmp_path, ["collector0.log", "learner.log"], max_lines=4, per_file=2) + + assert result.plain == "c2\nc3\nl2\nl3" + + +def test_tail_wrap_folds_long_lines_into_full_width_continuations(tmp_path: Path) -> None: + long_line = "x" * 100 + (tmp_path / "learner.log").write_text(long_line + "\n") + + result = worker_log_tail(tmp_path, ["learner.log"], max_lines=4, line_width=40, wrap=True) + + lines = result.plain.splitlines() + # 100 chars fold into ceil(100/40) = 3 full-width continuation lines + assert [len(ln) for ln in lines] == [40, 40, 20] + assert "".join(lines) == long_line + + +def test_tail_wrap_keeps_display_line_count_capped_at_max_lines(tmp_path: Path) -> None: + (tmp_path / "learner.log").write_text("y" * 200 + "\nshort\n") + + result = worker_log_tail(tmp_path, ["learner.log"], max_lines=2, line_width=40, wrap=True, per_file=2) + + lines = result.plain.splitlines() + # the long line folds into 5 display lines; the cap keeps only the last 2 + assert len(lines) == 2 + assert lines[-1].endswith("short") + + +def test_boot_panel_render_fits_terminal_height(tmp_path: Path, monkeypatch) -> None: + from rich.console import Console + + (tmp_path / "collector0.log").write_text("c-line\n") + monkeypatch.setattr("motrix_rl.fastsac.async_impl.panels.Console", lambda: Console(width=100, height=20)) + panel = BootPanel( + title="env/motrix.fastsac — worker startup", + log_dir=tmp_path, + log_names=["collector0.log", "learner.log"], + workers=[("collector", 0), ("learner", 0)], + ) + + probe = Console(width=100, height=20, force_terminal=True, file=open("/dev/null", "w")) + rendered = ["".join(seg.text for seg in row) for row in probe.render_lines(panel.render({("collector", 0)}))] + + # outer frame 2 lines; table region = 2 workers + 4 + gate 1 = 7; + # log region = 20 - 2 - 7 = 11; cell = 11 - 2 borders = 9 lines + assert panel.log_tail_lines == 9 + assert len(rendered) == 20 # full-height frame, never taller + assert any("ready" in ln for ln in rendered) + assert any("booting" in ln for ln in rendered) + assert any("c-line" in ln for ln in rendered) + + +def test_boot_panel_handoff_mode_renders_progress_and_gate(tmp_path: Path, monkeypatch) -> None: + from rich.console import Console + + (tmp_path / "collector0.log").write_text("c-line\n") + import motrix_rl.fastsac.async_impl.panels as panels + + devnull = open("/dev/null", "w") + monkeypatch.setattr(panels, "Console", lambda: Console(width=100, height=20, force_terminal=True, file=devnull)) + panel = panels.BootPanel( + "env/motrix.fastsac", + tmp_path, + ["collector0.log", "learner.log"], + [("collector", 0), ("learner", 0)], + ) + probe = Console(width=100, height=20, force_terminal=True, file=devnull) + boot = ["".join(s.text for s in row) for row in probe.render_lines(panel.render({("collector", 0)}))] + handoff = ["".join(s.text for s in row) for row in probe.render_lines(panel.render(set(), gate="g", starting=True))] + + # Boot phase: ready/booting statuses; the reserved gate line is blank. + boot_text = "\n".join(boot) + assert "ready" in boot_text and "booting" in boot_text and "c-line" in boot_text + assert "worker startup" not in boot_text # table title removed; task name is the frame title + # Handoff phase: same frame height, every worker "starting", gate line filled. + handoff_text = "\n".join(handoff) + assert len(handoff) == len(boot) == 20 + assert handoff_text.count("starting") == 2 + assert "│ g" in handoff_text # gate line filled inside the outer frame diff --git a/motrix_rl/tests/test_fastsac_collector.py b/motrix_rl/tests/test_fastsac_collector.py index a40d36d3..f24577a9 100644 --- a/motrix_rl/tests/test_fastsac_collector.py +++ b/motrix_rl/tests/test_fastsac_collector.py @@ -130,7 +130,6 @@ def _collector(device: str, *, compile: bool = False, amp: bool = False): env, cfg, _OBS_DIM, - _CRITIC_OBS_DIM, _ACT_DIM, source_actor.action_scale, source_actor.action_bias, @@ -232,7 +231,6 @@ def test_collector_reports_env_step_substage_timing() -> None: env, cfg, _OBS_DIM, - _CRITIC_OBS_DIM, _ACT_DIM, source_actor.action_scale, source_actor.action_bias, diff --git a/motrix_rl/tests/test_fastsac_ddp_equivalence.py b/motrix_rl/tests/test_fastsac_ddp_equivalence.py new file mode 100644 index 00000000..3713ffd3 --- /dev/null +++ b/motrix_rl/tests/test_fastsac_ddp_equivalence.py @@ -0,0 +1,170 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Layer-3 pipeline equivalence: the DDP update path must match single-process. + +Two ranks on gloo (CPU), each holding the SAME parameter clone and feeding +half of a fixed batch, run the real ``FastSacAgent`` update methods with +gradient averaging. A single-process reference runs the same methods on the +full batch. After every step the parameters must agree to float tolerance — +this pins down _allreduce_module_grads (parameter coverage, ordering), +log_alpha averaging and optimizer stepping, without any environment. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.multiprocessing as mp + +OBS, COBS, ACT = 6, 4, 2 + + +def _make_agent(world_size: int, seed: int) -> object: + from pathlib import Path + + from omegaconf import OmegaConf + + from motrix_rl.fastsac.agent import FastSacAgent + from motrix_rl.fastsac.config import FastSacAgentCfg + + yaml_path = Path(__file__).resolve().parents[2] / "configs" / "algo_base" / "motrix.fastsac.yaml" + base = OmegaConf.merge(OmegaConf.structured(FastSacAgentCfg), OmegaConf.load(yaml_path)["agent"]) + cfg = OmegaConf.to_object( + OmegaConf.merge( + base, + { + "actor_hidden_dim": 32, + "critic_hidden_dim": 32, + "num_atoms": 11, + "compile": False, + "learning_starts": 0, + }, + ) + ) + torch.manual_seed(seed) + agent = FastSacAgent( + obs_dim=OBS, + critic_obs_dim=COBS, + act_dim=ACT, + num_envs=8, + cfg=cfg, + device=torch.device("cpu"), + action_scale=torch.ones(ACT), + action_bias=torch.zeros(ACT), + world_size=world_size, + ) + return agent + + +def _fixed_batch(rows: int, seed: int) -> dict: + g = torch.Generator().manual_seed(seed) + obs = torch.randn(rows, OBS, generator=g) + return { + "obs": obs, + "next_obs": obs + 0.1 * torch.randn(rows, OBS, generator=g), + "critic_obs": torch.randn(rows, COBS, generator=g), + "next_critic_obs": torch.randn(rows, COBS, generator=g), + "actions": torch.randn(rows, ACT, generator=g), + "rewards": torch.randn(rows, generator=g), + "dones": (torch.rand(rows, generator=g) < 0.2).long(), + "truncations": (torch.rand(rows, generator=g) < 0.1).long(), + "effective_n_steps": torch.ones(rows, dtype=torch.long), + } + + +def _worker(rank: int, result_queue, seeds) -> None: + import torch.distributed as dist + + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(seeds["port"])) + dist.init_process_group("gloo", rank=rank, world_size=2) + import traceback + + # gloo has no ReduceOp.AVG — emulate it (NCCL supports it in production) + _orig_all_reduce = dist.all_reduce + + def _avg_all_reduce(tensor, op=None): + if op == dist.ReduceOp.AVG: + _orig_all_reduce(tensor, op=dist.ReduceOp.SUM) + tensor.div_(2) + return tensor + return _orig_all_reduce(tensor, op=op) + + dist.all_reduce = _avg_all_reduce + + try: + agent = _make_agent(2, seeds["init"]) + half = _fixed_batch(64, seeds["batch"]) + # rank r sees half r of the SAME global batch + half = {k: v[rank * 32 : (rank + 1) * 32] if v.shape[0] == 64 else v for k, v in half.items()} + rows = 32 + for step in range(seeds["steps"]): + # Align the RNG streams so the union of both ranks' draws equals + # the single-process draw: same seed everywhere, rank 1 skips + # rank 0's worth of normals first (CPU normal_ is a sequential + # stream, so half+half == full). + torch.manual_seed(10_000 * (step + 1)) + if rank == 1: + torch.randn(rows, ACT) # discard rank 0's draw + agent._update_main(half) + if step % 2 == 0: + torch.manual_seed(10_000 * (step + 1) + 1) + if rank == 1: + torch.randn(rows, ACT) + agent._update_pol(half) + import io + + buffer = io.BytesIO() + torch.save({k: v for k, v in agent.state_dict().items() if "optimizer" not in k}, buffer) + if rank == 0: + result_queue.put(buffer.getvalue()) + except Exception: + if rank == 0: + result_queue.put({"__error__": traceback.format_exc()}) + finally: + dist.destroy_process_group() + + +def test_ddp_update_matches_single_process() -> None: + """2-rank half-batch + grad AVG == 1-rank full batch, per update step.""" + seeds = {"init": 1234, "batch": 42, "steps": 3, "port": 29711} + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + procs = [ctx.Process(target=_worker, args=(r, result_queue, seeds)) for r in range(2)] + for p in procs: + p.start() + import io + + payload = result_queue.get(timeout=120) + if isinstance(payload, dict) and "__error__" in payload: + pytest.fail("DDP worker crashed:\n" + payload["__error__"]) + ddp_state = torch.load(io.BytesIO(payload), weights_only=False) + for p in procs: + p.join(timeout=60) + + single = _make_agent(1, seeds["init"]) + full = _fixed_batch(64, seeds["batch"]) + for step in range(seeds["steps"]): + torch.manual_seed(10_000 * (step + 1)) + single._update_main(full) + if step % 2 == 0: + torch.manual_seed(10_000 * (step + 1) + 1) + single._update_pol(full) + ref_state = {k: v for k, v in single.state_dict().items() if "optimizer" not in k} + + def _walk(prefix, a, b): + if torch.is_tensor(a): + if a.is_floating_point(): + torch.testing.assert_close(a, b, rtol=2e-4, atol=2e-4, msg=lambda m: f"{prefix}: {m}") + else: + assert torch.equal(a, b), prefix + elif isinstance(a, dict): + for k in a: + _walk(f"{prefix}.{k}", a[k], b[k]) + else: + assert a == b, prefix + + for key in ref_state: + _walk(key, ddp_state[key], ref_state[key]) diff --git a/motrix_rl/tests/test_fastsac_ipc_ring.py b/motrix_rl/tests/test_fastsac_ipc_ring.py index 6ad61e26..43d97cad 100644 --- a/motrix_rl/tests/test_fastsac_ipc_ring.py +++ b/motrix_rl/tests/test_fastsac_ipc_ring.py @@ -176,7 +176,7 @@ def test_ipc_and_host_ring_produce_identical_replay_buffers(): def test_use_ipc_transition_ring_gating(): - from motrix_rl.fastsac.async_impl.worker import use_ipc_transition_ring + from motrix_rl.fastsac.async_impl.topology import use_ipc_transition_ring def opts(mode): return SimpleNamespace(transition_ipc=mode) @@ -259,8 +259,12 @@ def test_learner_drains_ipc_ring_end_to_end(): async_options=SimpleNamespace(max_ingest_per_iter=CAPACITY, utd_mode="strict", weight_publish_interval=1) ) ) - learner = Learner(agent, cfg, receiver, SimpleNamespace(publish=lambda *a, **k: None), SimpleNamespace()) - assert learner._staging is None and learner._device_ring + from motrix_rl.fastsac.async_impl.learner import CollectorEndpoint + + stub_sender = SimpleNamespace(publish=lambda *a, **k: None) + control = SimpleNamespace(num_collectors=1, stop=False) + learner = Learner(agent, cfg, [CollectorEndpoint(ring=receiver, weight_sender=stub_sender)], control) + assert learner._staging is None # pure-IPC rank: no staging pipeline total = 40 # > rb cap 33 -> exercises the rb wrap batches = [_batch(t, seed=5) for t in range(total)] @@ -285,3 +289,103 @@ def test_learner_drains_ipc_ring_end_to_end(): torch.testing.assert_close(agent.rb.rewards[:, s].cpu(), batches[t][3]) torch.testing.assert_close(agent.rb.dones[:, s].cpu(), batches[t][4]) torch.testing.assert_close(agent.rb.truncations[:, s].cpu(), batches[t][5]) + + +@cuda_only +def test_learner_drains_mixed_host_and_ipc_rings(): + """A mixed rank (one host ring + one IPC ring) merges generations across both. + + The rank's drain pipeline pulls the IPC shard to host at assembly, so the + merged batch is assembled on a single device: env block [0:N_ENV] comes + from the host ring, [N_ENV:] from the IPC ring. + """ + from motrix_rl.fastsac.agent import FastSacAgent + from motrix_rl.fastsac.async_impl.learner import CollectorEndpoint, Learner + + total, merged_envs = 10, 2 * N_ENV + agent = FastSacAgent( + obs_dim=OBS, + critic_obs_dim=CRI, + act_dim=ACT, + num_envs=merged_envs, + cfg=SimpleNamespace( + actor_hidden_dim=16, + critic_hidden_dim=16, + num_q_networks=2, + actor_learning_rate=1e-3, + critic_learning_rate=1e-3, + alpha_learning_rate=1e-3, + weight_decay=0.0, + max_grad_norm=0.0, + use_layer_norm=False, + use_tanh=True, + log_std_max=0.0, + log_std_min=-5.0, + num_atoms=5, + v_min=-20.0, + v_max=20.0, + gamma=0.97, + tau=0.125, + alpha_init=0.001, + use_autotune=False, + target_entropy_ratio=0.0, + buffer_size=64, + num_steps=1, + batch_size=4, + learning_starts=1, + policy_frequency=4, + num_updates=1, + obs_normalization=False, + compile=False, + amp=False, + amp_dtype="bf16", + device=None, + ), + device=torch.device("cuda"), + ) + host_ring = SharedTransitionRing(CAPACITY, N_ENV, OBS, CRI, ACT) + owner, receiver = _owner_receiver() + stub = SimpleNamespace(publish=lambda *a, **k: None) + control = SimpleNamespace(num_collectors=2, stop=False) + cfg = SimpleNamespace( + trainer=SimpleNamespace( + async_options=SimpleNamespace(max_ingest_per_iter=CAPACITY, utd_mode="strict", weight_publish_interval=1) + ) + ) + learner = Learner( + agent, + cfg, + [CollectorEndpoint(ring=host_ring, weight_sender=stub), CollectorEndpoint(ring=receiver, weight_sender=stub)], + control, + ) + # the mixed rank lifts host shards to the device at assembly, so the + # merge runs D2D — no staging pipeline involved + assert learner._staging is None # mixed rank: device-side merge + + batches = [(_batch(t, seed=9), _batch(t, seed=500 + t)) for t in range(total)] + for t, (host_fields, ipc_fields) in enumerate(batches): + while not host_ring.push(*host_fields): + learner.drain() + while not owner.push(*ipc_fields): + owner.size() + learner.drain() + owner.size() + learner.drain() + + torch.cuda.synchronize() + while host_ring.has_next() or receiver.has_next(): + learner.drain() + learner.wait_ingest() + torch.cuda.synchronize() + + assert agent.rb.ptr == total + for t, (host_fields, ipc_fields) in enumerate(batches): + s = t % 64 + torch.testing.assert_close(agent.rb.observations[:N_ENV, s].cpu(), host_fields[0]) + torch.testing.assert_close(agent.rb.observations[N_ENV:, s].cpu(), ipc_fields[0]) + torch.testing.assert_close(agent.rb.critic_observations[:N_ENV, s].cpu(), host_fields[1]) + torch.testing.assert_close(agent.rb.critic_observations[N_ENV:, s].cpu(), ipc_fields[1]) + torch.testing.assert_close(agent.rb.actions[:N_ENV, s].cpu(), host_fields[2]) + torch.testing.assert_close(agent.rb.actions[N_ENV:, s].cpu(), ipc_fields[2]) + torch.testing.assert_close(agent.rb.rewards[:N_ENV, s].cpu(), host_fields[3]) + torch.testing.assert_close(agent.rb.rewards[N_ENV:, s].cpu(), ipc_fields[3]) diff --git a/motrix_rl/tests/test_fastsac_learner.py b/motrix_rl/tests/test_fastsac_learner.py index 9e292d8c..c895ca49 100644 --- a/motrix_rl/tests/test_fastsac_learner.py +++ b/motrix_rl/tests/test_fastsac_learner.py @@ -16,18 +16,6 @@ def _make_learner(utd_mode: str, num_updates: int) -> Learner: return learner -def test_strict_scales_num_updates_by_ingested_batches() -> None: - learner = _make_learner("strict", 4) - assert learner._num_updates_for(3) == 12 - assert learner._num_updates_for(0) == 0 - - -def test_learner_bound_runs_full_batch() -> None: - learner = _make_learner("learner_bound", 4) - assert learner._num_updates_for(0) == 4 - assert learner._num_updates_for(2) == 4 - - def test_own_copies_compiled_outputs_out_of_the_graph_pool() -> None: """Regression: metrics read at log time came from an invalidated CUDA graph. @@ -166,22 +154,22 @@ def test_update_metrics_survive_later_graph_generations() -> None: def test_nest_timing_path_merges_scalar_total_with_children_in_any_order() -> None: - from motrix_rl.fastsac.async_impl.worker import _nest_timing_path + from motrix_rl.fastsac.async_impl.stats import nest_timing_path # the collector emits a stage's scalar before its dotted sub-stages; the # rebuild must fold it into a "total" instead of nesting under a float tree: dict = {} - _nest_timing_path(tree, ("physics",), 15.0) - _nest_timing_path(tree, ("physics", "read"), 12.0) + nest_timing_path(tree, ("physics",), 15.0) + nest_timing_path(tree, ("physics", "read"), 12.0) assert tree == {"physics": {"total": 15.0, "read": 12.0}} # reverse order must converge to the same tree tree = {} - _nest_timing_path(tree, ("physics", "read"), 12.0) - _nest_timing_path(tree, ("physics",), 15.0) + nest_timing_path(tree, ("physics", "read"), 12.0) + nest_timing_path(tree, ("physics",), 15.0) assert tree == {"physics": {"total": 15.0, "read": 12.0}} # stages without sub-stages stay scalar tree = {} - _nest_timing_path(tree, ("apply_action",), 1.5) + nest_timing_path(tree, ("apply_action",), 1.5) assert tree == {"apply_action": 1.5} diff --git a/motrix_rl/tests/test_fastsac_pipeline_equivalence.py b/motrix_rl/tests/test_fastsac_pipeline_equivalence.py new file mode 100644 index 00000000..a053af11 --- /dev/null +++ b/motrix_rl/tests/test_fastsac_pipeline_equivalence.py @@ -0,0 +1,185 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Pipeline-equivalence tests: 2c/2l must be indistinguishable from 1c/1l. + +Three layers, in decreasing strictness: + +1. ring/generation-merge: the same transitions fed through two shard rings + (multi-collector path) must land in the replay buffer byte-identically to + the single-ring path — env-major order, n-step adjacency included; +2. normalizer statistics: per-rank update + cross-rank stat sync, cycled, + must always equal one global normalizer that saw ALL the data (this is + the exact code path that produced NaN on the cluster); +3. DDP update averaging: two ranks on half batches + gradient AVG must match + one rank on the full batch (CPU gloo, few steps) — covered by + test_fastsac_ddp_equivalence.py, which needs real process spawning. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from fastsac_async_mocks import make_mock_learner + +from motrix_rl.fastsac.async_impl.learner import Learner +from motrix_rl.fastsac.async_impl.transport import SharedTransitionRing +from motrix_rl.fastsac.buffer import EmpiricalNormalization, SimpleReplayBuffer + +_OBS, _COBS, _ACT = 7, 5, 3 + + +def _batch(envs: int, tag: float, seed: int) -> tuple: + g = torch.Generator().manual_seed(seed) + obs = torch.randn(envs, _OBS, generator=g) + tag + critic_obs = torch.randn(envs, _COBS, generator=g) + actions = torch.randn(envs, _ACT, generator=g) + rewards = torch.randn(envs, generator=g) + dones = (torch.rand(envs, generator=g) < 0.1).long() + truncations = (torch.rand(envs, generator=g) < 0.1).long() + return obs, critic_obs, actions, rewards, dones, truncations + + +_make_learner = make_mock_learner + + +# --------------------------------------------------------------------- layer 1 +def test_generation_merge_matches_single_ring_byte_for_byte() -> None: + """Same 8 batches: one ring drained directly vs two shard rings merged.""" + envs, shards, batches = 8, 2, 8 + whole = [_batch(envs, tag=float(t), seed=t) for t in range(batches)] + shard = lambda fields, r: tuple(f[r * (envs // shards) : (r + 1) * (envs // shards)] for f in fields) # noqa: E731 + + single_extends: list = [] + single = _make_learner([SharedTransitionRing(16, envs, _OBS, _COBS, _ACT)], single_extends) + for fields in whole: + assert single.rings[0].push(*fields) + single.drain() + + dual_extends: list = [] + rings = [SharedTransitionRing(16, envs // shards, _OBS, _COBS, _ACT) for _ in range(shards)] + dual = _make_learner(rings, dual_extends) + # generation k: both collectors push their shard of batch k (env blocks + # concatenate in collector order, matching the single-ring env order) + for fields in whole: + for r in range(shards): + assert rings[r].push(*shard(fields, r)) + while dual.drain() == 0: + pass + + assert len(single_extends) == len(dual_extends) + for a, b in zip(single_extends, dual_extends): + for x, y in zip(a, b): + torch.testing.assert_close(x, y) + # read cursors must advance by exactly the number of drained slots — a + # desync here stalls the generation merge permanently (the learner waits + # for a generation whose slots it already consumed) + assert single.rings[0].read_idx == batches + assert all(r.read_idx == batches for r in rings) + + +def test_generation_merge_preserves_n_step_adjacency() -> None: + """Per-env trajectories in the merged buffer match the single-ring path.""" + envs, shards, batches = 8, 2, 6 + + def make_rb(): + return SimpleReplayBuffer(n_env=envs, buffer_size=64, n_obs=_OBS, n_act=_ACT, n_critic_obs=_COBS, device="cpu") + + single_rb, dual_rb = make_rb(), make_rb() + single = _make_learner([SharedTransitionRing(16, envs, _OBS, _COBS, _ACT)], None) + single.agent.rb = single_rb + rings = [SharedTransitionRing(16, envs // shards, _OBS, _COBS, _ACT) for _ in range(shards)] + dual = _make_learner(rings, None) + dual.agent.rb = dual_rb + + whole = [_batch(envs, tag=float(t), seed=100 + t) for t in range(batches)] + for fields in whole: + assert single.rings[0].push(*fields) + for r in range(shards): + assert rings[r].push(*tuple(f[r * (envs // shards) : (r + 1) * (envs // shards)] for f in fields)) + single.drain() + while dual.drain() == 0: + pass + + for name in ("observations", "actions", "rewards", "dones", "truncations"): + torch.testing.assert_close(getattr(single_rb, name), getattr(dual_rb, name)) + assert single_rb.num_stored == dual_rb.num_stored + + +# --------------------------------------------------------------------- layer 2 +def _fake_dist_sum(flats: list[torch.Tensor]): + """Simulate a SUM all-reduce/broadcast over the given per-rank tensors.""" + + class _Dist: + def all_reduce(self, tensor, op=None): # noqa: ARG002 + total = sum(flats) + tensor.copy_(total) + return tensor + + def broadcast(self, tensor, src=0): # noqa: ARG002 + tensor.copy_(flats[0]) + return tensor + + return _Dist() + + +def test_normalizer_sync_cycles_equal_global_reference() -> None: + """Cycles of (per-shard update -> cross-rank sync) == one global normalizer.""" + torch.manual_seed(0) + dims, cycles, batch = 9, 50, 512 + ranks = [EmpiricalNormalization(dims, torch.device("cpu")) for _ in range(2)] + reference = EmpiricalNormalization(dims, torch.device("cpu")) + + learner = Learner.__new__(Learner) + learner.agent = SimpleNamespace(world_size=2, obs_normalizer=ranks[0], critic_obs_normalizer=ranks[1]) + learner.ddp_rank = 0 + + for cycle in range(cycles): + shards = [torch.randn(batch, dims) * (3.0 + cycle * 0.1) + cycle for _ in ranks] + for norm, shard in zip(ranks, shards): + if not norm.local_enabled: + norm.seed_local_accumulators() + norm.update(shard) + reference.update(torch.cat(shards)) + + # simulate the collective: both ranks run the SAME sync math on the + # SUM of their sufficient statistics + import torch.distributed as dist + + packs = [norm.local_sufficient_stats_flat() for norm in ranks] + + fake = _fake_dist_sum(packs) + orig_all_reduce, orig_broadcast = dist.all_reduce, dist.broadcast + dist.all_reduce, dist.broadcast = fake.all_reduce, fake.broadcast + try: + # run the real sync math on both ranks (each sees the SUM) + for norm in ranks: + learner.agent.obs_normalizer, learner.agent.critic_obs_normalizer = norm, None + learner._sync_normalizer_stats() + finally: + dist.all_reduce, dist.broadcast = orig_all_reduce, orig_broadcast + learner.agent.obs_normalizer, learner.agent.critic_obs_normalizer = ranks[0], ranks[1] + + for norm in ranks: + assert torch.isfinite(norm._std).all(), f"cycle {cycle}: std went NaN" + # float32 publics vs the reference's float32 incremental accumulation drift + # over dozens of cycles; the merged path keeps locals in float64 + torch.testing.assert_close(norm._mean, reference._mean, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(norm._std, reference._std, rtol=1e-2, atol=1e-2) + assert ranks[0].count == reference.count + + # zero-count start (the first-publish state) must stay finite too: both + # ranks at count=0 -> merged n=0 -> the sync leaves init stats untouched. + import torch.distributed as dist + + fresh = [EmpiricalNormalization(dims, torch.device("cpu")) for _ in range(2)] + learner.agent.obs_normalizer, learner.agent.critic_obs_normalizer = fresh[0], fresh[1] + zero = torch.zeros(1 + 2 * dims, dtype=torch.float64) + orig = dist.all_reduce + dist.all_reduce = lambda tensor, op=None: tensor.copy_(zero) # noqa: ARG005 + try: + learner._sync_normalizer_stats() + finally: + dist.all_reduce = orig + assert torch.isfinite(fresh[0]._std).all() and fresh[0].count == 0 diff --git a/motrix_rl/tests/test_rl_sim_backend.py b/motrix_rl/tests/test_rl_sim_backend.py index 353a11f3..0a1b6a10 100644 --- a/motrix_rl/tests/test_rl_sim_backend.py +++ b/motrix_rl/tests/test_rl_sim_backend.py @@ -19,7 +19,7 @@ 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.transport import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.transport import Control, SharedTransitionRing, StartupHandshake from motrix_rl.fastsac.async_impl.transport.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 @@ -249,7 +249,7 @@ def _async_cfg(): ) -def _collect_in_spawn(sim_backend: str) -> tuple[torch.Tensor, ...]: +def _collect_in_spawn(sim_backend: str, tmp_path) -> tuple[torch.Tensor, ...]: cfg = _async_cfg() dims = (_OBS_DIM, _OBS_DIM, _ACT_DIM) action_scale = torch.ones(_ACT_DIM) @@ -260,17 +260,36 @@ def _collect_in_spawn(sim_backend: str) -> tuple[torch.Tensor, ...]: 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)) - # Ship the real handshake message shape: (weight slots, ring slots). The - # host SharedTransitionRing needs no ring slots, so the ring field is None. - slot_queue.put((weight_tx.params, None)) # ship before the collector process starts + # Ship the real handshake messages: weight slots (publisher learner) and + # ring slots (draining learner; None for the host SharedTransitionRing). + handshake = StartupHandshake(ctx, num_collectors=1, barrier=False) + handshake.ship_weight_slots(0, weight_tx.params) + handshake.ship_ring_slots(0, None) + 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, slot_queue), + kwargs={ + "env_spec": env_spec, + "cfg": cfg, + "num_envs": _NUM_ENVS, + "dims": dims, + "action_scale": action_scale, + "action_bias": action_bias, + "ring": ring, + "weights": weights, + "control": control, + "stats_queue": stats_queue, + "error_queue": error_queue, + "num_iterations": 1, + "logging_interval": 1, + "is_resume": False, + "seed": 7, + "run_dir": str(tmp_path), + "handshake": handshake, + }, ) process.start() @@ -300,9 +319,9 @@ def _collect_in_spawn(sim_backend: str) -> tuple[torch.Tensor, ...]: platform.machine().lower() not in {"amd64", "x86_64"}, reason="FastSAC async shared memory currently supports x86-64 only", ) -def test_fastsac_async_supports_sim_backends() -> None: - np_slot = _collect_in_spawn("np") - torch_slot = _collect_in_spawn("torch") +def test_fastsac_async_supports_sim_backends(tmp_path) -> None: + np_slot = _collect_in_spawn("np", tmp_path) + torch_slot = _collect_in_spawn("torch", tmp_path) assert len(np_slot) == len(torch_slot) for np_tensor, torch_tensor in zip(np_slot, torch_slot): diff --git a/wiki/design/fastsac-async-heterogeneous-trainer.md b/wiki/design/fastsac-async-heterogeneous-trainer.md index 45f64d94..a23ed43d 100644 --- a/wiki/design/fastsac-async-heterogeneous-trainer.md +++ b/wiki/design/fastsac-async-heterogeneous-trainer.md @@ -2,7 +2,7 @@ ## 摘要 -异构 FastSAC 训练器把**仿真采样（collector）**与**网络训练（learner）**拆到两个进程，通过共享内存交换 transition 与权重，使 CPU 物理仿真与 GPU 梯度计算重叠，消除同步实现里「采样 → 训练 → 采样」串行循环中的 GPU 空转。 +异构 FastSAC 训练器把**仿真采样（collector）**与**网络训练（learner）**拆到两个进程，通过共享内存交换 transition 与权重，使 CPU 物理仿真与 GPU 梯度计算重叠，消除同步实现里「采样 → 训练 → 采样」串行循环中的 GPU 空转。多 NUMA node 服务器上可配置多个 collector 进程（每 node 一个），拓扑见 §8；本节先描述默认的 1 collector × 1 learner。 同步与异步执行共用 `motrix` framework 下唯一的 `fastsac` provider，对外方法名统一为 `motrix.fastsac`。`algo.asynchronous` 只选择执行拓扑，不改变算法、配置类型、run 身份或 checkpoint 格式。算法本身（`Actor`/`Critic`/`SimpleReplayBuffer`/`EmpiricalNormalization`/`FastSacAgent`）**原样复用、逐字节一致**。异构执行是默认模式，首要目标是让采样与训练各自满速。 @@ -78,6 +78,7 @@ motrix_rl/src/motrix_rl/fastsac/ │ └── train.py # 同步 Trainer └── async_impl/ ├── shm.py # 共享内存原语：SharedTransitionRing / WeightSnapshot / Control + ├── numa.py # 多 collector 的 NUMA/CPU 绑定（sched_setaffinity + libnuma set_membind） ├── collector.py # Collector：CPU 采样进程逻辑 ├── learner.py # Learner：GPU 训练进程逻辑 + UTD 治理 ├── worker.py # module-level 进程入口（可被 spawn pickle）+ 共享 builder @@ -179,7 +180,7 @@ producer lifetime 和 compiled collector 固定参数地址，不能只把 H2D ### 4.3 Control — 共享标量 -一小组共享 int64：`stop`（停止标志）、`global_step`（learner 迭代计数）、`collector_steps`（已产出的 env-step 批数，即训练进度基准）。seed 不放这里，作为进程入口参数直接传入。 +一小组共享 int64：`stop`（停止标志）、`global_step`（learner 迭代计数）、`collector_steps`（已产出的 env-step 批数，即训练进度基准；多 collector 下为 per-collector 计数器数组，每个 collector 单写自己的计数，聚合属性求和，见 §8）。seed 不放这里，作为进程入口参数直接传入（collector 的 seed 为 `seed + collector_id`）。 --- @@ -239,11 +240,11 @@ learner 启动即 `publish_weights()`，让 collector 在正式采样前拿到 ```python @dataclass class FastSacAsyncOptionsCfg: - ring_capacity: int = 64 # SharedTransitionRing slot 数 + ring_capacity: int = 64 # SharedTransitionRing slot 数（每 collector 一条独立 ring） utd_mode: str = "strict" # strict=精确比例；learner_bound=吞吐优先 - weight_publish_interval: int = 4 # learner 每 N 次更新发布一次权重 + weight_publish_interval: int = 4 # learner 每 N 次更新发布一次权重（逐 collector 广播） weight_poll_interval: int = 1 # collector 每 N 个 env-step 检查一次新权重 - max_ingest_per_iter: int = 8 # learner 每轮最多 drain 多少 slot + max_ingest_per_iter: int = 8 # learner 每轮对每条 ring 最多 drain 多少 slot idle_sleep_s: float = 0.0005 # 满环/欠数据时的退避睡眠 collector_inference_device: str = "cuda" # cpu / cuda / cuda:N；只控制 actor + policy normalizer collector_compile: bool = True # CUDA 固定 batch 推理使用 reduce-overhead @@ -251,6 +252,8 @@ class FastSacAsyncOptionsCfg: collector_amp_dtype: str = "fp16" # fp16 / bf16 transition_ipc: str = "auto" # 设备 transition 环：auto/on/off（见 §4.1） weight_ipc: str = "auto" # 权重快照 CUDA-IPC 传输：auto/on/off + 大小门控 + num_collectors: int = 1 # collector 进程数；num_envs 均分（须整除） + cpus_per_collector: int | None = None # 每 collector 从其（自动分配的）node 取的 CPU 数；None 用全部 @dataclass @@ -269,11 +272,74 @@ class FastSacCfg: --- -## 8. 不变量与关键取舍 +## 8. 多 collector 与 NUMA 绑定（多 NUMA node 服务器） + +默认拓扑仍是 1 collector × 1 learner。配置 `num_collectors > 1` 后，`num_envs` 均分给 N 个 collector 进程（要求整除），每个 collector 绑定一个 NUMA node 满速采样。核心原则：**learner 的 GPU 永不因数据断粮空转，collector 全速自由跑，一切同步点允许松弛**。 + +**拓扑扩展**： + +- **每 collector 一条独立 SPSC ring**（不做共享 MPSC 环）：ring 的 slot 形状为 `(capacity, num_envs/N, dim)`，所有无锁原语（单写方游标、seqlock）原样保留。满环背压按 ring 独立——快的 collector 阻塞在自己的满环上，不拖住别人。 +- **每 collector 一份独立 `WeightSnapshot`**：learner 逐份非阻塞 `publish`，某份正在被读（seqlock odd）不影响其他份；staleness 以 `async/policy_lag`（聚合 max）与 `async/policy_lag_collector{i}`（多 collector 时分列）监控，不做版本对齐屏障。 +- **learner 轮询多路 ring 并按「代」合并**：`drain()` 对每条 ring 至多消费 `max_ingest_per_iter` 个 slot；第 k 代（各 collector 的第 k 个 slot）拼成完整 `num_envs` 批（env 按 collector 连续分块映射，保证 n-step 的时间相邻性）。一个循环内所有完整批次先各自 `.to(device)`，再在 GPU 侧沿时间轴 stack，最后用 `SimpleReplayBuffer.extend_batch` 每字段一次连续列写入（H2D 次数与 GPU kernel 数不随批量增长——RTX 3090 + 双 CUDA collector 实测，逐批 `extend` 的阻塞式 CPU 源拷贝在多 CUDA context 争用下显著变慢）。未集齐的「代」以 CPU slot 视图挂起——慢 collector 自己的 ring 会先填满并背压它自己，不拖住别人；所有 ring 的读游标只在合并批到达 GPU 后推进，保持 ring 的 no-clobber 保证。`strict` 模式的 UTD 记账因此无需按 collector 数缩放（合并批就是完整 `num_envs` 批），任意 collector 数下长期 UTD 精确等于 `num_updates`；`learner_bound` 下 ring 全空时 learner 不等待，在已有 buffer 上继续 `num_updates`。 +- **数据顺序无关**：off-policy + i.i.d. 采样，多 collector 的 transition 入 buffer 顺序无关，不引入全局序号。 +- **seed 按 collector 划分**（`seed + collector_id`），保证可复现且样本不重复。 +- **进度与日志**：`Control.collector_steps` 是 per-collector 计数器数组（单写方不变），聚合求和后除以 collector 数得到「完整 num_envs 批等价步」，作为训练进度 / 日志 / checkpoint / resume 的统一基准——任意 collector 数下 `num_learning_iterations`、`learning_starts`、TensorBoard x 轴与同步版语义一致；每 collector 有独立 `StatsQueue`（快照携带 `collector_id`）、每个 learner rank 每 log 窗口发一份轻量 payload，全部汇到**父进程**——worker 完全对等，父进程聚合（return/ep_len/timing 取均值、episodes 求和、policy_lag 取 max）后渲染面板并写 TensorBoard。面板/系统采样由未绑定的父进程执行，CPU 视图是全机（所有 core）而非某个 node；多 collector 时额外输出 `async/policy_lag_collector{i}`、`async/ring_fill_collector{i}`，单 collector 的标量键保持不变。 +- **resume 语义**：多 collector 下 resume 与单 collector 相同——learner 从 checkpoint 恢复网络与优化器，所有 collector 重新 reset env 后从 checkpointed iteration 继续采样；ring / 在途 transition 不跨进程恢复。 + +**NUMA 绑定**（拓扑决策在 `async_impl/topology.py`，绑定原语在 `async_impl/numa.py`，等价 `numactl --cpunodebind= --membind=`，best-effort，**自动分配**）： + +- **分配策略**（`resolve_worker_topology`）：learner 绑到其 GPU 的 PCIe 本地 node（经 NVML/pynvml 查 PCI bus id + sysfs `numa_node`，与 system_metrics 同款进程内会话，不创建 CUDA context）；每个 collector 跟随其归属 learner 的 node——collector/learner 对绝不跨 NUMA 边界，transition ring、pinned staging 与权重快照全部 node 本地。单 node 宿主机 / CPU learner / GPU 拓扑未知时不绑定（OS 默认放置）。实测依据：双路双卡机器上整树 node 绑定较不绑 +58%（first-touch 内存页跨 node 是主要开销），collector 与 learner 同 node 的局部性收益远大于带宽减半的损失。 +- CPU affinity：`os.sched_setaffinity` 绑到 node 的 CPU 列表（sysfs `node{X}/cpulist`）；配置 `cpus_per_collector` 时进一步切成不重叠的连续分片，避免 collector 之间抢核。 +- 内存策略：libnuma 的 `set_membind`（ctypes 加载，numactl 同款调用）把该进程**后续**分配绑到本地 node——因此绑定发生在 worker 进程的第一行（父进程在 spawn 前预绑，保证 import 期分配也 node 本地），先于 env / staging / pinned buffer 的任何分配。libnuma 不可用或 node 未知时告警并沿用 OS 默认放置，单 NUMA 机器与容器行为不变。 + +**collector 推理设备**：多 collector 改变 CPU/CUDA 推理的权衡（N 个 CUDA collector 与 learner 产生 N 倍 H2D/D2H burst 争用）。首版保持默认 `collector_inference_device="cuda"`（与单 collector 一致），以实测吞吐决定多 collector 场景默认值是否调整。 + +**明确不做**：共享 MPSC 环 / 原子游标、collector 间同步、per-step 权重同步、多机。 + +--- + +## 9. 单机多卡（多 learner × DDP 数据并行） + +`num_learners > 1` 时每个 GPU 一个 learner 进程，`torch.distributed`（NCCL，file rendezvous 由 parent 注入）做梯度平均；collector 进一步按 `num_collectors % num_learners == 0` 均分给各 learner。所有共享内存原语零改动——每条 ring / 每条 weight channel / Control 仍是严格单写者单读者。 + +```text +parent (mp spawn, 创建全部 shm 原语) +│ +├── learner rank 0 @ cuda:0 ───────── DDP/NCCL 梯度同步（每 rank 本地 batch）──────────┐ +│ ├─ drain ring 0..k-1（本 rank 分片，按「代」合并成完整 num_envs 批） │ +│ ├─ 本地 sharded replay buffer（env 分片 → 每 env 完整轨迹在本 rank，n-step 邻接保持） │ +│ ├─ 唯一 weight 发布者：向【全部】collector publish │ +│ │ （CUDA-IPC 仅限与本 rank 同卡的 collector，其余走 host shm） │ +│ └─ 唯一 checkpoint / TensorBoard / 面板 writer（drain 全部 StatsQueue 聚合） │ +│ ▼ +├── learner rank 1 @ cuda:1 ── drain ring k..2k-1 ── 本地 rb ── 梯度 allreduce ── 与 rank 0 同步 +│ +├── collector 0 (env CPU, 推理 cuda:0) ──ring 0 (SPSC)──► learner 0 +├── collector 1 (env CPU, 推理 cuda:1) ──ring 1 (SPSC)──► learner 1 +└── ... ▲ + └── weight channel（rank 0 → 每个 collector 一份） + +Control（全局共享标量，所有进程可见）：collector_steps[] 聚合 // num_collectors + = 全局进度基准 → 迭代 / 日志 / checkpoint / update 次数全部由它推导（跨 rank 锁步） +``` + +**数据连接与 pipeline（一次循环）**：collector 本地采样（推理在归属 learner 的卡上）→ push 进自己的 SPSC ring，满环背压只作用于自己 → 归属 learner `drain`：把 k 条分片 ring 按代合并成完整 shard 批，一次性写入本地 replay buffer（按 shard env 数定尺寸）→ update：各 rank 以 `batch_size / num_learners` 采样，backward 后手动 all-reduce 梯度取均值（actor/qnet 每次反向后合并为一次扁平 all-reduce；标量 `log_alpha` 单独 all-reduce——不走 DDP 包装，因为更新边界调用的是 `get_actions_and_log_probs`/`projection`/`get_value` 等自定义方法而非 `module.forward`，DDP 既不代理也不同步它们）→ rank 0 按 `weight_publish_interval` 向全部 collector 发布权重 + normalizer 统计。 + +**关键规则**： + +- **update 次数锁步**：梯度同步要求各 rank backward 次数严格一致，因此多 learner 下 update 次数不从本地 `ingested` 推导，而从全局进度基准（`Control.collector_steps` 聚合）的增量推导——所有 rank 看到同一个值，天然锁步；分片慢的 rank 只是让全体多等，不会死锁。 +- **算法量全局、资源量本地**：`batch_size` 全局（内部按 rank 均分，DDP 平均后等价同步版全局 batch）；`num_envs` / `num_collectors` / NUMA 均为每节点量。本地 rb 容量 = 全局 `buffer_size`（总内存 × num_learners，是 env 分片保 n-step 的直接代价）。 +- **配置只给数量**：`num_learners` / `learner_devices`（null → 复制 `device`）；NUMA node 全部自动分配（collector round-robin、learner GPU 本地），无手工列表。`world_size == 1` 时跳过 `init_process_group` 与 DDP 包装，单卡行为逐字节不变；多 learner 下暂禁 learner 侧 `torch.compile`（reduce-overhead CUDA graph 与 DDP hook 的组合首版不碰）。 + +**已知取舍**：obs-normalizer 各 rank 只见本地分片，running stats 有轻微偏差，首版直接发布 rank 0 版本（严格一致可后续加发布前 all-reduce count/mean/var，量极小）；learner 侧 perf 指标只反映 rank 0。多机（torchrun / 网络 rendezvous / 跨机传输）本期不做——逻辑拓扑与部署机制分离，rendezvous 换来源即可平移。 + +--- + +## 10. 不变量与关键取舍 - **算法不变**：transition 构造、`rb.extend` 调用与 `sample` 语义、更新数学与同步版逐字节一致；异构只改「谁在哪个进程执行」。 - **normalizer 单写方**：只有 learner 以 `update=True` 更新 normalizer；collector 只读快照。权重通道因此是单向的。 -- **单生产者单消费者**：环与权重快照都建立在「每个共享量只有一个写方」之上，这是无锁 / 无 CAS 的前提；游标与数据之间的顺序则依赖 x86/TSO（不插屏障，故当前 **x86-only**，弱内存序 ISA 暂不支持）。当前只支持 1 collector × 1 learner、单机；不做多机分布式。 +- **单生产者单消费者**：环与权重快照都建立在「每个共享量只有一个写方」之上，这是无锁 / 无 CAS 的前提；多 collector 通过「每 collector 一条独立 ring + 一份独立 WeightSnapshot」保持该前提，而不是引入 MPSC / 原子游标。游标与数据之间的顺序则依赖 x86/TSO（不插屏障，故当前 **x86-only**，弱内存序 ISA 暂不支持）。当前支持 1..N collector × 1..M learner（单机）、不做多机分布式。 - **背压优先于放开比例**：collector 快时阻塞采样而非无界缓冲，用 `ring_capacity` 吸收抖动，避免 off-policy 失真。 - **checkpoint 与同步版字节兼容**：保证 play / resume 互通与 A/B 对比有效。 - **非确定性**：两进程相对速度随机，逐步复现不可能；正确性以「固定 seed 下 `strict` 模式收敛曲线落在同步版 run-to-run 方差带内」在统计层面成立。seed 同时播撒 collector（env + 采样噪声）与 learner（网络初始化 + 采样噪声）。 @@ -282,6 +348,6 @@ class FastSacCfg: --- -## 9. 一句话总结 +## 11. 一句话总结 FastSAC 的 off-policy 属性 + normalizer 为 learner 独占，使「collector/learner 分进程 + 共享内存」在不改算法、不改同步版的前提下成立；唯一的 `motrix.fastsac` provider 通过 `asynchronous` 字段选择 Trainer，并共用 env/config/checkpoint。默认异步执行；基础配置用 `strict` 验证算法等价，吞吐优先的任务用 `learner_bound` 让采样与训练各自满速。三个必须做对的点是：**SPSC 有界背压环**（防内存失控 / off-policy 失真）、**UTD 比例治理**（吞吐模式下监控并标注实际 UTD）、**seqlock 双缓冲权重快照**（无锁读、杜绝撕裂、靠短发布间隔压低 staleness）。 diff --git a/wiki/plan/fastsac-async-multi-learner.md b/wiki/plan/fastsac-async-multi-learner.md new file mode 100644 index 00000000..3ed679e0 --- /dev/null +++ b/wiki/plan/fastsac-async-multi-learner.md @@ -0,0 +1,17 @@ +# FastSAC 单机多卡（多 learner × DDP）实施计划 + +## 摘要 + +落地 [design 文档 §9](../design/fastsac-async-heterogeneous-trainer.md)：`num_learners > 1` 时每 GPU 一个 learner 进程做 DDP 数据并行，collector 按 `num_collectors % num_learners == 0` 归属到各 learner；`world_size == 1` 时行为与现有单 learner 逐字节一致。 + +## TODO + +- [x] config：`num_learners` / `learner_devices`，yaml 同步；NUMA node 自动分配（collector round-robin、learner GPU 本地），无手工列表 +- [x] 校验：整除关系、设备索引唯一、`batch_size % num_learners == 0` +- [x] agent：`world_size` 参数；DDP 包装 actor/qnet；`log_alpha` 梯度手动 allreduce；本地 batch = `batch_size // num_learners`；多 learner 禁 learner 侧 compile +- [x] learner：`maybe_train_global(gstep)`——update 次数从全局进度基准增量推导（跨 rank 锁步） +- [x] worker：`run_learner_process` 增加 rank / world / rendezvous / device；rank 0 建全部 weight sender 并 ship 全部 slot queue、drain 全部 stats queue、独占 checkpoint/TB；退出前 barrier +- [x] train：spawn N 个 learner；ring 按 `i // k` 归属分区；`collector_inference_device: "cuda"` 解析为归属 learner 的卡 +- [x] 测试：分区映射、锁步 due 计算、config 校验、collector 设备解析（CPU 可跑，不依赖 N 卡） +- [x] ruff + 全量 motrix_rl 测试 + 提交推送 +- [ ] 真实多卡（≥2 GPU）环境端到端验证（本机单卡，仅验证了校验路径与单/多 collector 全链路） diff --git a/wiki/plan/index.md b/wiki/plan/index.md index 169e1079..02b44968 100644 --- a/wiki/plan/index.md +++ b/wiki/plan/index.md @@ -13,4 +13,4 @@ ## 文档列表 -当前暂无进行中的计划。新的功能请先完成 `wiki/design/` 中的设计确认，再在此目录建立对应的实现计划。 +- [fastsac-async-multi-learner.md](fastsac-async-multi-learner.md) — FastSAC 单机多卡（多 learner × DDP）实施计划