From 65587c24d28615c8bce8c6b2a4e9507af7c9ab2c Mon Sep 17 00:00:00 2001 From: zilch Date: Wed, 23 Sep 2026 01:09:51 +0800 Subject: [PATCH] =?UTF-8?q?perf(fastsac):=20fast=20async=20learner=20pipel?= =?UTF-8?q?ine=20=E2=80=94=20batched=20drain,=20CUDA-IPC=20ring,=20fused?= =?UTF-8?q?=20updates?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Learner-side ingest and update-path performance work, plus training panel observability: * Batched async drain: contiguous ring runs (read_span/commit_reads) move through pinned staging as one non-blocking H2D per field on a dedicated copy stream that overlaps gradient updates; replay buffer gains extend_batch for slot-major multi-transition writes. * CUDA-IPC transition ring (transition_ipc auto/on/off): when collector inference and the learner share one GPU, slots live in one fused device tensor — the collector publishes each env-step batch with a single fused H2D, the learner ingests pure D2D strided copies. GPU writes are ordered behind per-slot CUDA events before cursor bumps; endpoint-local issued counters index slots (the published cursor lags in-flight copies). Host ring remains the fallback transport. * Fused learner updates: the target-net soft update moved into the compiled critic step (replays in the same CUDA graph); _own output clones are restricted to what outlives a graph generation, with the gated actor pair always owned (metrics must survive future replays — regression-tested). Multi-step compile blocks and module-only compile variants were benchmarked and rejected (torch.compile splits regions at backward/optimizer.step; per-step boundary is the optimum). * Timing stats report only first-level env_step sub-stages (panel and TensorBoard), removing duplicate child paths. * Training panel: per-GPU utilization/VRAM samplers (sample_per_device + sample_gpu_devices), a dedicated System tab (keyboard 1/2/3) with separate CPU and GPU blocks — aggregate CPU summary in the card title, per-core utilization spectrum (height-block glyphs with a connected ceiling cap; MOTRIX_PANEL_CPU_SPECTRUM selects the style for fonts that flatten block glyphs), per-GPU device table — while overview/timing and the plain-text fallback stay aggregate-only. Measured (microduck-walk-flat scale): learner drain 0.46-0.48 -> 0.14-0.16 ms/batch at 256 envs (3x), 1.86 -> 0.34 ms per 8-slot run at 2048 envs (5.5x); collector-side fused push costs +0.03 ms/batch. --- configs/algo_base/motrix.fastsac.yaml | 10 +- motrix_rl/src/motrix_rl/console.py | 245 ++++++++++++--- motrix_rl/src/motrix_rl/fastsac/agent.py | 52 ++-- .../motrix_rl/fastsac/async_impl/collector.py | 13 +- .../motrix_rl/fastsac/async_impl/learner.py | 122 ++++++-- .../motrix_rl/fastsac/async_impl/shm/ring.py | 146 --------- .../src/motrix_rl/fastsac/async_impl/train.py | 14 +- .../async_impl/{shm => transport}/__init__.py | 26 +- .../async_impl/{shm => transport}/common.py | 2 +- .../fastsac/async_impl/transport/ipc_ring.py | 232 ++++++++++++++ .../fastsac/async_impl/transport/ring.py | 192 ++++++++++++ .../{shm => transport}/weight_channel.py | 4 +- .../motrix_rl/fastsac/async_impl/worker.py | 126 ++++++-- motrix_rl/src/motrix_rl/fastsac/buffer.py | 30 ++ motrix_rl/src/motrix_rl/fastsac/config.py | 8 + motrix_rl/src/motrix_rl/fastsac/sync/train.py | 9 +- motrix_rl/src/motrix_rl/system_metrics.py | 199 +++++++++--- motrix_rl/tests/test_console.py | 109 ++++++- motrix_rl/tests/test_fastsac_buffer.py | 86 ++++++ motrix_rl/tests/test_fastsac_collector.py | 13 +- motrix_rl/tests/test_fastsac_ipc_ring.py | 287 ++++++++++++++++++ motrix_rl/tests/test_fastsac_learner.py | 92 ++++++ motrix_rl/tests/test_rl_sim_backend.py | 14 +- motrix_rl/tests/test_system_metrics.py | 60 +++- scripts/bench_fastsac_collector_inference.py | 221 -------------- .../fastsac-async-heterogeneous-trainer.md | 32 +- 26 files changed, 1790 insertions(+), 554 deletions(-) delete mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py rename motrix_rl/src/motrix_rl/fastsac/async_impl/{shm => transport}/__init__.py (51%) rename motrix_rl/src/motrix_rl/fastsac/async_impl/{shm => transport}/common.py (98%) create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ipc_ring.py create mode 100644 motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ring.py rename motrix_rl/src/motrix_rl/fastsac/async_impl/{shm => transport}/weight_channel.py (99%) create mode 100644 motrix_rl/tests/test_fastsac_ipc_ring.py delete mode 100644 scripts/bench_fastsac_collector_inference.py diff --git a/configs/algo_base/motrix.fastsac.yaml b/configs/algo_base/motrix.fastsac.yaml index d43b6a07..68974902 100644 --- a/configs/algo_base/motrix.fastsac.yaml +++ b/configs/algo_base/motrix.fastsac.yaml @@ -101,11 +101,19 @@ trainer: # machine's CPU/GPU balance (see issue #62's sweep table). learner_cpu_cores: null collector_cpu_cores: null + # Transition-ring transport: "auto" (default) puts the ring slots in + # CUDA-IPC device memory when learner and collector inference share one + # GPU (fused single H2D on the collector, D2D-only learner ingest) and + # falls back to host shared memory otherwise; "on"/"off" force the + # device/host path ("on" warns and falls back to the host ring when the + # two sides are not on one GPU). Keep the values quoted: unquoted on/off + # parse as booleans in YAML. + transition_ipc: "auto" # Weight-snapshot transport: "auto" (default) enables CUDA-IPC device slots # only when learner and collector share one GPU and the actor params reach # weight_ipc_min_bytes; "on"/"off" force the device/host path ("on" warns # and falls back to the host path when learner and collector are not on one # GPU). Keep the values quoted: unquoted on/off parse as booleans in YAML. - weight_ipc: auto + weight_ipc: "auto" # Minimum flattened parameter bytes for the CUDA-IPC path under "auto". weight_ipc_min_bytes: 16777216 diff --git a/motrix_rl/src/motrix_rl/console.py b/motrix_rl/src/motrix_rl/console.py index 79cde678..9b6715ed 100644 --- a/motrix_rl/src/motrix_rl/console.py +++ b/motrix_rl/src/motrix_rl/console.py @@ -10,11 +10,11 @@ import re import shutil import sys -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any -from motrix_rl.system_metrics import CpuLoad, MemoryUsage +from motrix_rl.system_metrics import CpuLoad, GpuDeviceUsage, MemoryUsage try: # cbreak keyboard input needs a POSIX terminal; other platforms keep a plain Live import select @@ -69,6 +69,9 @@ class TrainingPanelStats: gpu_utilization_percent: float | None = None memory_usage: MemoryUsage | None = None gpu_memory_usage: MemoryUsage | None = None + # Per-device GPU utilization/memory; when present the system card renders + # one line per accelerator instead of the aggregate GPU/VRAM fields. + gpu_devices: Sequence[GpuDeviceUsage] | None = None checkpoint_path: str | None = None @@ -131,9 +134,9 @@ def open_training_live(): def emit_training_panel(live, stats: TrainingPanelStats, *, title: str = "rl") -> None: - """Render one training panel; 1/2 switch overview and timing views.""" + """Render one training panel; 1/2/3 switch overview, timing and system views.""" if live is not None: - detail = bool(getattr(live, "_motrix_detail", False)) + view = getattr(live, "_motrix_view", "overview") input_fd = getattr(live, "_input_fd", None) if input_fd is not None: ready, _, _ = select.select([input_fd], [], [], 0) @@ -142,12 +145,10 @@ def emit_training_panel(live, stats: TrainingPanelStats, *, title: str = "rl") - command = os.read(input_fd, 1).decode("utf-8", errors="ignore").lower() except OSError: command = "" - if command == "2": - detail = True - elif command == "1": - detail = False - live._motrix_detail = detail - panel = render_training_panel(stats, title=title, detail=detail) + if command in ("1", "2", "3"): + view = {"1": "overview", "2": "timing", "3": "system"}[command] + live._motrix_view = view + panel = render_training_panel(stats, title=title, view=view) live.update(panel, refresh=True) else: print(format_training_panel(stats, title=title)) @@ -221,6 +222,16 @@ def si(n: float) -> str: f"({load.used_logical_cpus:.1f}/{load.logical_cpu_count}T{cores}) " f"iowait {load.iowait_percent:.1f}% steal {load.steal_percent:.1f}%" ) + # Plain-text fallback mirrors the overview card: aggregate GPU stats only + # (per-device detail is the interactive System view's job). + devices = stats.gpu_devices or () + if devices: + gpu_util, gpu_mem = _aggregate_gpu_devices(devices) + else: + gpu_util, gpu_mem = stats.gpu_utilization_percent, stats.gpu_memory_usage + if gpu_util is not None or gpu_mem is not None: + parts = [f"{gpu_util:.0f}%" if gpu_util is not None else "n/a", _format_memory(gpu_mem)] + lines.append(" system gpu " + " ".join(parts)) metrics = stats.training_metrics buf = f"{si(stats.buffer_size)}/{si(stats.buffer_capacity)}" if metrics is None: @@ -466,23 +477,165 @@ def _prototype_bar(fraction: float, *, width: int = 22, style: str = "cyan"): return result -def _prototype_tabs(stats: TrainingPanelStats, detail: bool): +def _cpu_spectrum_rows(per_core: Sequence[float], load_style, 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 | + bitmap; default height). The code cannot probe the terminal font, so the + override exists for fonts that render partial-height or shade block + glyphs as indistinguishable equal blocks: + + * ``height`` — one ``▁▂▃▄▅▆▇█`` glyph per core (prettiest; needs a font + with correct block-element metrics, e.g. Cascadia Mono, Sarasa Term, + any Nerd Font); + * ``shade`` — density glyphs ``░▒▓█``, same height by design; + * ``bitmap`` — btop-style columns of ``█``/spaces on 4 rows, only two + distinct characters so it survives every font. + + One space separates adjacent cores so same-height neighbours don't fuse + into a solid bar; the default row of 48 cores + 47 gaps spans the same + ~95 columns as the previous unspaced 96-core row. + """ + from rich.text import Text + + style = os.environ.get("MOTRIX_PANEL_CPU_SPECTRUM", "height").lower() + if style not in ("height", "shade", "bitmap"): + style = "height" + rows: list = [] + 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 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 + # like it touches the 100% reference). + levels = " ▁▂▃▄▅▆▇█" if style == "height" else "░▒▓█" + spectrum = Text() + for i, value in enumerate(chunk): + if i: + spectrum.append(" ") + glyph = levels[min(int(value / 100.0 * (len(levels) - 1)), len(levels) - 1)] + spectrum.append(glyph, style=load_style(value)) + if style == "height": + # Ceiling line above the skyline: marks each column's 100% + # reference. The lower-one-eighth block is used because it sits + # at the BOTTOM of its own cell — a full-height column below + # touches it, forming one connected bar (▔ would leave a 7/8-row + # gap between the cap and the skyline). + rows.append(Text.assemble((" " * label_width, "dim"), (" ".join("▁" for _ in chunk), "dim"))) + rows.append(Text.assemble((label, "dim"), (spectrum,))) + else: + height = 4 + for level in range(height, 0, -1): # top row first + threshold = 100.0 * (level - 1) / height + row = Text() + for i, value in enumerate(chunk): + if i: + row.append(" ") + row.append( + "█" if value > threshold else " ", style=load_style(value) if value > threshold else "dim" + ) + rows.append(Text.assemble((label if level == height else " " * label_width, "dim"), (row,))) + return rows + + +def _prototype_system_page(stats: TrainingPanelStats, load_style, memory_style): + """Dedicated per-device system view: host summary + one row per GPU. + + The overview/timing cards stay aggregate-only, so 8-GPU hosts get their + per-device detail here instead of an overflowing System health card. CPU + load is shown as a compact per-core spectrum (one block glyph per logical + core, height and color by utilization) — readable at 192 cores without a + per-core table. + """ + from rich.table import Table + from rich.text import Text + + load = stats.cpu_load + cpu_parts: list[Any] = [] + 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)) + else: + # RAM is not repeated here — the always-visible System health card in + # the summary row already carries it. + cpu_parts.append(Text("no per-core samples", style="dim")) + + devices = stats.gpu_devices or () + columns = Table(header_style="dim", expand=True, show_edge=False, box=None) + columns.add_column("Device", no_wrap=True) + columns.add_column("Utilization", justify="right", no_wrap=True) + columns.add_column("VRAM", justify="right", no_wrap=True) + columns.add_column("VRAM load", ratio=1) + if not devices: + columns.add_row( + Text("no per-device GPU samples", style="dim"), + Text( + f"{stats.gpu_utilization_percent:.0f}%" if stats.gpu_utilization_percent is not None else "n/a", + style=load_style(stats.gpu_utilization_percent), + ), + Text(_format_memory(stats.gpu_memory_usage), style=memory_style(stats.gpu_memory_usage)), + "", + ) + for device in devices: + util = device.utilization_percent + bar = "" + if device.memory is not None and device.memory.total_bytes > 0: + width = 16 + filled = round(width * device.memory.used_bytes / device.memory.total_bytes) + bar = "█" * filled + "░" * (width - filled) + device_cell = Text(f"GPU{device.index}", style="white") + if device.name: + device_cell.append(f" {device.name}", style="dim") + columns.add_row( + device_cell, + Text(f"{util:.0f}%" if util is not None else "n/a", style=load_style(util)), + Text(_format_memory(device.memory), style=memory_style(device.memory)), + Text(bar, style=memory_style(device.memory)), + ) + + # CPU and GPU are separate blocks: the aggregate CPU summary sits in the + # CPU card's title (top-right corner), RAM heads the body next to the + # per-core spectrum, and the device table fills the GPU card. + if load is not None: + summary = ( + f"CPU {load.utilization_percent:.0f}% " + f"({load.used_logical_cpus:.1f}/{load.logical_cpu_count}T" + + (f", {load.physical_core_count}C" if load.physical_core_count is not None else "") + + ")" + ) + cpu_title = Text(summary, justify="right", style=f"bold {load_style(load.utilization_percent)}") + else: + cpu_title = Text("CPU n/a", justify="right", style="dim") + cpu_card = Panel( + Group(*cpu_parts), + title=cpu_title, + border_style="grey37", + padding=(0, 1), + ) + gpu_card = Panel( + columns, + title=f"GPU ({len(devices)} devices)" if devices else "GPU", + border_style="grey37", + padding=(0, 1), + ) + return Group(cpu_card, gpu_card) + + +def _prototype_tabs(stats: TrainingPanelStats, view: str): from rich.text import Text tabs = Table.grid(expand=True, padding=(0, 2)) tabs.add_column() tabs.add_column() + tabs.add_column() tabs.add_column(ratio=1, justify="right") - active = "Timing" if detail else "Overview" - labels = (("Overview", ""), ("Timing", "")) - row = [] - for label, count in labels: - item = Text(label, style="bold cyan" if label == active else "dim") - if count: - item.append(f" {count}", style="grey50") - row.append(item) + labels = ("Overview", "Timing", "System") + row = [Text(label, style="bold cyan" if label.lower() == view else "dim") for label in labels] if _POSIX_TTY: # key handling is POSIX-only; don't advertise it elsewhere - row.append(Text("keyboard: 1/2 switch tabs", style="dim")) + row.append(Text("keyboard: 1/2/3 switch tabs", style="dim")) tabs.add_row(*row) return tabs @@ -494,6 +647,19 @@ def _format_memory(memory: MemoryUsage | None) -> str: return f"{memory.used_bytes / gib:.1f}/{memory.total_bytes / gib:.1f} GiB" +def _aggregate_gpu_devices(devices: Sequence[GpuDeviceUsage]) -> tuple[float | None, MemoryUsage | None]: + """Mean utilization and summed memory across the reported devices.""" + utils = [d.utilization_percent for d in devices if d.utilization_percent is not None] + mean_util = sum(utils) / len(utils) if utils else None + known = [d.memory for d in devices if d.memory is not None and d.memory.total_bytes > 0] + total_mem = ( + MemoryUsage(used_bytes=sum(m.used_bytes for m in known), total_bytes=sum(m.total_bytes for m in known)) + if known + else None + ) + return mean_util, total_mem + + def _run_progress_time_text(stats: TrainingPanelStats): """One ``elapsed · ~remaining`` line for the Run progress card.""" from rich.text import Text @@ -506,11 +672,13 @@ def _run_progress_time_text(stats: TrainingPanelStats): return line -def render_training_panel(stats: TrainingPanelStats, *, title: str = "rl", detail: bool = False): +def render_training_panel(stats: TrainingPanelStats, *, title: str = "rl", view: str = "overview"): if not _RICH: raise RuntimeError("rich is not available") from rich.text import Text + if view not in ("overview", "timing", "system"): + raise ValueError(f"unknown panel view {view!r}") progress = max(0.0, min(1.0, stats.iteration / max(stats.total_iterations, 1))) collect, learn = _timing_totals(stats) @@ -546,22 +714,25 @@ def memory_style(memory: MemoryUsage | None) -> str: progress_row = Table.grid(expand=True, padding=(0, 1)) progress_row.add_column(ratio=1) progress_row.add_row(_prototype_bar(progress, width=18)) + # The overview/timing cards stay aggregate-only; per-device detail lives on + # the dedicated System view so 8-GPU hosts don't blow up the card layout. + devices = stats.gpu_devices or () + if devices: + gpu_util, gpu_mem = _aggregate_gpu_devices(devices) + else: + gpu_util, gpu_mem = stats.gpu_utilization_percent, stats.gpu_memory_usage + left_group = Group( + Text(cpu_text, style=f"bold {load_style(load.utilization_percent if load else None)}"), + Text(f"GPU {gpu_util:.0f}%" if gpu_util is not None else "GPU n/a", style=f"{load_style(gpu_util)}"), + ) + right_group = Group( + Text(f"RAM {_format_memory(stats.memory_usage)}", style=memory_style(stats.memory_usage)), + Text(f"VRAM {_format_memory(gpu_mem)}", style=memory_style(gpu_mem)), + ) system_health = Table.grid(expand=True, padding=(0, 1)) system_health.add_column(ratio=1) system_health.add_column(ratio=1) - system_health.add_row( - Group( - Text(cpu_text, style=f"bold {load_style(load.utilization_percent if load else None)}"), - Text( - f"GPU {stats.gpu_utilization_percent:.0f}%" if stats.gpu_utilization_percent is not None else "GPU n/a", - style=f"{load_style(stats.gpu_utilization_percent)}", - ), - ), - Group( - Text(f"RAM {_format_memory(stats.memory_usage)}", style=memory_style(stats.memory_usage)), - Text(f"VRAM {_format_memory(stats.gpu_memory_usage)}", style=memory_style(stats.gpu_memory_usage)), - ), - ) + system_health.add_row(left_group, right_group) summary.add_row( card( f"Run progress ({progress * 100:.1f}%)", @@ -635,7 +806,9 @@ def memory_style(memory: MemoryUsage | None) -> str: Group(*env_parts), title=f"Environment metrics ({len(env_items)})", border_style="grey37", padding=(0, 1) ) - if detail: + if view == "system": + lower = Group(_prototype_system_page(stats, load_style, memory_style)) + elif view == "timing": timing_blocks: list[Any] = [] columns = Table.grid(expand=True, padding=(0, 1)) columns.add_column(ratio=1) @@ -681,7 +854,7 @@ def memory_style(memory: MemoryUsage | None) -> str: lower = Group(*timing_blocks) if timing_blocks else Text("timing details unavailable", style="grey70") else: lower = Group(*left_blocks, environment) - blocks = [summary, lower, _prototype_tabs(stats, detail)] + blocks = [summary, lower, _prototype_tabs(stats, view)] if stats.checkpoint_path: blocks.insert(-1, Text(f"✓ saved checkpoint {stats.checkpoint_path}", style="green")) # Let Rich use the actual terminal width. The card bodies are intentionally diff --git a/motrix_rl/src/motrix_rl/fastsac/agent.py b/motrix_rl/src/motrix_rl/fastsac/agent.py index 24e47ace..332f7ef1 100644 --- a/motrix_rl/src/motrix_rl/fastsac/agent.py +++ b/motrix_rl/src/motrix_rl/fastsac/agent.py @@ -171,7 +171,12 @@ def __init__( inductor_config.compile_threads = 1 # B3 experiment: reduce-overhead wraps each compiled update in a # CUDA graph (trees), removing per-kernel launch and region-gap CPU - # time inside the hot learner loop. + # time inside the hot learner loop. (Coarser units were measured at + # microduck scale and do NOT help: torch.compile splits regions + # containing backward()/optimizer.step() regardless of the outer + # boundary, and module-only/default-mode compilation is ~50% slower + # or crashes on cudagraph output pools. The per-step boundary is + # the measured optimum.) self._update_main_runtime = torch.compile(self._update_main, mode="reduce-overhead") self._update_pol_runtime = torch.compile(self._update_pol, mode="reduce-overhead") @@ -226,6 +231,19 @@ def _update_main(self, b: dict): torch.nn.utils.clip_grad_norm_(self.qnet.parameters(), cfg.max_grad_norm) self.q_optimizer.step() + # Target-network soft update, fused into the compiled region: the + # foreach ops write the same kind of param storage the optimizer step + # above already mutates in place, so the whole thing replays as one + # CUDA graph instead of two eager kernel launches per step. Ordering + # constraints: after the target read at the top of this function and + # after q_optimizer.step() (both satisfied here); _update_pol reads + # neither qnet_target nor writes qnet, so it stays order-independent. + with torch.no_grad(): + tau = cfg.tau + tgt = [p.data for p in self.qnet_target.parameters()] + torch._foreach_mul_(tgt, 1.0 - tau) + torch._foreach_add_(tgt, [p.data for p in self.qnet.parameters()], alpha=tau) + alpha_loss = torch.zeros((), device=self.device) if cfg.use_autotune: alpha_loss = (-self.log_alpha.exp() * (next_logp.detach() + self.target_entropy)).mean() @@ -255,14 +273,6 @@ def _update_pol(self, b: dict): self.actor_optimizer.step() return actor_loss.detach().float(), (-log_probs.mean()).detach().float() - @torch.no_grad() - def _soft_update(self): - tau = self.cfg.tau - src = [p.data for p in self.qnet.parameters()] - tgt = [p.data for p in self.qnet_target.parameters()] - torch._foreach_mul_(tgt, 1.0 - tau) - torch._foreach_add_(tgt, src, alpha=tau) - def update(self, num_updates: int): """Run ``num_updates`` gradient steps, each on a fresh batch. @@ -288,7 +298,7 @@ def update(self, num_updates: int): return None batch_per_env = max(cfg.batch_size // self.num_envs, 1) last = (torch.zeros((), device=self.device),) * 5 - timing_s = {key: 0.0 for key in ("sample_normalize", "critic_alpha", "actor", "soft_update")} + timing_s = {key: 0.0 for key in ("sample_normalize", "critic_alpha", "actor")} update_started = time.perf_counter() # Batched data preparation (Holosoma-style): sample once and normalize # once per update() call, then slice views into per-gradient-step @@ -319,22 +329,28 @@ def update(self, num_updates: int): # Required with reduce-overhead (CUDA graph trees): open a new graph # generation for this update. NOTE it does not preserve anything -- # it *invalidates* the previous generation's outputs, which is why - # every output that outlives its own iteration goes through `_own`. + # any output that outlives its own iteration goes through `_own`. torch.compiler.cudagraph_mark_step_begin() stage_started = time.perf_counter() - qf_loss, alpha_loss, qf_max, qf_min = _own(self._update_main_runtime(b)) + outputs = self._update_main_runtime(b) timing_s["critic_alpha"] += time.perf_counter() - stage_started - actor_loss, entropy = last[3], last[4] + actor_pair = (last[3], last[4]) if (self.update_idx + i) % cfg.policy_frequency == 0: stage_started = time.perf_counter() - actor_loss, entropy = _own(self._update_pol_runtime(b)) + pol_outputs = self._update_pol_runtime(b) timing_s["actor"] += time.perf_counter() - stage_started + # Always own: the pair is carried across later generations + # within this call AND the returned metrics must stay readable + # after future update() calls replay the pol graph. + actor_pair = _own(pol_outputs) + + # Only the final step's main outputs feed the returned metrics; own + # them so they survive future update() calls' replays. Earlier + # steps' outputs are never read — only replaced below. + main_outputs = _own(outputs) if i == num_updates - 1 else outputs + last = (*main_outputs, *actor_pair) - stage_started = time.perf_counter() - self._soft_update() - timing_s["soft_update"] += time.perf_counter() - stage_started - last = (qf_loss, alpha_loss, qf_max, actor_loss, entropy) self.update_idx += num_updates timing_s["total"] = time.perf_counter() - update_started self._last_update_timing_ms = {key: value * 1000.0 for key, value in timing_s.items()} 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 96f8fdca..48559c9a 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/collector.py @@ -5,7 +5,7 @@ Holds its own :class:`~motrix_rl.fastsac.networks.Actor` and read-only :class:`~motrix_rl.fastsac.buffer.EmpiricalNormalization`, both refreshed from -the learner via its :class:`~motrix_rl.fastsac.async_impl.shm.WeightReceiver` endpoint. Each step +the learner via its :class:`~motrix_rl.fastsac.async_impl.transport.WeightReceiver` endpoint. Each step mirrors the sync collector phase (``agent.py`` collect phase) exactly: decide action -> ``env.step`` -> push the transition batch to the shared ring -> update episode bookkeeping. The normalizer is used read-only (``update=False``), matching @@ -20,8 +20,8 @@ import torch from torch import nn -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, bind_flat_params -from motrix_rl.fastsac.async_impl.shm.weight_channel import WeightReceiver +from motrix_rl.fastsac.async_impl.transport import Control, IpcTransitionRing, SharedTransitionRing, bind_flat_params +from motrix_rl.fastsac.async_impl.transport.weight_channel import WeightReceiver from motrix_rl.fastsac.buffer import EmpiricalNormalization from motrix_rl.fastsac.config import FastSacAgentCfg, FastSacCfg from motrix_rl.fastsac.networks import Actor @@ -76,7 +76,7 @@ def __init__( act_dim: int, action_scale: torch.Tensor, action_bias: torch.Tensor, - ring: SharedTransitionRing, + ring: SharedTransitionRing | IpcTransitionRing, weights: WeightReceiver, control: Control, is_resume: bool = False, @@ -350,6 +350,11 @@ def snapshot_stats(self) -> dict: env_perf = self._env_perf if env_perf is not None: for path, mean_ms in env_perf.stage_mean_ms("step").items(): + if "." in path: + # Only the first sub-stage level is reported: a parent's + # total already includes its children, so deeper paths + # would duplicate time without adding actionable signal. + continue stats["timing_ms"][f"env_step.{path}"] = mean_ms env_perf.reset() self.term_accum, self.term_count = {}, 0 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 2f444103..5086b5cd 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/learner.py @@ -4,10 +4,11 @@ """Learner: owns a full ``FastSacAgent`` and drives training off the shared ring. Unlike the sync trainer it does NOT step the env. It drains raw transitions from -:class:`~motrix_rl.fastsac.async_impl.shm.SharedTransitionRing` into the agent's GPU -replay buffer, runs gradient updates governed by ``utd_mode`` (§6 of the -design), and periodically publishes actor weights + obs-normalizer stats to the -collector via its :class:`~motrix_rl.fastsac.async_impl.shm.WeightSender` endpoint. +the shared transition ring (host fields or CUDA-IPC device fields, see +``transport/ring.py`` / ``transport/ipc_ring.py``) into the agent's GPU replay buffer, runs +gradient updates governed by ``utd_mode`` (§6 of the design), and periodically +publishes actor weights + obs-normalizer stats to the collector via its +:class:`~motrix_rl.fastsac.async_impl.transport.WeightSender` endpoint. The update math is reused unchanged from the sync agent: this module delegates the per-step gradient work to ``agent.update(n)`` and only owns the @@ -18,9 +19,11 @@ import time +import torch + from motrix_rl.fastsac.agent import FastSacAgent -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing -from motrix_rl.fastsac.async_impl.shm.weight_channel import WeightSender +from motrix_rl.fastsac.async_impl.transport import Control, IpcTransitionRing, SharedTransitionRing +from motrix_rl.fastsac.async_impl.transport.weight_channel import WeightSender from motrix_rl.fastsac.config import FastSacCfg @@ -29,7 +32,7 @@ def __init__( self, agent: FastSacAgent, cfg: FastSacCfg, - ring: SharedTransitionRing, + ring: SharedTransitionRing | IpcTransitionRing, weights: WeightSender, control: Control, ): @@ -42,6 +45,31 @@ def __init__( self._learning_starts = agent.cfg.learning_starts self._last_publish_ms = 0.0 + # Host-ring async-ingest plumbing (CUDA only): contiguous runs are + # staged through pinned buffers and moved to the GPU with non-blocking + # H2D copies on a dedicated stream, so ingestion overlaps gradient + # updates. The CUDA-IPC device ring needs none of this — its slots are + # already on the device and ingest is a same-stream D2D copy. + self._device_ring = isinstance(ring, IpcTransitionRing) + self._copy_stream = None + self._copy_event = None + self._pending_copy = False + self._staging = None + if agent.device.type == "cuda" and not self._device_ring: + self._copy_stream = torch.cuda.Stream(device=agent.device) + self._copy_event = torch.cuda.Event() + rb = agent.rb + chunk = max(self.async_options.max_ingest_per_iter, 1) + pin = lambda *shape: torch.empty(*shape, pin_memory=True) # noqa: E731 + self._staging = ( + pin(chunk, rb.n_env, rb.n_obs), + pin(chunk, rb.n_env, rb.n_critic_obs), + pin(chunk, rb.n_env, rb.n_act), + pin(chunk, rb.n_env), + torch.empty(chunk, rb.n_env, dtype=torch.int64, pin_memory=True), + torch.empty(chunk, rb.n_env, dtype=torch.int64, pin_memory=True), + ) + # keep normalizers/actor in train mode: the learner is the update side. self.agent.set_train_mode() @@ -56,33 +84,62 @@ def update_idx(self) -> int: def drain(self) -> int: """Move up to ``max_ingest_per_iter`` ring slots into the replay buffer. - Returns the number of slots ingested. Read cursor advances only after the - GPU copy, so the collector cannot clobber an in-flight slot. The replay - buffer derives each transition's ``next_obs`` from the following slot's - stored observation, so no successor peek is needed and a slot is - ingested as soon as it is committed. + Returns the number of slots ingested. Slots are consumed in contiguous + runs (``ring.read_span()``): + + * CUDA-IPC device ring: the strided device views go straight into the + replay buffer with D2D copies on the current stream (ordered before + any subsequent sample by stream order); the read cursor is released + behind an event once those copies complete. + * host ring on CUDA: each run is memcpy'd into pinned staging and moved + to the GPU as one non-blocking H2D copy per field on the copy stream + (overlapping the next gradient update), then the read cursor + advances — the producer cannot clobber in-flight data because the + ring slot was already fully copied to staging. + * host ring on CPU: the views copy directly into the buffer. + + The replay buffer derives each transition's ``next_obs`` from the + following slot's stored observation, so no successor peek is needed and + a slot is ingested as soon as it is committed. """ - device = self.agent.device + budget = max(self.async_options.max_ingest_per_iter, 1) + if self._staging is not None: + # Staging may still be the source of the previous drain's in-flight + # H2D copies; wait before overwriting it. Any pending copy was + # already joined by the intervening maybe_train(), so this is + # normally a no-op sync. + self._copy_stream.synchronize() ingested = 0 - for _ in range(max(self.async_options.max_ingest_per_iter, 1)): - if not self.ring.has_next(): - break - slot = self.ring.read_slot() - assert slot is not None # has_next implies a readable slot - obs, critic_obs, actions, rewards, dones, truncations = slot - self.agent.rb.extend( - obs.to(device), - critic_obs.to(device), - actions.to(device), - rewards.to(device), - dones.to(device), - truncations.to(device), - ) - self.ring.commit_read() - ingested += 1 + while ingested < budget and self.ring.has_next(): + k, views = self.ring.read_span() + k = min(k, budget - ingested) + if k < views[0].shape[0]: + views = tuple(v[:k] for v in views) + if self._staging is not None: + for stage, view in zip(self._staging, views): + stage[:k].copy_(view) # ring -> pinned (plain CPU memcpy) + with torch.cuda.stream(self._copy_stream): + self.agent.rb.extend_batch(*(stage[:k] for stage in self._staging)) + self._pending_copy = True + else: + # CUDA-IPC device ring or CPU host ring: consume the views + # directly (D2D strided copy, or plain CPU copy). + self.agent.rb.extend_batch(*views) + self.ring.commit_reads(k) + ingested += k + if self._pending_copy: + self._copy_event.record(self._copy_stream) return ingested # ------------------------------------------------------------------ update + def wait_ingest(self) -> None: + """Block until every issued ingest copy has landed in the buffer.""" + if self._copy_stream is not None: + self._copy_stream.synchronize() + elif self._device_ring: + torch.cuda.synchronize(self.agent.device) + self._pending_copy = False + def _ready(self) -> bool: return self.control.collector_steps >= self._learning_starts and self.agent.rb.num_stored > 0 @@ -101,6 +158,13 @@ def maybe_train(self, ingested: int) -> dict | None: """Run ratio-governed updates. Returns last metrics dict or ``None``.""" if not self._ready(): return None + if self._pending_copy: + # Sampling reads slots the copy stream may still be filling; make + # the compute stream wait for the in-flight H2D ingest copies. + # (The device-ring path needs no event: its D2D copies are on the + # same stream as sampling.) + self._copy_event.wait() + self._pending_copy = False n = self._num_updates_for(ingested) # Delegate the per-step work to the agent; this module no longer keeps # its own update-loop / update_idx / _last_actor — the agent's diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py deleted file mode 100644 index c6cfc0e7..00000000 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/ring.py +++ /dev/null @@ -1,146 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""SPSC transition ring between the collector and learner processes.""" - -from __future__ import annotations - -import torch - -from motrix_rl.fastsac.async_impl.shm.common import _shared - - -# ---------------------------------------------------------------- transition ring -class SharedTransitionRing: - """SPSC ring of transition batches with bounded backpressure. - - Each slot holds one env-step batch: the six tensors produced by one - collector step, with the leading dimension being ``num_envs`` (so a slot - is ``(num_envs, dim)``). ``next_obs``/``next_critic_obs`` are NOT stored: - with auto-reset envs the observation returned by step ``t`` is exactly the - stored observation of step ``t+1`` (at episode ends it is the reset - observation), so the consumer's replay buffer derives them from the - successive slot. This halves the per-slot copy volume and shared-memory - footprint. - - Producer (collector) calls :meth:`push`; when the ring is full it returns - ``False`` and the caller must retry/backoff — that is the backpressure that - keeps the collector from outrunning the learner and flooding memory. - - Consumer (learner) calls :meth:`read_slot` to get zero-copy CPU views of the - oldest unread slot, moves them to its device, then calls :meth:`commit_read`. - The read cursor only advances after the copy, so the producer can never - clobber a slot that is still being ingested (``push`` blocks while the ring - is full). - - Memory ordering - ~~~~~~~~~~~~~~~ - Only ``_write`` is read by the consumer, only ``_read`` is read by the - producer, and each cursor has a single writer — so an aligned int64 store - is enough to publish progress and no atomic RMW is needed. Correctness also - needs the consumer to not observe the ``_write`` bump before the slot's data - stores have landed (and symmetrically for ``_read``); on x86/TSO that - ordering is free, so no memory barrier is used and this path is x86-only - (see the module "Memory ordering" note). - """ - - FIELDS = ( - "obs", - "critic_obs", - "actions", - "rewards", - "dones", - "truncations", - ) - - def __init__( - self, - capacity: int, - num_envs: int, - obs_dim: int, - critic_obs_dim: int, - act_dim: int, - ): - if capacity < 1: - raise ValueError(f"ring_capacity must be >= 1, got {capacity}") - self.capacity = capacity - self.num_envs = num_envs - f32, i64 = torch.float32, torch.int64 - self.obs = _shared((capacity, num_envs, obs_dim), f32) - self.critic_obs = _shared((capacity, num_envs, critic_obs_dim), f32) - self.actions = _shared((capacity, num_envs, act_dim), f32) - self.rewards = _shared((capacity, num_envs), f32) - self.dones = _shared((capacity, num_envs), i64) - self.truncations = _shared((capacity, num_envs), i64) - # cursors are shared so the two processes see each other's progress. - self._write = _shared((1,), i64) - self._read = _shared((1,), i64) - - @property - def write_idx(self) -> int: - # Consumer reads this; single-writer producer, so a plain aligned int64 - # load is coherent. On x86/TSO the data loads that follow cannot be - # reordered ahead of it, so no acquire barrier is needed (x86-only). - return int(self._write[0]) - - @property - def read_idx(self) -> int: - # Producer reads this; single-writer consumer. On x86/TSO the following - # is_full() decision is based on an up-to-date value without a barrier. - return int(self._read[0]) - - def size(self) -> int: - """Number of unread slots currently buffered.""" - return self.write_idx - self.read_idx - - def is_full(self) -> bool: - return self.size() >= self.capacity - - def push(self, obs, critic_obs, actions, rewards, dones, truncations) -> bool: - """Copy one env-step batch into the next slot. Returns False if full. - - Inputs are CPU tensors shaped ``(num_envs, dim)``; ``dones``/``truncations`` - are int64 to match ``SimpleReplayBuffer.extend`` semantics. - """ - if self.is_full(): - return False - slot = self.write_idx % self.capacity - self.obs[slot].copy_(obs) - self.critic_obs[slot].copy_(critic_obs) - self.actions[slot].copy_(actions) - self.rewards[slot].copy_(rewards) - self.dones[slot].copy_(dones) - self.truncations[slot].copy_(truncations) - # Publish the slot. On x86/TSO the field copies above are guaranteed - # visible before this cursor bump, so a consumer that reads the new - # write_idx also sees the data (x86-only; ARM would need a release here). - self._write[0] += 1 - return True - - def read_slot(self): - """Return CPU views of the oldest unread slot, or ``None`` if empty. - - Does NOT advance the read cursor; call :meth:`commit_read` after the - consumer has finished copying the data elsewhere. - """ - if self.size() <= 0: - return None - slot = self.read_idx % self.capacity - return ( - self.obs[slot], - self.critic_obs[slot], - self.actions[slot], - self.rewards[slot], - self.dones[slot], - self.truncations[slot], - ) - - def has_next(self) -> bool: - """Whether at least one unread slot is committed (i.e. readable).""" - return self.size() > 0 - - def commit_read(self) -> None: - # Free the slot. On x86/TSO our reads above complete before this cursor - # bump, so the producer's is_full() cannot reuse a slot we are still - # copying out (x86-only; ARM would need a release here). - self._read[0] += 1 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 72006e8f..59e8989a 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py @@ -28,12 +28,13 @@ from motrix_env_core.renderer import RenderConfig 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.shm import Control, SharedTransitionRing -from motrix_rl.fastsac.async_impl.shm.weight_channel import WeightChannelShared +from motrix_rl.fastsac.async_impl.transport import Control, RingCursors, SharedTransitionRing +from motrix_rl.fastsac.async_impl.transport.weight_channel import WeightChannelShared from motrix_rl.fastsac.async_impl.worker import ( build_env, run_collector_process, run_learner_process, + use_ipc_transition_ring, ) from motrix_rl.fastsac.config import FastSacCfg from motrix_rl.fastsac.wrap import FastSacEnvWrap @@ -131,8 +132,15 @@ def train(self) -> None: 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_envs = self._context.num_envs - ring = SharedTransitionRing(async_options.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim) + if use_ipc_transition_ring(async_options, learner_device, collector_device): + ring: SharedTransitionRing | RingCursors = RingCursors() + else: + ring = SharedTransitionRing(async_options.ring_capacity, num_envs, obs_dim, critic_obs_dim, act_dim) weights = WeightChannelShared(obs_dim=obs_dim) control = Control() diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py similarity index 51% rename from motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py rename to motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py index b6275c77..d30f8bc8 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/__init__.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/__init__.py @@ -1,31 +1,35 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Shared-memory primitives for the async FastSAC trainer. +"""Cross-process transport primitives for the async FastSAC trainer. -This package is the established internal API surface of the former ``shm`` -module; the facade keeps the ``motrix_rl.fastsac.async_impl.shm`` import path -stable across the worker/collector/learner modules and tests. Import from the -defining submodules for new internal uses; import from here only when relying -on the stable package namespace. +This package is the established internal API surface; the facade keeps the +``motrix_rl.fastsac.async_impl.transport`` import path stable across the +worker/collector/learner modules and tests. Import from the defining +submodules for new internal uses; import from here only when relying on the +stable package namespace. Submodules: * :mod:`.common` — the shared-memory allocator, shared scalars (:class:`Control`) and flat-parameter helpers. -* :mod:`.ring` — the SPSC transition ring (collector -> learner). +* :mod:`.ring` — the SPSC transition ring (collector -> learner): + host-shm fields plus the shared :class:`~.ring.RingCursors` protocol. +* :mod:`.ipc_ring` — CUDA-IPC device-fields variant of the ring, used + when collector inference and the learner share one GPU. * :mod:`.weight_channel` — the seqlock weight channel (learner -> collector) with host-shm and CUDA-IPC endpoint implementations. """ -from motrix_rl.fastsac.async_impl.shm.common import ( +from motrix_rl.fastsac.async_impl.transport.common import ( Control, bind_flat_params, flatten_params, load_flat_params, ) -from motrix_rl.fastsac.async_impl.shm.ring import SharedTransitionRing -from motrix_rl.fastsac.async_impl.shm.weight_channel import ( +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 ( GpuIpcWeightReceiver, GpuIpcWeightSender, HostWeightReceiver, @@ -42,6 +46,8 @@ "GpuIpcWeightSender", "HostWeightReceiver", "HostWeightSender", + "IpcTransitionRing", + "RingCursors", "SharedTransitionRing", "WeightChannelShared", "WeightReceiver", diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py similarity index 98% rename from motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py rename to motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py index 39189914..afcf4d35 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/common.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/common.py @@ -12,7 +12,7 @@ * :class:`SharedTransitionRing` — single-producer / single-consumer ring of raw transition batches (collector -> learner) with bounded backpressure. -* the weight channel (see :mod:`motrix_rl.fastsac.async_impl.shm.weight_channel`) +* the weight channel (see :mod:`motrix_rl.fastsac.async_impl.transport.weight_channel`) — double-buffered actor weights + obs-normalizer stats (learner -> collector) guarded by a seqlock so readers always see a complete, consistent snapshot even when the writer publishes twice during a read. diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ipc_ring.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ipc_ring.py new file mode 100644 index 00000000..85b3f87b --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ipc_ring.py @@ -0,0 +1,232 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""CUDA-IPC transition ring: SPSC ring with device-resident slots. + +Used when collector inference and the learner share one GPU. The six +transition fields are fused into a single ``(capacity, num_envs, feat)`` +float32 device tensor (``feat = obs_dim + critic_obs_dim + act_dim + 3``; +reward / done / truncation ride as floats, with 0/1 exactly representable and +converted back to int64 by the replay buffer's ``copy_``). This lets the +collector publish one env-step batch with a **single** H2D copy from one +pinned staging buffer, and lets the learner ingest with pure D2D strided +copies — no host transfer stays on the learner's critical path. + +Ownership and handoff +~~~~~~~~~~~~~~~~~~~~~ +The owner (learner process, which owns the CUDA context) allocates the fused +slot tensor and must keep it alive for the process lifetime. The tensor is +shipped to the collector through the existing one-shot ``slot_queue`` +handshake — ``torch.multiprocessing`` registers CUDA-IPC reducers for CUDA +tensors, so the queue transfer maps the same device memory into the collector +process. The receiver constructs its endpoint from the arrived tensor. Both +sides share the parent-created host :class:`~.ring.RingCursors`. + +Publish ordering (the correctness core) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Host cursors are bumped by CPU stores, but slot data is written by GPU copies +whose completion is NOT ordered by x86/TSO. A cursor may therefore only be +advanced once the corresponding device writes have completed, proven by CUDA +events: + +* producer: after the slot's H2D is enqueued, record a per-slot event; + ``_write`` advances only up to the boundary of completed events (lazily + flushed by :meth:`push` / :meth:`is_full`). The single pinned staging + buffer is likewise not reused until the previous H2D's event completed. +* consumer: ``commit_reads`` records an event after the run's D2D reads are + enqueued; ``_read`` advances only up to the boundary of completed events + (lazily flushed by :meth:`read_span` / :meth:`has_next` / :meth:`size`). + +Both cursor stores remain single-writer aligned int64 host stores, so the +host-side protocol of the shared-memory ring carries over unchanged. +""" + +from __future__ import annotations + +from collections import deque, namedtuple + +import torch + +from motrix_rl.fastsac.async_impl.transport.ring import RingCursors + +# A slot-boundary event and the cursor value it certifies once complete: +# the write (or read) cursor value the entry advances the shared cursor to +# when its event is observed finished. +_PendingEvent = namedtuple("_PendingEvent", ["cursor", "event"]) + + +class IpcTransitionRing: + """SPSC transition ring backed by one fused CUDA-IPC device tensor. + + The producer endpoint is built in the collector process from the shipped + ``slots`` tensor; the consumer endpoint in the learner process may be the + owner (it allocated ``slots``). The two roles use disjoint APIs: + + producer: :meth:`is_full`, :meth:`push` + consumer: :meth:`has_next`, :meth:`read_span`, :meth:`commit_reads` + """ + + def __init__( + self, + cursors: RingCursors, + slots: torch.Tensor, + capacity: int, + num_envs: int, + obs_dim: int, + critic_obs_dim: int, + act_dim: int, + ): + if not slots.is_cuda: + raise ValueError(f"IPC transition ring requires CUDA slots, got device {slots.device}") + expected = (capacity, num_envs, obs_dim + critic_obs_dim + act_dim + 3) + if tuple(slots.shape) != expected: + raise ValueError(f"IPC slot tensor shape {tuple(slots.shape)} does not match {expected}") + self.cursors = cursors + self.slots = slots + self.capacity = capacity + self.num_envs = num_envs + self.obs_dim = obs_dim + self.critic_obs_dim = critic_obs_dim + self.act_dim = act_dim + # Field layout along the fused last dimension. + self._o0 = 0 + self._o1 = obs_dim + self._o2 = obs_dim + critic_obs_dim + self._o3 = self._o2 + act_dim + self._o4 = self._o3 + 1 # reward + self._o5 = self._o4 + 1 # done + # Endpoint-local issued counts. The shared cursors hold only the + # event-proven (published) values, which LAG the issued counts while + # copies are in flight — so slot indexing and publish targets must come + # from the issued counts, never from the cursors. + self._issued_writes = self.cursors.write_idx + self._issued_reads = self.cursors.read_idx + # Producer-side: events of pushed-but-not-yet-published slots, in + # order. (published cursor value, event) — the write cursor value the + # entry certifies once its event completes. Deques: flush pops from + # the head, and list.pop(0) would shift the remaining entries. + self._write_events: deque[_PendingEvent] = deque() + # Consumer-side: events of committed-but-not-yet-advanced reads. + self._read_events: deque[_PendingEvent] = deque() + # Single pinned staging reused by every push; guarded by _staging_event. + self._staging = torch.empty((num_envs, expected[2]), dtype=torch.float32, pin_memory=True) + self._staging_event = torch.cuda.Event() + self._staging_used = False + + # -------------------------------------------------------------- producer + def _flush_writes(self) -> None: + """Advance ``_write`` to the boundary of completed push events.""" + events = self._write_events + while events and events[0].event.query(): + self.cursors._write[0] = events.popleft().cursor + + @property + def write_idx(self) -> int: + self._flush_writes() + return self.cursors.write_idx + + @property + def read_idx(self) -> int: + self._flush_reads() + return self.cursors.read_idx + + def size(self) -> int: + """Published (event-proven) unread count — the observable ring fill.""" + return self.write_idx - self.read_idx + + def is_full(self) -> bool: + # Backpressure must count the producer's issued (in-flight included) + # writes, so a burst of pushes can never exceed capacity. The flush + # publishes completed slots, mirroring the collector's every-step poll. + self._flush_writes() + return self._issued_writes - self.read_idx >= self.capacity + + def push( + self, + obs: torch.Tensor, + critic_obs: torch.Tensor, + actions: torch.Tensor, + rewards: torch.Tensor, + dones: torch.Tensor, + truncations: torch.Tensor, + ) -> bool: + """Publish one env-step batch into the next slot. Returns False if full. + + Inputs are the collector's CPU tensors (``dones``/``truncations`` may + be any numeric dtype; they are stored as 0/1 floats). One fused H2D + lands the whole slot; the cursor publishes lazily once it completes. + """ + if self.is_full(): + return False + if self._staging_used and not self._staging_event.query(): + # The previous H2D still reads the staging buffer; it completes in + # well under one env step, so this is normally a no-op wait. + self._staging_event.synchronize() + s = self._staging + s[:, self._o0 : self._o1].copy_(obs) + s[:, self._o1 : self._o2].copy_(critic_obs) + s[:, self._o2 : self._o3].copy_(actions) + s[:, self._o3 : self._o4].copy_(rewards.reshape(-1, 1)) + s[:, self._o4 : self._o5].copy_(dones.reshape(-1, 1)) + s[:, self._o5 :].copy_(truncations.reshape(-1, 1)) + issued = self._issued_writes + slot = issued % self.capacity + stream = torch.cuda.current_stream(self.slots.device) + self.slots[slot].copy_(s, non_blocking=True) + event = torch.cuda.Event() + event.record(stream) + self._write_events.append(_PendingEvent(issued + 1, event)) + self._issued_writes = issued + 1 + self._staging_event.record(stream) + self._staging_used = True + return True + + # -------------------------------------------------------------- consumer + def _flush_reads(self) -> None: + """Advance ``_read`` to the boundary of completed commit events.""" + events = self._read_events + while events and events[0].event.query(): + self.cursors._read[0] = events.popleft().cursor + + def has_next(self) -> bool: + # Flush this endpoint's commit events so the producer's is_full sees + # freed slots (the learner's every-loop poll). + self._flush_reads() + return self.write_idx > self._issued_reads + + def read_span(self) -> tuple[int, tuple[torch.Tensor, ...]]: + """Longest contiguous unread run, as ``(count, field_views)``. + + The views are strided slices of the fused slot tensor — device tensors + of shape ``(count, num_envs, dim)`` (``(count, num_envs)`` for the + scalar fields). Their data is complete: the producer only publishes a + slot after its H2D event completed. Returns ``(0, ())`` when empty. + """ + self._flush_reads() + available = self.write_idx - self._issued_reads + if available <= 0: + return 0, () + slot = self._issued_reads % self.capacity + count = min(available, self.capacity - slot) + run = self.slots[slot : slot + count] + return count, ( + run[:, :, self._o0 : self._o1], + run[:, :, self._o1 : self._o2], + run[:, :, self._o2 : self._o3], + run[:, :, self._o3], + run[:, :, self._o4], + run[:, :, self._o5], + ) + + def commit_reads(self, n: int) -> None: + """Release ``n`` slots of a consumed run once their reads completed. + + Call AFTER the consumer's D2D copies reading the run have been + enqueued on the current stream: the recorded event orders the cursor + advance behind those copies, so the producer can never overwrite a + slot the learner is still reading. + """ + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(self.slots.device)) + self._read_events.append(_PendingEvent(self._issued_reads + n, event)) + self._issued_reads += n diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ring.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ring.py new file mode 100644 index 00000000..7c168d3e --- /dev/null +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/ring.py @@ -0,0 +1,192 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""SPSC transition ring between the collector and learner processes. + +The ring protocol (cursor semantics, backpressure) is shared by two physical +transports: + +* host shared-memory fields (this module, :class:`SharedTransitionRing`) — + the universal fallback; +* CUDA-IPC device fields (``ipc_ring.py``, :class:`IpcTransitionRing`) — + used when collector inference and the learner share one GPU. + +Both transports advance the same parent-created :class:`RingCursors`, so the +producer/consumer contract below is transport-independent. + +Producer (collector) calls :meth:`push`; when the ring is full it returns +``False`` and the caller must retry/backoff — that is the backpressure that +keeps the collector from outrunning the learner and flooding memory. + +Consumer (learner) calls :meth:`read_span` to get zero-copy views of the +longest contiguous unread run of slots, moves them to its replay buffer, then +calls :meth:`commit_reads`. The read cursor only advances after the copy, so +the producer can never clobber a slot that is still being ingested. + +Memory ordering +~~~~~~~~~~~~~~~ +Only ``_write`` is read by the consumer, only ``_read`` is read by the +producer, and each cursor has a single writer — so an aligned int64 store +is enough to publish progress and no atomic RMW is needed. Correctness also +needs the consumer to not observe the ``_write`` bump before the slot's data +stores have landed (and symmetrically for ``_read``); on x86/TSO that +ordering is free, so no memory barrier is used and this path is x86-only +(see the module "Memory ordering" note). The CUDA-IPC transport instead +orders GPU writes behind per-slot events before bumping a cursor (see +``ipc_ring.py``). +""" + +from __future__ import annotations + +import torch + +from motrix_rl.fastsac.async_impl.transport.common import _shared + + +class RingCursors: + """Host shared-memory read/write cursors shared by both ring transports. + + Created once by the parent process and inherited by both workers. Each + cursor has a single writer (the producer owns ``_write``, the consumer + owns ``_read``), so plain aligned int64 stores publish progress coherently + on x86/TSO without barriers. + """ + + def __init__(self): + self._write = _shared((1,), torch.int64) + self._read = _shared((1,), torch.int64) + + @property + def write_idx(self) -> int: + # Consumer reads this; single-writer producer, so a plain aligned int64 + # load is coherent. On x86/TSO the data loads that follow cannot be + # reordered ahead of it, so no acquire barrier is needed (x86-only). + return int(self._write[0]) + + @property + def read_idx(self) -> int: + # Producer reads this; single-writer consumer. On x86/TSO the following + # is_full() decision is based on an up-to-date value without a barrier. + return int(self._read[0]) + + +# ---------------------------------------------------------------- transition ring +class SharedTransitionRing: + """SPSC ring of transition batches with bounded backpressure (host fields). + + Each slot holds one env-step batch: the six tensors produced by one + collector step, with the leading dimension being ``num_envs`` (so a slot + is ``(num_envs, dim)``). ``next_obs``/``next_critic_obs`` are NOT stored: + with auto-reset envs the observation returned by step ``t`` is exactly the + stored observation of step ``t+1`` (at episode ends it is the reset + observation), so the consumer's replay buffer derives them from the + successive slot. This halves the per-slot copy volume and shared-memory + footprint. + + Fields live in host shared memory, usable by any collector/learner device + combination. See ``ipc_ring.py`` for the CUDA-IPC device-fields variant + sharing the same cursor protocol. + """ + + FIELDS = ( + "obs", + "critic_obs", + "actions", + "rewards", + "dones", + "truncations", + ) + + def __init__( + self, + capacity: int, + num_envs: int, + obs_dim: int, + critic_obs_dim: int, + act_dim: int, + cursors: RingCursors | None = None, + ): + if capacity < 1: + raise ValueError(f"ring_capacity must be >= 1, got {capacity}") + self.capacity = capacity + self.num_envs = num_envs + f32, i64 = torch.float32, torch.int64 + self.obs = _shared((capacity, num_envs, obs_dim), f32) + self.critic_obs = _shared((capacity, num_envs, critic_obs_dim), f32) + self.actions = _shared((capacity, num_envs, act_dim), f32) + self.rewards = _shared((capacity, num_envs), f32) + self.dones = _shared((capacity, num_envs), i64) + self.truncations = _shared((capacity, num_envs), i64) + self.cursors = cursors if cursors is not None else RingCursors() + + @property + def write_idx(self) -> int: + return self.cursors.write_idx + + @property + def read_idx(self) -> int: + return self.cursors.read_idx + + def size(self) -> int: + """Number of unread slots currently buffered.""" + return self.write_idx - self.read_idx + + def is_full(self) -> bool: + return self.size() >= self.capacity + + def push(self, obs, critic_obs, actions, rewards, dones, truncations) -> bool: + """Copy one env-step batch into the next slot. Returns False if full. + + Inputs are CPU tensors shaped ``(num_envs, dim)``; ``dones``/``truncations`` + are int64 to match ``SimpleReplayBuffer.extend`` semantics. + """ + if self.is_full(): + return False + slot = self.write_idx % self.capacity + self.obs[slot].copy_(obs) + self.critic_obs[slot].copy_(critic_obs) + self.actions[slot].copy_(actions) + self.rewards[slot].copy_(rewards) + self.dones[slot].copy_(dones) + self.truncations[slot].copy_(truncations) + # Publish the slot. On x86/TSO the field copies above are guaranteed + # visible before this cursor bump, so a consumer that reads the new + # write_idx also sees the data (x86-only; ARM would need a release here). + self.cursors._write[0] += 1 + return True + + def has_next(self) -> bool: + """Whether at least one unread slot is committed (i.e. readable).""" + return self.size() > 0 + + def read_span(self) -> tuple[int, tuple]: + """Longest contiguous unread run, as ``(count, field_views)``. + + ``field_views`` is the six per-field tensors of shape + ``(count, num_envs, dim)`` (``(count, num_envs)`` for the scalar + fields) — contiguous in the slot (leading) dimension, so the consumer + can move the whole run with one copy per field. Returns ``(0, ())`` + when empty. Does NOT advance the read cursor; free the run with + :meth:`commit_reads` once the data has been copied out. + """ + size = self.size() + if size <= 0: + return 0, () + slot = self.read_idx % self.capacity + count = min(size, self.capacity - slot) + return count, ( + self.obs[slot : slot + count], + self.critic_obs[slot : slot + count], + self.actions[slot : slot + count], + self.rewards[slot : slot + count], + self.dones[slot : slot + count], + self.truncations[slot : slot + count], + ) + + def commit_reads(self, n: int) -> None: + """Advance the read cursor by ``n`` slots of a consumed contiguous run. + + On x86/TSO the consumer's reads complete before this cursor bump, so + the producer's is_full() cannot reuse a slot still being copied out. + """ + self.cursors._read[0] += n diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/weight_channel.py similarity index 99% rename from motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py rename to motrix_rl/src/motrix_rl/fastsac/async_impl/transport/weight_channel.py index 6dbac0a5..fd4a4ba6 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/shm/weight_channel.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/transport/weight_channel.py @@ -29,7 +29,7 @@ *even* when done; readers retry on any inconsistency. Single writer + single reader means plain aligned int64 stores suffice and no atomic RMW is needed. On x86/TSO the data-before-counter ordering is free; ARM would need real -release/acquire (see the shm module's "Memory ordering" note). +release/acquire (see transport.common's "Memory ordering" note). Transport choice ---------------- @@ -50,7 +50,7 @@ import torch.multiprocessing # noqa: F401 registers CUDA-IPC reducers in every importing process from torch import nn -from motrix_rl.fastsac.async_impl.shm.common import _NORM_KEYS, _shared, flatten_params, load_flat_params +from motrix_rl.fastsac.async_impl.transport.common import _NORM_KEYS, _shared, flatten_params, load_flat_params # obs-normalizer stat buffers that the collector needs (read-only) to reproduce # the sync ``act()`` path (normalize with update=False, see agent.act). 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 1dbea998..86a13e63 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -38,8 +38,13 @@ 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.shm import Control, SharedTransitionRing -from motrix_rl.fastsac.async_impl.shm.weight_channel import ( +from motrix_rl.fastsac.async_impl.transport import ( + Control, + IpcTransitionRing, + RingCursors, + SharedTransitionRing, +) +from motrix_rl.fastsac.async_impl.transport.weight_channel import ( GpuIpcWeightSender, HostWeightSender, WeightChannelShared, @@ -50,7 +55,13 @@ 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 +from motrix_rl.system_metrics import ( + CpuLoadSampler, + GpuMemoryUsageSampler, + GpuUtilizationSampler, + MemoryUsageSampler, + sample_gpu_devices, +) def _timing_mean(values: list[float]) -> float: @@ -146,6 +157,59 @@ def build_agent(cfg: FastSacCfg, dims, num_envs, device, action_scale, action_bi ) +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).""" @@ -277,14 +341,7 @@ def _build_weight_sender( 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 = device.type == "cuda" and collector_device.type == "cuda" - if same_gpu: - # ``learner=`` without an index means the current device; resolve it so - # the comparison never treats "cuda" as matching an explicit different - # index. Only resolve under the cuda branch: a CPU learner has no CUDA - # context and torch.cuda.current_device() would raise. - learner_index = device.index if device.index is not None else torch.cuda.current_device() - same_gpu = learner_index == collector_device.index + same_gpu = same_cuda_device(device, collector_device) if mode == "on" and not same_gpu: reason = ( "collector inference device is not CUDA" @@ -310,7 +367,7 @@ def run_collector_process( dims: tuple[int, int, int], action_scale: torch.Tensor, action_bias: torch.Tensor, - ring: SharedTransitionRing, + ring: SharedTransitionRing | RingCursors, weights: WeightChannelShared, control: Control, stats_queue: Queue, @@ -330,17 +387,24 @@ def run_collector_process( obs_dim, critic_obs_dim, act_dim = dims device = torch.device("cpu") env = build_env(env_spec, num_envs, device, seed=seed) - # Handshake: build the receiver from the slot tensors the learner - # shipped (host shm or CUDA-IPC), before the collector is wired up. + # 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: - slots = slot_queue.get(timeout=60.0) + 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 weight-slot tensors " + "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 - weight_rx = weight_receiver_for(weights, slots) + 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 + ) collector = Collector( env, cfg, @@ -384,7 +448,7 @@ def run_learner_process( dims: tuple[int, int, int], action_scale: torch.Tensor, action_bias: torch.Tensor, - ring: SharedTransitionRing, + ring: SharedTransitionRing | RingCursors, weights: WeightChannelShared, control: Control, stats_queue: Queue, @@ -423,9 +487,23 @@ def run_learner_process( # 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. + # 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) - slot_queue.put(weight_tx.params) + 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 + ) + 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 @@ -536,8 +614,9 @@ def _drain_stats(): # 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.read``); nesting is rebuilt with one rule, - # and a stage's own total folds into its node. + # ``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) @@ -585,6 +664,7 @@ def _drain_stats(): 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") @@ -622,6 +702,7 @@ def _drain_stats(): if 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" torch.save(agent.state_dict(), path) checkpoints.record_checkpoint_artifact( @@ -639,6 +720,7 @@ def _drain_stats(): # 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) diff --git a/motrix_rl/src/motrix_rl/fastsac/buffer.py b/motrix_rl/src/motrix_rl/fastsac/buffer.py index ad39b3e1..58eaa4de 100644 --- a/motrix_rl/src/motrix_rl/fastsac/buffer.py +++ b/motrix_rl/src/motrix_rl/fastsac/buffer.py @@ -90,6 +90,36 @@ def extend(self, obs, critic_obs, actions, rewards, dones, truncations) -> None: self.truncations[:, slot] = truncations self.ptr += 1 + def extend_batch(self, obs, critic_obs, actions, rewards, dones, truncations) -> None: + """Append ``k`` transitions per environment from slot-major tensors. + + Inputs are shaped ``(k, n_env, dim)`` for obs / critic_obs / actions + and ``(k, n_env)`` for the scalar fields — the contiguous multi-slot + layout produced by the async transition ring's :meth:`read_span`. + Writes go directly into the strided buffer slices (one copy per field + per contiguous chunk instead of one per slot) on the caller's current + stream, so an async ingest can issue non-blocking pinned H2D copies on + a side stream. + """ + k = obs.shape[0] + start = self.ptr % self._cap + n = min(k, self._cap - start) + self._write_span(start, (t[:n] for t in (obs, critic_obs, actions, rewards, dones, truncations))) + if k > n: + self._write_span(0, (t[n:] for t in (obs, critic_obs, actions, rewards, dones, truncations))) + self.ptr += k + + def _write_span(self, slot: int, fields) -> None: + obs, critic_obs, actions, rewards, dones, truncations = fields + n = obs.shape[0] + end = slot + n + self.observations[:, slot:end].copy_(obs.permute(1, 0, 2), non_blocking=True) + self.critic_observations[:, slot:end].copy_(critic_obs.permute(1, 0, 2), non_blocking=True) + self.actions[:, slot:end].copy_(actions.permute(1, 0, 2), non_blocking=True) + self.rewards[:, slot:end].copy_(rewards.permute(1, 0), non_blocking=True) + self.dones[:, slot:end].copy_(dones.permute(1, 0), non_blocking=True) + self.truncations[:, slot:end].copy_(truncations.permute(1, 0), non_blocking=True) + @torch.no_grad() def sample(self, batch_size: int) -> dict: if self.num_stored == 0: diff --git a/motrix_rl/src/motrix_rl/fastsac/config.py b/motrix_rl/src/motrix_rl/fastsac/config.py index 2f25e167..e5dfcd35 100644 --- a/motrix_rl/src/motrix_rl/fastsac/config.py +++ b/motrix_rl/src/motrix_rl/fastsac/config.py @@ -77,6 +77,14 @@ class FastSacAsyncOptionsCfg: # collector. learner_cpu_cores: str | None = None collector_cpu_cores: str | None = None + # Transition-ring transport between the async collector and learner: + # "auto" places the ring slots in CUDA-IPC device memory when learner and + # collector inference share one GPU (single fused H2D on the collector, + # D2D-only learner ingest) and falls back to host shared memory otherwise; + # "on"/"off" force the device/host path ("on" warns and falls back when + # the two sides are not on the same GPU). Keep the values quoted: unquoted + # on/off parse as booleans in YAML. + transition_ipc: str = "auto" # Weight-snapshot transport between the async learner and collector: # "auto" picks CUDA-IPC device slots when both sides share one GPU and the # actor parameters are large enough (>= weight_ipc_min_bytes) for the diff --git a/motrix_rl/src/motrix_rl/fastsac/sync/train.py b/motrix_rl/src/motrix_rl/fastsac/sync/train.py index 98da5cd2..3c456f28 100644 --- a/motrix_rl/src/motrix_rl/fastsac/sync/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/sync/train.py @@ -20,7 +20,13 @@ from motrix_rl.fastsac.wrap_np import FastSacNpEnvWrap from motrix_rl.fastsac.wrap_torch import FastSacTorchEnvWrap from motrix_rl.frameworks import TrainerBase, TrainerContext -from motrix_rl.system_metrics import CpuLoadSampler, GpuMemoryUsageSampler, GpuUtilizationSampler, MemoryUsageSampler +from motrix_rl.system_metrics import ( + CpuLoadSampler, + GpuMemoryUsageSampler, + GpuUtilizationSampler, + MemoryUsageSampler, + sample_gpu_devices, +) # Enable TF32 matmul on Ampere+ GPUs. SAC training has no precision concern with # TF32 (10 mantissa bits), and the speedup is meaningful when AMP is off. @@ -328,6 +334,7 @@ def emit_msg(msg: str) -> None: 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"{self._env_name}/motrix.fastsac") diff --git a/motrix_rl/src/motrix_rl/system_metrics.py b/motrix_rl/src/motrix_rl/system_metrics.py index 0fbb687d..b840b4d9 100644 --- a/motrix_rl/src/motrix_rl/system_metrics.py +++ b/motrix_rl/src/motrix_rl/system_metrics.py @@ -25,7 +25,10 @@ class CpuLoad: Utilization excludes idle, I/O-wait, and stolen virtual-CPU time. The equivalent logical CPU count makes the normalized percentage unambiguous - on SMT systems. + on SMT systems. ``per_core_percent`` carries the same utilization broken + down per logical CPU (sorted by cpu id) for the system view's per-core + display; it comes free — the sampler already reads per-CPU counters to + compute the aggregate. """ utilization_percent: float @@ -34,6 +37,9 @@ class CpuLoad: physical_core_count: int | None iowait_percent: float steal_percent: float + per_core_percent: tuple[float, ...] | None = None + # Static "model name" from /proc/cpuinfo (Linux); None where unavailable. + model_name: str | None = None @dataclass(frozen=True) @@ -52,12 +58,14 @@ def __init__( *, stat_path: str | Path = "/proc/stat", topology_root: str | Path = "/sys/devices/system/cpu", + cpuinfo_path: str | Path = "/proc/cpuinfo", cpu_ids: set[int] | None = None, ) -> None: self._stat_path = Path(stat_path) self._topology_root = Path(topology_root) self._cpu_ids = cpu_ids if cpu_ids is not None else self._available_cpu_ids() self._physical_core_count = self._read_physical_core_count() + self._model_name = self._read_model_name(Path(cpuinfo_path)) self._previous = self._read_times() def sample(self) -> CpuLoad | None: @@ -68,16 +76,23 @@ def sample(self) -> CpuLoad | None: self._previous = current return None - total = sum(current[cpu].total - self._previous[cpu].total for cpu in common_ids) - executing = sum(current[cpu].executing - self._previous[cpu].executing for cpu in common_ids) - iowait = sum(current[cpu].iowait - self._previous[cpu].iowait for cpu in common_ids) - steal = sum(current[cpu].steal - self._previous[cpu].steal for cpu in common_ids) + previous = self._previous + total = sum(current[cpu].total - previous[cpu].total for cpu in common_ids) + executing = sum(current[cpu].executing - previous[cpu].executing for cpu in common_ids) + iowait = sum(current[cpu].iowait - previous[cpu].iowait for cpu in common_ids) + steal = sum(current[cpu].steal - previous[cpu].steal for cpu in common_ids) self._previous = current if total <= 0: return None logical_cpu_count = len(common_ids) utilization_percent = 100.0 * executing / total + per_core = tuple( + 100.0 + * (current[cpu].executing - previous[cpu].executing) + / max(current[cpu].total - previous[cpu].total, 1) + for cpu in sorted(common_ids) + ) return CpuLoad( utilization_percent=utilization_percent, used_logical_cpus=logical_cpu_count * utilization_percent / 100.0, @@ -85,6 +100,8 @@ def sample(self) -> CpuLoad | None: physical_core_count=self._physical_core_count, iowait_percent=100.0 * iowait / total, steal_percent=100.0 * steal / total, + per_core_percent=per_core, + model_name=self._model_name, ) @staticmethod @@ -119,6 +136,17 @@ def _read_times(self) -> dict[int, _CpuTimes]: ) return times + def _read_model_name(self, cpuinfo_path: Path) -> str | None: + """First ``model name`` entry from /proc/cpuinfo; None off-Linux or unreadable.""" + try: + for line in cpuinfo_path.read_text().splitlines(): + if line.startswith("model name"): + _, _, value = line.partition(":") + return value.strip() or None + except OSError: + return None + return None + def _read_physical_core_count(self) -> int | None: cores: set[tuple[int, int]] = set() try: @@ -140,6 +168,18 @@ class MemoryUsage: total_bytes: int +@dataclass(frozen=True) +class GpuDeviceUsage: + """One accelerator's utilization and memory, identified by device index.""" + + index: int + utilization_percent: float | None + memory: MemoryUsage | None + # Marketing/model name from NVML or AMD SMI; None when the backend or + # driver does not expose one. + name: str | None = None + + class _MemoryStatusEx(ctypes.Structure): """``MEMORYSTATUSEX`` layout for the Windows ``GlobalMemoryStatusEx`` call.""" @@ -226,72 +266,153 @@ def _nvml() -> tuple[Any, list[Any]] | None: class GpuMemoryUsageSampler: - """Read aggregate accelerator memory via AMD SMI or NVML.""" + """Read accelerator memory via AMD SMI or NVML, per device or summed.""" - def sample(self) -> MemoryUsage | None: + def sample_per_device(self) -> list[MemoryUsage] | None: + """One :class:`MemoryUsage` per device (backend enumeration order).""" session = _amd_smi() if session is not None: amdsmi, devices = session - used = total = 0 + usages: list[MemoryUsage] = [] try: for device in devices: - used += int(amdsmi.amdsmi_get_gpu_memory_usage(device, amdsmi.AmdSmiMemoryType.VRAM)) - total += int(amdsmi.amdsmi_get_gpu_memory_total(device, amdsmi.AmdSmiMemoryType.VRAM)) + used = int(amdsmi.amdsmi_get_gpu_memory_usage(device, amdsmi.AmdSmiMemoryType.VRAM)) + total = int(amdsmi.amdsmi_get_gpu_memory_total(device, amdsmi.AmdSmiMemoryType.VRAM)) + usages.append(MemoryUsage(used_bytes=used, total_bytes=total)) except Exception: return None - return MemoryUsage(used_bytes=used, total_bytes=total) if total > 0 else None + return usages session = _nvml() if session is None: return None pynvml, devices = session - used = total = 0 try: - for device in devices: - info = pynvml.nvmlDeviceGetMemoryInfo(device) - used += info.used - total += info.total + return [ + MemoryUsage(used_bytes=info.used, total_bytes=info.total) + for info in (pynvml.nvmlDeviceGetMemoryInfo(device) for device in devices) + ] except pynvml.NVMLError: return None + + def sample(self) -> MemoryUsage | None: + usages = self.sample_per_device() + if not usages: + return None + used = sum(usage.used_bytes for usage in usages) + total = sum(usage.total_bytes for usage in usages) return MemoryUsage(used_bytes=used, total_bytes=total) if total > 0 else None class GpuUtilizationSampler: - """Read aggregate accelerator utilization via AMD SMI or NVML.""" + """Read accelerator utilization via AMD SMI or NVML, per device or averaged.""" - def sample(self) -> float | None: + def sample_per_device(self) -> list[float | None] | None: + """One utilization percentage per device, or ``None`` per unavailable device. + + Returns ``None`` when no per-device source exists (e.g. the AMD sysfs + fallback reports unlabeled cards); callers fall back to the aggregate. + """ global _amd_activity_supported session = _amd_smi() if session is not None: amdsmi, devices = session - values: list[float] = [] - if _amd_activity_supported is not False: - try: - for device in devices: - activity = amdsmi.amdsmi_get_gpu_activity(device) - value = activity.get("gfx_activity", activity.get("gpu_busy_percent")) - if value is not None: - values.append(float(value)) - except Exception: - _amd_activity_supported = False if _amd_activity_supported is False: - # Some integrated AMD GPUs (including Radeon 890M) expose - # VRAM through AMD SMI but return AMDSMI_STATUS_UNEXPECTED_DATA - # for ``amdsmi_get_gpu_activity``. The kernel's DRM sysfs - # counter is available on those devices and reports the same - # busy percentage used by rocm-smi. - values = _sysfs_gpu_busy_percent() - return sum(values) / len(values) if values else None + return None + values: list[float | None] = [] + try: + for device in devices: + activity = amdsmi.amdsmi_get_gpu_activity(device) + value = activity.get("gfx_activity", activity.get("gpu_busy_percent")) + values.append(float(value) if value is not None else None) + except Exception: + _amd_activity_supported = False + return None + return values session = _nvml() if session is None: return None pynvml, devices = session - values: list[int] = [] try: - for device in devices: - values.append(pynvml.nvmlDeviceGetUtilizationRates(device).gpu) + return [float(pynvml.nvmlDeviceGetUtilizationRates(device).gpu) for device in devices] except pynvml.NVMLError: return None - return sum(values) / len(values) if values else None + + def sample(self) -> float | None: + amd_session = _amd_smi() is not None + values = self.sample_per_device() + if values is None: + if not amd_session: + return None + # Some integrated AMD GPUs (including Radeon 890M) expose VRAM + # through AMD SMI but return AMDSMI_STATUS_UNEXPECTED_DATA for + # ``amdsmi_get_gpu_activity``. The kernel's DRM sysfs counter is + # available on those devices and reports the same busy percentage + # used by rocm-smi (per-card but unlabeled, hence aggregate-only). + sysfs = _sysfs_gpu_busy_percent() + return sum(sysfs) / len(sysfs) if sysfs else None + known = [value for value in values if value is not None] + return sum(known) / len(known) if known else None + + +def _gpu_device_names() -> list[str | None] | None: + """One name per device (backend enumeration order), or ``None``. + + Names are static, so the query result is cached alongside the backend + session state. AMD SMI product info varies by driver; unrecognized + shapes degrade to a ``None`` entry rather than failing the whole list. + """ + session = _amd_smi() + if session is not None: + amdsmi, devices = session + names: list[str | None] = [] + try: + for device in devices: + info = amdsmi.amdsmi_get_processor_info(device) + # Parenthesized: the isinstance guard selects the whole `or` + # chain (conditional expressions bind loosest), so a non-dict + # info degrades to None instead of touching .get(). + name = (info.get("market_name") or info.get("product_name")) if isinstance(info, dict) else None + names.append(str(name) if name else None) + except Exception: + return None + return names + session = _nvml() + if session is None: + return None + pynvml, devices = session + try: + # Older pynvml returns bytes; modern versions return str. Any failure + # (unsupported handle, driver quirk) degrades to unnamed rows. + names = [pynvml.nvmlDeviceGetName(device) for device in devices] + return [name.decode() if isinstance(name, bytes) else name for name in names] + except Exception: + return None + + +def sample_gpu_devices( + utilization: GpuUtilizationSampler, + memory: GpuMemoryUsageSampler, +) -> list[GpuDeviceUsage] | None: + """Combine the two samplers into one per-device usage list, or ``None``. + + Devices come from the same backend enumeration in both samplers, so + indices pair up. A failing memory source degrades to utilization-only; + a failing name source degrades to unnamed rows. + """ + utils = utilization.sample_per_device() + if utils is None: + return None + memories = memory.sample_per_device() + names = _gpu_device_names() + return [ + GpuDeviceUsage( + index=index, + utilization_percent=value, + memory=memories[index] if memories is not None and index < len(memories) else None, + name=names[index] if names is not None and index < len(names) else None, + ) + for index, value in enumerate(utils) + ] def _sysfs_gpu_busy_percent() -> list[float]: diff --git a/motrix_rl/tests/test_console.py b/motrix_rl/tests/test_console.py index 9ab33652..c64c0f07 100644 --- a/motrix_rl/tests/test_console.py +++ b/motrix_rl/tests/test_console.py @@ -168,10 +168,10 @@ def _panel_stats(**overrides: Any) -> TrainingPanelStats: return TrainingPanelStats(**values) -def _render_panel(stats: TrainingPanelStats, *, detail: bool = False, width: int = 200) -> str: +def _render_panel(stats: TrainingPanelStats, *, view: str = "overview", width: int = 200) -> str: console = Console(width=width) with console.capture() as capture: - console.print(render_training_panel(stats, detail=detail)) + console.print(render_training_panel(stats, view=view)) return capture.get() @@ -218,7 +218,7 @@ def test_render_training_panel_detail_view_shows_timing_tree_with_shares() -> No ) overview = _render_panel(stats) - detail = _render_panel(stats, detail=True) + detail = _render_panel(stats, view="timing") assert "STAGE" not in overview assert "STAGE" in detail @@ -258,6 +258,107 @@ def test_render_training_panel_reports_system_health() -> None: assert "VRAM n/a" in panel +def test_overview_aggregates_gpu_devices_and_system_view_shows_per_gpu() -> None: + from motrix_rl.system_metrics import GpuDeviceUsage + + stats = _panel_stats( + gpu_devices=[ + GpuDeviceUsage( + index=0, + utilization_percent=85.0, + memory=MemoryUsage(used_bytes=1024**3, total_bytes=2 * 1024**3), + name="NVIDIA GeForce RTX 4090", + ), + GpuDeviceUsage(index=1, utilization_percent=40.0, memory=None), + ] + ) + + overview = _render_panel(stats, view="overview") + + # Overview/timing cards stay aggregate-only (mean util, summed VRAM) + assert "GPU 62%" in overview + assert "VRAM 1.0/2.0 GiB" in overview + assert "GPU0" not in overview + assert "GPU1" not in overview + + from motrix_rl.system_metrics import CpuLoad + + stats = _panel_stats( + gpu_devices=stats.gpu_devices, + cpu_load=CpuLoad( + utilization_percent=62.5, + used_logical_cpus=125.0, + logical_cpu_count=200, + physical_core_count=100, + iowait_percent=0.0, + steal_percent=0.0, + per_core_percent=tuple(100.0 if i % 50 == 0 else 0.0 for i in range(200)), + model_name="AMD EPYC 9654 96-Core Processor", + ), + ) + + system = _render_panel(stats, view="system") + + # The CPU card leads with the host's model name + assert "EPYC 9654" in system + + # The dedicated system view lists every device, with its model name + assert "GPU0" in system and "85%" in system + assert "RTX 4090" in system + assert "GPU1" in system and "40%" in system + assert "n/a" in system + assert "1.0/2.0 GiB" in system + # per-core spectrum: one glyph per logical core, wrapped at 48 per row + # (48 glyphs + 47 inter-core gaps span the same width as the old 96-glyph row) + assert system.count("cores 0-47") == 1 + assert system.count("cores 48-95") == 1 + assert system.count("cores 192-199") == 1 + # default style is the partial-height glyph spectrum: idle cores show the + # lowest glyph, busy ones the full block, and no bitmap rows appear + assert "▁" in system + assert "▄" not in system + # a dim ceiling line marks each column's 100% reference: lower-eighth + # blocks on the row above, touching full-height columns below + assert "▔" not in system + # overview shows none of the per-core detail + assert "cores" not in _render_panel(stats, view="overview") + + +def test_render_training_panel_falls_back_to_aggregate_gpu_without_device_stats() -> None: + stats = _panel_stats(gpu_utilization_percent=85.0) + + panel = _render_panel(stats) + + assert "GPU 85%" in panel + assert "VRAM n/a" in panel + + +def test_cpu_spectrum_style_env_override(monkeypatch) -> None: + from motrix_rl.system_metrics import CpuLoad + + stats = _panel_stats( + cpu_load=CpuLoad( + utilization_percent=50.0, + used_logical_cpus=1.0, + logical_cpu_count=2, + physical_core_count=1, + iowait_percent=0.0, + steal_percent=0.0, + per_core_percent=(100.0, 25.0), + ) + ) + + monkeypatch.setenv("MOTRIX_PANEL_CPU_SPECTRUM", "bitmap") + bitmap = _render_panel(stats, view="system") + assert "▄" not in bitmap + assert bitmap.count("cores 0-1") == 1 + # bitmap mode: 4 stacked rows for one chunk + assert sum(1 for line in bitmap.splitlines() if "█" in line and "cores" not in line) >= 3 + + monkeypatch.setenv("MOTRIX_PANEL_CPU_SPECTRUM", "bogus") + assert "▂" in _render_panel(stats, view="system") # unknown value falls back to height + + def test_format_memory_renders_gib_and_missing_values() -> None: assert _format_memory(None) == "n/a" assert _format_memory(MemoryUsage(used_bytes=1024**3, total_bytes=4 * 1024**3)) == "1.0/4.0 GiB" @@ -292,7 +393,7 @@ def test_render_training_panel_only_advertises_keyboard_on_posix_tty(monkeypatch # 1/2 key handling needs a POSIX TTY; other platforms get a plain Live stats = _panel_stats() monkeypatch.setattr(console_module, "_POSIX_TTY", True) - assert "keyboard: 1/2 switch tabs" in _render_panel(stats) + assert "keyboard: 1/2/3 switch tabs" in _render_panel(stats) monkeypatch.setattr(console_module, "_POSIX_TTY", False) assert "keyboard" not in _render_panel(stats) diff --git a/motrix_rl/tests/test_fastsac_buffer.py b/motrix_rl/tests/test_fastsac_buffer.py index c9292127..7cb272cc 100644 --- a/motrix_rl/tests/test_fastsac_buffer.py +++ b/motrix_rl/tests/test_fastsac_buffer.py @@ -163,3 +163,89 @@ def test_partial_buffer_n_step_only_full_windows(): t0 = int(batch["obs"][row, 0]) assert t0 <= total - 3 assert int(batch["effective_n_steps"][row]) == 3 + + +def test_extend_batch_matches_extend_across_wrap(): + """Batched multi-slot ingest must be indistinguishable from per-slot extend. + + Covers chunk sizes that split the (k, n_env, ·) span across the buffer's + physical wrap and spans longer than the capacity. + """ + for total, chunk in ( + (BUFFER_SIZE * 3, 2), + (BUFFER_SIZE * 3, BUFFER_SIZE - 1), + (BUFFER_SIZE * 4 + 3, BUFFER_SIZE + 2), + ): + per_slot = SimpleReplayBuffer(N_ENV, BUFFER_SIZE, N_OBS, N_ACT, N_CRITIC_OBS, device="cpu") + batched = SimpleReplayBuffer(N_ENV, BUFFER_SIZE, N_OBS, N_ACT, N_CRITIC_OBS, device="cpu") + dones = [1 if t % 4 == 3 else 0 for t in range(total)] + truncs = [1 if t % 5 == 4 else 0 for t in range(total)] + _fill(per_slot, total, dones, truncs) + for start in range(0, total, chunk): + n = min(chunk, total - start) + fields = ( + torch.full((n, N_ENV, N_OBS), float(0)), + torch.full((n, N_ENV, N_CRITIC_OBS), float(0)), + torch.full((n, N_ENV, N_ACT), float(0)), + torch.full((n, N_ENV), float(0)), + torch.zeros(n, N_ENV, dtype=torch.long), + torch.zeros(n, N_ENV, dtype=torch.long), + ) + for i, t in enumerate(range(start, start + n)): + fields[0][i] = float(t) + fields[1][i] = float(100 + t) + fields[2][i] = float(t) + fields[3][i] = float(t) + fields[4][i] = dones[t] + fields[5][i] = truncs[t] + batched.extend_batch(*fields) + assert batched.ptr == per_slot.ptr + for name in ("observations", "critic_observations", "actions", "rewards", "dones", "truncations"): + assert torch.equal(getattr(batched, name), getattr(per_slot, name)) + + +def test_ring_read_span_fifo_and_no_wrap(): + """Contiguous-run consumption is FIFO, in-order, and never crosses the wrap.""" + from motrix_rl.fastsac.async_impl.transport import SharedTransitionRing + + capacity, num_envs, obs_dim, act_dim = 5, 3, 4, 2 + + def push(ring, t): + return ring.push( + torch.full((num_envs, obs_dim), float(t)), + torch.full((num_envs, obs_dim), float(t)), + torch.full((num_envs, act_dim), float(t)), + torch.full((num_envs,), float(t)), + torch.full((num_envs,), t % 2, dtype=torch.long), + torch.full((num_envs,), t % 3 == 0, dtype=torch.long), + ) + + ring = SharedTransitionRing(capacity, num_envs, obs_dim, obs_dim, act_dim) + next_expected = 0 + for t in range(capacity * 4): + assert push(ring, t) + # consume at varying granularity; each run's views must carry the + # oldest unread values in push order + k = min((t % 3) + 1, ring.size()) + if k: + count, views = ring.read_span() + assert count >= k + assert 1 <= count <= capacity - (ring.read_idx % capacity) # no wrap + for i in range(k): + assert torch.all(views[0][i] == next_expected + i) + assert torch.all(views[1][i] == next_expected + i) # critic dim reuse ok + assert torch.all(views[2][i] == next_expected + i) + assert torch.all(views[3][i] == next_expected + i) + assert torch.all(views[4][i] == (next_expected + i) % 2) + assert torch.all(views[5][i] == ((next_expected + i) % 3 == 0)) + ring.commit_reads(k) + next_expected += k + # drain the remainder + while ring.has_next(): + count, views = ring.read_span() + for i in range(count): + assert torch.all(views[0][i] == next_expected + i) + ring.commit_reads(count) + next_expected += count + assert next_expected == capacity * 4 + assert ring.size() == 0 diff --git a/motrix_rl/tests/test_fastsac_collector.py b/motrix_rl/tests/test_fastsac_collector.py index c90a42b7..a40d36d3 100644 --- a/motrix_rl/tests/test_fastsac_collector.py +++ b/motrix_rl/tests/test_fastsac_collector.py @@ -10,8 +10,12 @@ from motrix_env_core.perf import Perf from motrix_rl.fastsac.async_impl.collector import Collector, resolve_collector_inference_device -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing -from motrix_rl.fastsac.async_impl.shm.weight_channel import HostWeightReceiver, HostWeightSender, WeightChannelShared +from motrix_rl.fastsac.async_impl.transport import Control, SharedTransitionRing +from motrix_rl.fastsac.async_impl.transport.weight_channel import ( + HostWeightReceiver, + HostWeightSender, + WeightChannelShared, +) from motrix_rl.fastsac.buffer import EmpiricalNormalization from motrix_rl.fastsac.networks import Actor @@ -245,9 +249,10 @@ def test_collector_reports_env_step_substage_timing() -> None: assert "env_step" in stats["timing_ms"] assert "env_step.apply_action" in stats["timing_ms"] assert "env_step.physics" in stats["timing_ms"] - # nested sub-stages arrive as dotted paths for the panel's tree rebuild - assert "env_step.physics.read" in stats["timing_ms"] assert stats["timing_ms"]["env_step.apply_action"] >= 0.0 + # only the first sub-stage level is reported — a parent's total already + # includes its children, deeper dotted paths are dropped at the stats layer + assert not any(key.count(".") > 1 for key in stats["timing_ms"]) # sub-stage aggregation is windowed like the other timings assert env.perf.snapshot() == () diff --git a/motrix_rl/tests/test_fastsac_ipc_ring.py b/motrix_rl/tests/test_fastsac_ipc_ring.py new file mode 100644 index 00000000..6ad61e26 --- /dev/null +++ b/motrix_rl/tests/test_fastsac_ipc_ring.py @@ -0,0 +1,287 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Correctness tests for the CUDA-IPC transition ring. + +Single-process owner + receiver pairs share one CUDA tensor (the in-process +stand-in for the cross-process IPC mapping), so every protocol property under +test — fused-slot layout, event-ordered cursor publishing, lazy commit, +backpressure, replay-buffer equivalence with the host ring — behaves exactly +as it does between the two worker processes. + +GPU-less environments skip the ring tests; the host ring keeps its own tests +in ``test_fastsac_buffer.py``. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from motrix_rl.fastsac.async_impl.transport import ( + IpcTransitionRing, + RingCursors, + SharedTransitionRing, +) +from motrix_rl.fastsac.buffer import SimpleReplayBuffer + +CAPACITY, N_ENV, OBS, CRI, ACT = 5, 3, 4, 6, 2 +FEAT = OBS + CRI + ACT + 3 + +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="IPC ring requires a GPU") + + +def _batch(t: int, seed: int = 0): + """Deterministic per-step field tensors (the collector's CPU outputs).""" + g = torch.Generator().manual_seed(seed * 1000 + t) + return ( + torch.randn(N_ENV, OBS, generator=g), + torch.randn(N_ENV, CRI, generator=g), + torch.randn(N_ENV, ACT, generator=g), + torch.randn(N_ENV, generator=g), + torch.randint(0, 2, (N_ENV,), generator=g, dtype=torch.long), + torch.randint(0, 2, (N_ENV,), generator=g, dtype=torch.long), + ) + + +def _owner_receiver(): + cursors = RingCursors() + slots = torch.zeros(CAPACITY, N_ENV, FEAT, dtype=torch.float32, device="cuda") + owner = IpcTransitionRing(cursors, slots, CAPACITY, N_ENV, OBS, CRI, ACT) + receiver = IpcTransitionRing(cursors, slots, CAPACITY, N_ENV, OBS, CRI, ACT) + return owner, receiver + + +@cuda_only +def test_ipc_ring_rejects_bad_slots(): + cursors = RingCursors() + with pytest.raises(ValueError, match="CUDA"): + IpcTransitionRing(cursors, torch.zeros(CAPACITY, N_ENV, FEAT), CAPACITY, N_ENV, OBS, CRI, ACT) + with pytest.raises(ValueError, match="shape"): + IpcTransitionRing( + cursors, torch.zeros(CAPACITY, N_ENV, FEAT + 1, device="cuda"), CAPACITY, N_ENV, OBS, CRI, ACT + ) + + +@cuda_only +def test_ipc_ring_roundtrip_matches_pushed_data(): + """Every field survives the fused H2D roundtrip, in FIFO order, across wraps.""" + owner, receiver = _owner_receiver() + pushed = [_batch(t, seed=7) for t in range(CAPACITY * 3)] + expected = 0 + + def drain_and_verify() -> None: + nonlocal expected + # Emulate the cross-process roles: only the producer's own calls flush + # its publish events (the collector polls is_full every step), so nudge + # the owner side before the receiver looks. + owner.size() + while receiver.has_next(): + k, views = receiver.read_span() + for i in range(k): + f = pushed[expected + i] + torch.testing.assert_close(views[0][i], f[0].cuda()) + torch.testing.assert_close(views[1][i], f[1].cuda()) + torch.testing.assert_close(views[2][i], f[2].cuda()) + torch.testing.assert_close(views[3][i], f[3].cuda()) + torch.testing.assert_close(views[4][i].long().cpu(), f[4]) + torch.testing.assert_close(views[5][i].long().cpu(), f[5]) + receiver.commit_reads(k) + expected += k + + for fields in pushed: + while not owner.push(*fields): + drain_and_verify() # free slots (the learner's role) and retry + drain_and_verify() + torch.cuda.synchronize() # land the last push's event before the final drain + drain_and_verify() + assert expected == len(pushed) + + +@cuda_only +def test_ipc_ring_cursor_publishes_only_after_event(): + """A pushed slot is invisible until its event completes; size() flushes.""" + owner, receiver = _owner_receiver() + fields = _batch(0, seed=1) + assert owner.push(*fields) + torch.cuda.synchronize() # land the H2D; the cursor must catch up lazily + assert owner.size() == 1 + assert receiver.has_next() + + +@cuda_only +def test_ipc_ring_backpressure_bounds_in_flight(): + owner, receiver = _owner_receiver() + for t in range(CAPACITY): + assert owner.push(*_batch(t, seed=3)) + assert owner.is_full() + assert not owner.push(*_batch(CAPACITY, seed=3)) + # consumer frees everything -> producer can push again (across the wrap) + while receiver.has_next(): + k, _ = receiver.read_span() + receiver.commit_reads(k) + assert not owner.is_full() + assert owner.push(*_batch(CAPACITY, seed=3)) + + +@cuda_only +def test_ipc_ring_back_to_back_pushes_do_not_collide(): + """Regression: pushes issued while the cursor lags (events in flight) must + still target distinct slots and publish to distinct cursor values.""" + owner, receiver = _owner_receiver() + pushed = [_batch(t, seed=9) for t in range(CAPACITY)] + for fields in pushed: # NO producer-side flush between pushes + assert owner.push(*fields) + torch.cuda.synchronize() + owner.size() + expected = 0 + while receiver.has_next(): + k, views = receiver.read_span() + for i in range(k): + torch.testing.assert_close(views[0][i], pushed[expected + i][0].cuda()) + receiver.commit_reads(k) + expected += k + assert expected == CAPACITY + + +@cuda_only +def test_ipc_and_host_ring_produce_identical_replay_buffers(): + """The learner's ingest path is transport-equivalent at the buffer level.""" + total = CAPACITY * 4 + host = SharedTransitionRing(total, N_ENV, OBS, CRI, ACT) + rb_host = SimpleReplayBuffer(N_ENV, 2 * total, OBS, ACT, CRI, device="cpu") + rb_ipc = SimpleReplayBuffer(N_ENV, 2 * total, OBS, ACT, CRI, device="cuda") + owner, receiver = _owner_receiver() + batches = [_batch(t, seed=11) for t in range(total)] + for fields in batches: + assert host.push(*fields) + while not owner.push(*fields): + while receiver.has_next(): + k, views = receiver.read_span() + rb_ipc.extend_batch(*views) + receiver.commit_reads(k) + torch.cuda.synchronize() # land the last push's event before the final drains + owner.size() # producer-side flush, as the collector's polling would do + while host.has_next(): + _k, views = host.read_span() + rb_host.extend_batch(*views) + host.commit_reads(_k) + while receiver.has_next(): + k, views = receiver.read_span() + rb_ipc.extend_batch(*views) + receiver.commit_reads(k) + assert rb_ipc.ptr == rb_host.ptr == total + for name in ("observations", "critic_observations", "actions", "rewards", "dones", "truncations"): + assert torch.equal(getattr(rb_ipc, name).cpu(), getattr(rb_host, name)), name + + +def test_use_ipc_transition_ring_gating(): + from motrix_rl.fastsac.async_impl.worker import use_ipc_transition_ring + + def opts(mode): + return SimpleNamespace(transition_ipc=mode) + + cpu, gpu, gpu0, gpu1 = ( + torch.device("cpu"), + torch.device("cuda"), + torch.device("cuda", 0), + torch.device("cuda", 1), + ) + assert use_ipc_transition_ring(opts("off"), gpu0, gpu0) is False + assert use_ipc_transition_ring(opts("auto"), gpu0, gpu0) is True + # unspecified index means the default current device (cuda:0), NOT a + # wildcard: it matches cuda:0 but never an explicit other GPU + assert use_ipc_transition_ring(opts("auto"), gpu, gpu0) is True + assert use_ipc_transition_ring(opts("auto"), gpu, gpu1) is False + assert use_ipc_transition_ring(opts("auto"), cpu, gpu0) is False + assert use_ipc_transition_ring(opts("auto"), gpu0, gpu1) is False + assert use_ipc_transition_ring(opts("on"), gpu0, gpu1) is False # warns + falls back + assert use_ipc_transition_ring(opts("on"), cpu, gpu0) is False + assert use_ipc_transition_ring(opts(True), gpu0, gpu0) is True # YAML boolean form + with pytest.raises(ValueError, match="transition_ipc"): + use_ipc_transition_ring(opts("nope"), gpu0, gpu0) + + +@cuda_only +def test_learner_drains_ipc_ring_end_to_end(): + """Learner.drain on the IPC ring fills the GPU replay buffer correctly. + + Covers the full consumer path — strided fused views, float->int64 done + conversion, rb wrap — interleaved with publish/commit laziness, mirroring + the earlier batched-drain smoke test for the host ring. + """ + from motrix_rl.fastsac.agent import FastSacAgent + from motrix_rl.fastsac.async_impl.learner import Learner + + agent = FastSacAgent( + obs_dim=OBS, + critic_obs_dim=CRI, + act_dim=ACT, + num_envs=N_ENV, + 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=True, + target_entropy_ratio=0.0, + buffer_size=32, + 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"), + ) + owner, receiver = _owner_receiver() + cfg = SimpleNamespace( + trainer=SimpleNamespace( + 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 + + total = 40 # > rb cap 33 -> exercises the rb wrap + batches = [_batch(t, seed=5) for t in range(total)] + for t, fields in enumerate(batches): + while not owner.push(*fields): + owner.size() + learner.drain() + owner.size() + torch.cuda.synchronize() # land the last publish event before the final drain + owner.size() + while receiver.has_next(): + learner.drain() + learner.wait_ingest() + torch.cuda.synchronize() + assert agent.rb.ptr == total + cap = 33 + for t in range(total - cap + 1, total): # surviving slots after the wrap + s = t % cap + torch.testing.assert_close(agent.rb.observations[:, s].cpu(), batches[t][0]) + torch.testing.assert_close(agent.rb.critic_observations[:, s].cpu(), batches[t][1]) + torch.testing.assert_close(agent.rb.actions[:, s].cpu(), batches[t][2]) + 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]) diff --git a/motrix_rl/tests/test_fastsac_learner.py b/motrix_rl/tests/test_fastsac_learner.py index 5155948c..9e292d8c 100644 --- a/motrix_rl/tests/test_fastsac_learner.py +++ b/motrix_rl/tests/test_fastsac_learner.py @@ -73,6 +73,98 @@ def test_own_leaves_non_tensors_alone() -> None: assert _own((1, "a", None)) == (1, "a", None) +def _tiny_agent_cfg(**overrides): + """Minimal valid FastSacAgentCfg namespace for a real (tiny) CUDA agent.""" + values = dict( + 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=True, + target_entropy_ratio=0.0, + buffer_size=64, + num_steps=1, + batch_size=8, + learning_starts=1, + policy_frequency=4, + num_updates=4, + obs_normalization=True, + compile=True, + amp=True, + amp_dtype="bf16", + device=None, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def test_update_metrics_survive_later_graph_generations() -> None: + """Outputs carried across `cudagraph_mark_step_begin` must stay readable. + + Regression for the `_own` placement in ``FastSacAgent.update``: the actor + pair is carried across policy-frequency gating (produced at a non-final + iteration), and the main outputs are owned only at the final iteration. + Both must remain readable after FURTHER update calls open new CUDA graph + generations — reading graph-pool memory invalidated by a later replay + raises instead of returning stale numbers. + + Needs a GPU: reduce-overhead compiles to CUDA graphs and is a no-op + without one, so there is nothing to invalidate on CPU. + """ + import pytest + import torch + + if not torch.cuda.is_available(): + pytest.skip("CUDA graphs require a GPU") + + from motrix_rl.fastsac.agent import FastSacAgent + + n_env, obs, cri, act = 8, 5, 7, 3 + agent = FastSacAgent( + obs_dim=obs, + critic_obs_dim=cri, + act_dim=act, + num_envs=n_env, + cfg=_tiny_agent_cfg(), + device=torch.device("cuda"), + ) + for _ in range(40): + agent.rb.extend( + torch.randn(n_env, obs), + torch.randn(n_env, cri), + torch.randn(n_env, act), + torch.randn(n_env), + torch.zeros(n_env, dtype=torch.long), + torch.zeros(n_env, dtype=torch.long), + ) + + # warmup/compile, then exercise both gating positions: + agent.update(4) # update_idx 0 -> 4: actor gated at i=0 (non-final, carried) + first = agent.update(4) # update_idx 4 -> 8: actor gated at i=0 (non-final) + agent.update_idx = 5 + second = agent.update(4) # update_idx 5 -> 9: actor gated at i=3 (final iteration) + # Further generations have invalidated earlier graph pools; the returned + # metrics must be owned copies and therefore still readable. + agent.update(4) + for metrics in (first, second): + values = {key: float(value) for key, value in metrics.items()} + assert all(value == value for value in values.values()) # no NaNs from torn reads + + 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 diff --git a/motrix_rl/tests/test_rl_sim_backend.py b/motrix_rl/tests/test_rl_sim_backend.py index c908f00c..353a11f3 100644 --- a/motrix_rl/tests/test_rl_sim_backend.py +++ b/motrix_rl/tests/test_rl_sim_backend.py @@ -19,8 +19,8 @@ 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.shm import Control, SharedTransitionRing -from motrix_rl.fastsac.async_impl.shm.weight_channel import HostWeightSender, WeightChannelShared +from motrix_rl.fastsac.async_impl.transport import Control, SharedTransitionRing +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 @@ -262,7 +262,9 @@ def _collect_in_spawn(sim_backend: str) -> tuple[torch.Tensor, ...]: 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)) - slot_queue.put(weight_tx.params) # ship before the collector process starts + # 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 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) @@ -289,9 +291,9 @@ def _collect_in_spawn(sim_backend: str) -> tuple[torch.Tensor, ...]: assert child_error is None assert process.exitcode == 0 assert control.collector_steps == 1 - slot = ring.read_slot() - assert slot is not None - return tuple(tensor.clone() for tensor in slot) + count, views = ring.read_span() + assert count >= 1 + return tuple(tensor.clone() for tensor in views) @pytest.mark.skipif( diff --git a/motrix_rl/tests/test_system_metrics.py b/motrix_rl/tests/test_system_metrics.py index 09676563..3e808741 100644 --- a/motrix_rl/tests/test_system_metrics.py +++ b/motrix_rl/tests/test_system_metrics.py @@ -8,10 +8,12 @@ import motrix_rl.system_metrics as system_metrics from motrix_rl.system_metrics import ( CpuLoadSampler, + GpuDeviceUsage, GpuMemoryUsageSampler, GpuUtilizationSampler, MemoryUsage, MemoryUsageSampler, + sample_gpu_devices, ) @@ -20,7 +22,10 @@ def test_cpu_load_sampler_uses_counter_deltas_for_available_cpus(tmp_path) -> No stat_path.write_text( "cpu 200 0 0 1800 0 0 0 0\ncpu0 100 0 0 900 0 0 0 0\ncpu1 100 0 0 900 0 0 0 0\ncpu2 100 0 0 900 0 0 0 0\n" ) - sampler = CpuLoadSampler(stat_path=stat_path, topology_root=tmp_path, cpu_ids={0, 1}) + (tmp_path / "cpuinfo").write_text("processor\t: 0\nmodel name\t: AMD EPYC 9654 96-Core Processor\nflags\t: fpu\n") + sampler = CpuLoadSampler( + stat_path=stat_path, topology_root=tmp_path, cpuinfo_path=tmp_path / "cpuinfo", cpu_ids={0, 1} + ) stat_path.write_text( "cpu 330 0 0 1850 20 0 0 0\ncpu0 150 0 0 950 0 0 0 0\ncpu1 180 0 0 900 20 0 0 0\ncpu2 300 0 0 900 0 0 0 0\n" @@ -35,6 +40,10 @@ def test_cpu_load_sampler_uses_counter_deltas_for_available_cpus(tmp_path) -> No assert load.physical_core_count is None assert load.iowait_percent == 10.0 assert load.steal_percent == 0.0 + # per-core breakdown comes from the same counter deltas, sorted by cpu id + assert load.per_core_percent == (50.0, 80.0) + # the static model name read once from cpuinfo rides along on every sample + assert load.model_name == "AMD EPYC 9654 96-Core Processor" def test_cpu_load_sampler_returns_none_without_elapsed_cpu_time(tmp_path) -> None: @@ -84,6 +93,55 @@ def test_gpu_samplers_aggregate_utilization_mean_and_memory_sum(monkeypatch) -> assert GpuMemoryUsageSampler().sample() == MemoryUsage(used_bytes=400 * 1024**2, total_bytes=600 * 1024**2) +def test_gpu_samplers_report_per_device_utilization_and_memory(monkeypatch) -> None: + handles = ["gpu0", "gpu1"] + _fake_nvml( + monkeypatch, + handles, + utilization={"gpu0": 10, "gpu1": 30}, + memory={"gpu0": (100 * 1024**2, 200 * 1024**2), "gpu1": (300 * 1024**2, 400 * 1024**2)}, + ) + + assert GpuUtilizationSampler().sample_per_device() == [10.0, 30.0] + assert GpuMemoryUsageSampler().sample_per_device() == [ + MemoryUsage(used_bytes=100 * 1024**2, total_bytes=200 * 1024**2), + MemoryUsage(used_bytes=300 * 1024**2, total_bytes=400 * 1024**2), + ] + assert sample_gpu_devices(GpuUtilizationSampler(), GpuMemoryUsageSampler()) == [ + GpuDeviceUsage( + index=0, + utilization_percent=10.0, + memory=MemoryUsage(used_bytes=100 * 1024**2, total_bytes=200 * 1024**2), + ), + GpuDeviceUsage( + index=1, + utilization_percent=30.0, + memory=MemoryUsage(used_bytes=300 * 1024**2, total_bytes=400 * 1024**2), + ), + ] + + +def test_sample_gpu_devices_degrades_to_utilization_only(monkeypatch) -> None: + handles = ["gpu0"] + _fake_nvml(monkeypatch, handles, utilization={"gpu0": 55}, memory={}) + + class _BrokenMemorySampler(GpuMemoryUsageSampler): + def sample_per_device(self): + return None + + assert sample_gpu_devices(GpuUtilizationSampler(), _BrokenMemorySampler()) == [ + GpuDeviceUsage(index=0, utilization_percent=55.0, memory=None) + ] + + +def test_sample_gpu_devices_is_none_without_a_per_device_source(monkeypatch) -> None: + # no backend session at all + monkeypatch.setattr(system_metrics, "_amd_smi_state", ()) + monkeypatch.setattr(system_metrics, "_nvml_state", ()) + + assert sample_gpu_devices(GpuUtilizationSampler(), GpuMemoryUsageSampler()) is None + + def test_gpu_samplers_return_none_on_nvml_error(monkeypatch) -> None: handles = ["gpu0"] _fake_nvml( diff --git a/scripts/bench_fastsac_collector_inference.py b/scripts/bench_fastsac_collector_inference.py deleted file mode 100644 index 3872ec44..00000000 --- a/scripts/bench_fastsac_collector_inference.py +++ /dev/null @@ -1,221 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Benchmark the FastSAC collector actor-only inference boundary. - -The measured interval is exactly ``Collector._infer``: CPU policy observation -staging, optional H2D, read-only normalization, stochastic actor inference, -and optional D2H plus synchronization. Environment stepping, transition-ring -push, weight synchronization and learner work are intentionally excluded. -""" - -from __future__ import annotations - -import argparse -import json -import statistics -import time -from pathlib import Path -from types import SimpleNamespace - -import torch - -from motrix_rl.fastsac.async_impl.collector import Collector -from motrix_rl.fastsac.async_impl.shm import Control, SharedTransitionRing, WeightSnapshot -from motrix_rl.fastsac.buffer import EmpiricalNormalization -from motrix_rl.fastsac.networks import Actor - - -class _BenchmarkEnv: - def __init__(self, num_envs: int, obs_dim: int, critic_obs_dim: int): - self.num_envs = num_envs - self._obs = torch.randn(num_envs, obs_dim) - self._critic_obs = torch.zeros(num_envs, critic_obs_dim) - self.last_info = {} - - def reset(self): - return self._obs, self._critic_obs - - -def _percentile(values: list[float], percentile: float) -> float: - ordered = sorted(values) - index = max(0, min(len(ordered) - 1, int((len(ordered) - 1) * percentile + 0.5))) - return ordered[index] - - -def _source_policy(args): - action_scale = torch.ones(args.act_dim) - action_bias = torch.zeros(args.act_dim) - actor = Actor( - n_obs=args.obs_dim, - n_act=args.act_dim, - hidden_dim=args.hidden_dim, - log_std_max=0.0, - log_std_min=-5.0, - use_tanh=True, - use_layer_norm=True, - action_scale=action_scale, - action_bias=action_bias, - device="cpu", - ) - with torch.no_grad(): - actor.fc_mu.weight.normal_(0.0, 0.02) - actor.fc_mu.bias.normal_(0.0, 0.02) - actor.fc_logstd.weight.normal_(0.0, 0.02) - actor.fc_logstd.bias.normal_(0.0, 0.02) - normalizer = EmpiricalNormalization(args.obs_dim, device="cpu") - normalizer._mean.normal_(0.0, 0.1) - normalizer._std.uniform_(0.8, 1.2) - normalizer._var.copy_(normalizer._std.square()) - normalizer.count.fill_(1_000_000) - return actor, normalizer, action_scale, action_bias - - -def _collector(args, source_actor, source_normalizer, action_scale, action_bias): - cfg = SimpleNamespace( - agent=SimpleNamespace( - actor_hidden_dim=args.hidden_dim, - log_std_max=0.0, - log_std_min=-5.0, - use_tanh=True, - use_layer_norm=True, - obs_normalization=True, - learning_starts=0, - ), - trainer=SimpleNamespace( - async_options=SimpleNamespace( - collector_inference_device=args.device, - collector_compile=args.compile, - collector_amp=args.amp, - collector_amp_dtype=args.amp_dtype, - weight_poll_interval=1, - ) - ), - ) - env = _BenchmarkEnv(args.num_envs, args.obs_dim, args.critic_obs_dim) - # Stub ring (never pushed: the bench only exercises inference/sync); capacity - # 2 satisfies the ring's successor-slot contract. - ring = SharedTransitionRing(2, args.num_envs, args.obs_dim, args.critic_obs_dim, args.act_dim) - weights = WeightSnapshot(sum(p.numel() for p in source_actor.parameters()), args.obs_dim) - weights.publish(source_actor, source_normalizer) - collector = Collector( - env, - cfg, - args.obs_dim, - args.critic_obs_dim, - args.act_dim, - action_scale, - action_bias, - ring, - weights, - Control(), - ) - collector.reset() - collector.sync_weights() - collector.warmup_inference() - return collector, weights - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cpu") - parser.add_argument("--compile", action="store_true") - parser.add_argument("--amp", action="store_true") - parser.add_argument("--amp-dtype", choices=("fp16", "bf16"), default="fp16") - parser.add_argument("--num-envs", type=int, default=4096) - parser.add_argument("--obs-dim", type=int, default=154) - parser.add_argument("--critic-obs-dim", type=int, default=154) - parser.add_argument("--act-dim", type=int, default=29) - parser.add_argument("--hidden-dim", type=int, default=512) - parser.add_argument("--warmup", type=int, default=20) - parser.add_argument("--samples", type=int, default=200) - parser.add_argument("--sync-samples", type=int, default=50) - parser.add_argument("--cpu-threads", type=int, default=12) - parser.add_argument("--output", type=Path) - args = parser.parse_args() - - torch.manual_seed(7) - torch.set_num_threads(args.cpu_threads) - source_actor, source_normalizer, action_scale, action_bias = _source_policy(args) - collector, weights = _collector(args, source_actor, source_normalizer, action_scale, action_bias) - obs = torch.randn(args.num_envs, args.obs_dim) - - for _ in range(args.warmup): - collector._infer(obs) - - samples_ms = [] - first = None - last = None - for index in range(args.samples): - start = time.perf_counter() - actions = collector._infer(obs) - samples_ms.append((time.perf_counter() - start) * 1000.0) - if index == 0: - first = actions.clone() - if index == args.samples - 1: - last = actions.clone() - - with torch.no_grad(): - normalized = source_normalizer(obs, update=False) - cpu_deterministic = source_actor.explore(normalized, deterministic=True) - with collector._autocast(): - device_deterministic = collector._policy.deterministic(obs.to(collector.device)).cpu() - abs_error = (device_deterministic - cpu_deterministic).abs() - sync_samples_ms = [] - for _ in range(args.sync_samples): - weights.publish(source_actor, source_normalizer) - start = time.perf_counter() - collector.sync_weights() - sync_samples_ms.append((time.perf_counter() - start) * 1000.0) - staging_ptrs = [ - buffer.data_ptr() - for buffer in (collector._obs_host, collector._obs_device, collector._actions_host) - if buffer is not None - ] - result = { - "torch_version": torch.__version__, - "device": str(collector.device), - "device_name": torch.cuda.get_device_name(collector.device) if collector.device.type == "cuda" else None, - "compile": args.compile, - "amp": args.amp, - "amp_dtype": args.amp_dtype if args.amp else None, - "num_envs": args.num_envs, - "obs_dim": args.obs_dim, - "critic_obs_dim": args.critic_obs_dim, - "act_dim": args.act_dim, - "hidden_dim": args.hidden_dim, - "cpu_threads": torch.get_num_threads(), - "warmup": args.warmup, - "samples": args.samples, - "latency_ms": { - "median": statistics.median(samples_ms), - "p90": _percentile(samples_ms, 0.90), - "mean": statistics.fmean(samples_ms), - "min": min(samples_ms), - "max": max(samples_ms), - }, - "weight_sync_ms": { - "median": statistics.median(sync_samples_ms), - "p90": _percentile(sync_samples_ms, 0.90), - "mean": statistics.fmean(sync_samples_ms), - "min": min(sync_samples_ms), - "max": max(sync_samples_ms), - }, - "deterministic_vs_cpu": { - "max_abs_error": float(abs_error.max()), - "mean_abs_error": float(abs_error.mean()), - }, - "stochastic_samples_differ": not torch.equal(first, last), - "actions_finite": bool(torch.isfinite(last).all()), - "actions_in_bounds": bool(torch.all(last >= -1.0) and torch.all(last <= 1.0)), - "staging_ptrs": staging_ptrs, - } - rendered = json.dumps(result, indent=2) - print(rendered) - if args.output: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(rendered + "\n") - - -if __name__ == "__main__": - main() diff --git a/wiki/design/fastsac-async-heterogeneous-trainer.md b/wiki/design/fastsac-async-heterogeneous-trainer.md index 0ed753eb..45f64d94 100644 --- a/wiki/design/fastsac-async-heterogeneous-trainer.md +++ b/wiki/design/fastsac-async-heterogeneous-trainer.md @@ -49,15 +49,19 @@ - **Collector（进程 A）**：拥有唯一的 CPU `DirectEnv` 与一个**推理专用**的 `Actor` + `obs_normalizer` 副本（CUDA 默认，也可显式选择 CPU；两者均 `eval`）。只做前向、`env.step`、把 transition 批推入共享环、读权重快照、维护 episode 记账。不持有 optimizer、qnet、replay buffer。 - **Learner（进程 B，GPU）**：就是现有的 `FastSacAgent`，但**不驱动 env**。它从共享环把 transition 灌进自己的 GPU replay buffer，照常调用 `agent.update(n)`；周期性把 actor 权重 + normalizer 快照发布到共享内存；并独占日志与 checkpoint。 -**为什么 transition 走「CPU 共享环 + learner 端 ingest」，而不是两进程直接共享 GPU replay buffer？** +**为什么 transition 走「共享环 + learner 端 ingest」，而不是两进程直接共享 GPU replay buffer？** - `SimpleReplayBuffer.sample()` 含大量 gather / n-step 计算，是 learner 独占的读路径；让 collector 也触碰会引入跨进程锁与 GPU 上下文共享。 -- CUDA IPC 共享 GPU tensor 复杂且脆弱（context、stream、生命周期）；而 CPU 共享内存 tensor（`Tensor.share_memory_()` / `torch.multiprocessing`）成熟稳定。 -- collector 本就是 CPU 负载，transition 先落 CPU 共享内存零成本；learner ingest 时一次性 `.to(device)` 批量上 GPU，比同步版「每 env-step 一次小拷贝」更 GPU 友好。 +- 环与 buffer 解耦：环只做无锁搬运，buffer 的环形索引、n-step 语义完全归属 learner。 -**collector 的环境固定在 CPU，Actor 的设备独立配置**：基于 2048/4096 环境的服务器端到端吞吐结果，默认把 actor 与只读 observation normalizer 放到 CUDA；CPU 保留为显式兼容配置。CUDA 路径不改变 env wrapper 的 device，也不移动 critic observation、reward/done、bookkeeping 或共享 transition ring。learner 与 collector 同卡时仍需结合具体任务确认资源竞争边界。 +环本身有两种物理传输，由 `transition_ipc` 选择（见 §4.1）： -CUDA 推理的数据边界为：CPU policy observation 先复制到预分配 pinned host buffer，再异步 H2D 到固定 shape device buffer；actor 输出立即异步 D2H 到预分配 pinned action buffer，并在返回 CPU env 前同步。若 `torch.compile(mode="reduce-overhead")` 复用 CUDA Graph output storage，D2H 已在下一次 replay 前完成，因此环境不会持有随后被覆盖的 device output 引用。权重仍由 learner 发布到 CPU `WeightSnapshot`；collector actor 参数绑定到一个 contiguous device flat buffer，每个新版本只做一次 pinned H2D，而不是逐参数传输。 +- **host 共享内存环**（默认回退路径）：CPU `share_memory_()` tensor，任何设备组合都可用。 +- **CUDA-IPC 设备环**：collector 与 learner 推理/训练在同一 GPU 时，slot 直接放在显存里。env 的 CPU 输出在 collector 侧一次融合 H2D 直写 slot，learner 的 ingest 变为纯 D2D——env 产生的字节只过一次 PCIe，且被 collector 的 env step 时间掩盖；learner 关键路径上不再有任何 H2D/staging memcpy。 + +**collector 的环境固定在 CPU，Actor 的设备独立配置**：基于 2048/4096 环境的服务器端到端吞吐结果，默认把 actor 与只读 observation normalizer 放到 CUDA；CPU 保留为显式兼容配置。CUDA 路径不改变 env wrapper 的 device，也不移动 critic observation、reward/done、bookkeeping；transition 环的物理位置由 `transition_ipc` 决定，与推理设备解耦。learner 与 collector 同卡时仍需结合具体任务确认资源竞争边界。 + +CUDA 推理的数据边界为：CPU policy observation 先复制到预分配 pinned host buffer，再异步 H2D 到固定 shape device buffer；actor 输出立即异步 D2H 到预分配 pinned action buffer，并在返回 CPU env 前同步。若 `torch.compile(mode="reduce-overhead")` 复用 CUDA Graph output storage，D2H 已在下一次 replay 前完成，因此环境不会持有随后被覆盖的 device output 引用。权重由 learner 发布到权重通道（seqlock 双 buffer），物理位置由 `weight_ipc` 选择：`auto`（默认）在 learner 与 collector 推理同卡且 actor 参数达到 `weight_ipc_min_bytes` 时使用 CUDA-IPC 设备槽（publish 是一次 D2D 拷贝），否则使用 host 共享内存槽（一次融合 pinned D2H + shm 写入）。collector actor 参数绑定到一个 contiguous device flat buffer：host 路径每个新版本只做一次 pinned H2D，而不是逐参数传输；device 路径则完全不经 host。 --- @@ -115,10 +119,24 @@ uv run scripts/train.py task=g1-walk-flat/motrix.fastsac algo.asynchronous=false 两个共享游标 `_write` / `_read` 实现无锁环： - **生产**：`push()` 满环（`write - read >= C`）时返回 `False`，**不 step env、不丢数据**；写入 slot 各字段后再 `_write += 1`（x86/TSO 下字段写保证先于游标 bump 可见）。 -- **消费**：`read_slot()` 返回最旧未读 slot 的零拷贝 CPU 视图但**不推进** `_read`；learner 把数据 `.to(device)` 后再 `commit_read()`（`_read += 1`）。读游标只在拷贝完成后前进，故生产者永不覆盖仍在 ingest 的 slot。 +- **消费**：`read_span()` 返回最长连续未读 run 的长度与六个字段的 `(count, num_envs, dim)` 视图但**不推进** `_read`；learner 拷入 replay buffer 后再 `commit_reads(count)`（`_read += count`，同様由 event 定序，见 IPC 环）。读游标只在拷贝完成后前进，故生产者永不覆盖仍在 ingest 的 slot。 +- 游标（`RingCursors`）是独立于字段存储的 host 共享 tensor，由父进程创建——host 环与 IPC 环复用同一游标协议。 **背压方向是核心旋钮**：环满（collector 快）→ collector 阻塞采样，天然把采样速率压到 learner 消费速率，防止无界内存增长、防止 replay buffer 被过新数据刷爆而 off-policy 失真；环空（learner 快）→ learner 无新数据可 ingest，由 §5 的 UTD 治理决定「等数据」还是「在已有 buffer 上继续更新」。`C` 需足够吸收两进程抖动（一次 GC、一次 CUDA sync），但不宜过大以免抬高在途 staleness。 +#### CUDA-IPC 设备环（`transition_ipc`） + +host 环下 learner 的 drain 承担全部搬运：ring→pinned 的 CPU memcpy 与 H2D 都在 learner 关键路径上。设备环把这段搬运整体移到 collector 侧并隐藏： + +- **融合 slot 布局**：全部字段放进一块 `(C, num_envs, obs_dim + critic_obs_dim + act_dim + 3)` 的 f32 设备 tensor（reward/done/truncation 以 0/1 float 存储，learner 拷入 i64 buffer 时由 `copy_` 顺手完成数值恒等的类型转换）。collector 每步把六个 CPU 字段拼进一块 pinned staging，**一次 H2D** 写入整个 slot；learner 侧按最后一维 offset 切出六个 strided 视图直接喂 `extend_batch`（纯 D2D）。 +- **归属与交接**：设备字段由 learner 进程分配（IPC handle 导出方需要 CUDA context 且必须保活），经既有的 `slot_queue` 握手随权重 slot 一起 ship 给 collector（`torch.multiprocessing` 已注册 CUDA tensor 的跨进程 reducer）；collector 侧从到达的 tensor 建立接收端。游标沿用父进程创建的 host `RingCursors`，两进程可见。 +- **发布定序（正确性核心）**：GPU 写入不受 x86/TSO 保障，游标 bump 前必须让设备写入落定——生产者在 H2D 入队后对该 slot record 一个 CUDA event，`_write` 只推进到「event 已完成的最旧 slot + 1」（惰性 flush：`is_full`/`push` 前查询队首 event，完成即推进）。消费侧对称：`commit_reads` 在 D2D 读取入队后 record event，`_read` 只推进到「event 已完成」的边界。staging 复用同样由 event 守护：覆写前确认上一次 H2D 已完成（正常节奏下 env.step 的毫秒级间隔使其成为 no-op）。 +- **采样无竞争**：learner 的 D2D ingest 与 `sample()` 在同一条默认流上，流内天然有序，不需要额外 event；event 只服务于跨进程游标推进。 +- **传输选择**：`transition_ipc: auto/on/off`。`auto` 要求 learner 设备与 `collector_inference_device` 解析后同为 CUDA（未显式给 index 时视为同卡；显式 index 不同则回退）；`on` 在不满足时告警回退 host 环。host 环是永久保留的回退路径，覆盖 CPU collector、异卡与 IPC 建立失败的边界。 +- **代价**：显存增加约 `C × num_envs × feat × 4` 字节（microduck-walk-flat 量级约 130MB），高维 critic_obs 任务需把 `ring_capacity` 纳入显存预算。 + + + **语义一致性**：collector 只是把同步版 collect 相位原样搬到另一进程，transition 的构造代码相同——调用顺序、dtype（dones/truncations 用 long）、auto-reset 后的 next_obs 语义与同步版逐字节一致。这是「算法未变、只变执行拓扑」的基础。 ### 4.2 WeightSnapshot — seqlock 双缓冲（learner → collector） @@ -231,6 +249,8 @@ class FastSacAsyncOptionsCfg: collector_compile: bool = True # CUDA 固定 batch 推理使用 reduce-overhead collector_amp: bool = True # 默认使用实测吞吐最优的 FP16 collector autocast 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 + 大小门控 @dataclass