diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md index f1e19efcf..cdd80dd00 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md @@ -7,7 +7,7 @@ current log name is `fast_sac`. ## Runtime Model -The off-policy runner decouples CPU simulation from accelerator learning through +The off-policy runner decouples simulation collection from accelerator learning through bounded shared memory. A collector subprocess publishes packed transitions through two ingress slots, while the complete replay ring is authoritative on one CUDA or Apple MPS learner device. Host replay allocation therefore does not @@ -44,3 +44,17 @@ uv run train --algo sac --task g1_walk_flat --sim mujoco \ algo.max_iterations=1000 \ training.no_play=true ``` + +## Single-node multi-GPU device placement + +`training.devices` assigns rank i's learner to `cuda:devices[i]`; each rank owns one +collector. For mjwarp, the rank process and its collector process explicitly bind Warp's +default/current device to that same learner device before probe or production environment +materialization. The collector therefore does not fall back to Warp's fresh-process default +of `cuda:0`. The local binding is recorded as `collector_backend_device` in the runtime +manifest. + +MuJoCo has a committed multi-GPU scaling benchmark. The mjwarp per-rank placement contract is +covered by `tests/base/backend/test_process_device.py` and the off-policy runner/worker unit +tests; the repository does not currently contain an mjwarp multi-GPU throughput or convergence +benchmark. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md index ab75031a7..bebe85295 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md @@ -6,7 +6,7 @@ SAC 通过 `scripts/train_sac.py` 运行；TD3 与 FlashSAC 各有独立的入 ## 运行模型 -off-policy runner 通过有界 shared memory 把 CPU 仿真与 accelerator 学习解耦。 +off-policy runner 通过有界 shared memory 把仿真采集与 accelerator 学习解耦。 collector 子进程通过两个 packed ingress slot 发布 transition，完整 replay ring 只由 一个 CUDA 或 Apple MPS learner device 持有。因此 host replay 分配不随 replay capacity 增长；slot 的 device copy 完成后才推进 `ptr` 与 `size`。CUDA 通过 side @@ -51,6 +51,12 @@ rank 子进程。启动时 rank 0 一次性广播 actor、critic、target critic 执行阻塞式 flat-gradient `all_reduce(SUM) / world_size`。各 rank 不交换 replay 数据， 在相同初值、平均梯度和更新顺序下各自维护一致的 optimizer 状态。 +每个 rank 当前只创建一个 collector。使用 mjwarp 时，rank i 会在 probe env 和 collector +正式 env materialization 之前，把 Warp 的进程默认/当前 device 显式绑定到该 rank 的 +learner device `cuda:devices[i]`；因此 collector 不依赖 Warp 新进程默认的 `cuda:0`，也不 +会跨 rank 集中到同一张卡。runtime manifest 的 `collector_backend_device` 记录本 rank 的 +实际绑定。 + off-policy 只公开 `training.devices` 这一个设备字段：`null` 或 `[]` 自动选择单个 learner device，`[0]` 显式选择 `cuda:0`，两个以上索引才启动多卡拓扑。 @@ -86,7 +92,9 @@ collector 的 CPU 亲和按 rank 自动均分（`cpu_count // world_size` 一段 默认 stream 与 side stream 均通过；跳过 warmup 时首个 all-reduce 会在 capture 中报 `operation not permitted when stream is capturing`。有限超时的最小复现见 `scripts/benchmark/rl/reproduce_nccl_cuda_graph_capture.py`。 -- 仅验证过 `mujoco` backend。 +- MuJoCo 有已提交的多卡 scaling benchmark；mjwarp 的 per-rank device placement 有 + `tests/base/backend/test_process_device.py` 与 off-policy runner/worker 单测覆盖，但仓库中 + 尚无 mjwarp 多卡吞吐或收敛 benchmark。 - 仅单节点：rank 之间通过 run 目录里的 FileStore rendezvous，NCCL 走 TCP loopback（默认 `NCCL_P2P_DISABLE=1` / `NCCL_SHM_DISABLE=1`，环境变量显式设置 时优先）——部分机型（如 RTX 6000D）的 NCCL P2P/SHM peer transport 不可靠， diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index a4774a03a..46a732b01 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -20,6 +20,7 @@ sys.path.append(str(ROOT_DIR)) from unilab.base.backend.base import log_playback_plan +from unilab.base.backend.process_device import configure_backend_process_device from unilab.base.config_adapter import create_env from unilab.ipc.dp_launcher import ( UNILAB_DP_LOG_DIR, @@ -106,6 +107,10 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): from unilab.utils.device import get_default_device rank_device = resolve_dp_rank_device(dp_devices, dp_rank) or get_default_device() + # Bind backend-global device state before algorithm builders materialize + # their probe envs. The spawned collector repeats this binding in its own + # process using the same rank-local device. + configure_backend_process_device(str(cfg.training.sim_backend), rank_device) host_cpu_count = os.cpu_count() or 1 explicit_cpu_ids = getattr(cfg.training, "dp_collector_cpu_ids", None) if explicit_cpu_ids is not None: diff --git a/src/unilab/algos/offpolicy/double_buffer_runner.py b/src/unilab/algos/offpolicy/double_buffer_runner.py index 5d58e5c65..f689811ba 100644 --- a/src/unilab/algos/offpolicy/double_buffer_runner.py +++ b/src/unilab/algos/offpolicy/double_buffer_runner.py @@ -27,6 +27,7 @@ torch_thread_env, ) from unilab.algos.offpolicy.worker import off_policy_collector_fn, sample_offpolicy_actions +from unilab.base.backend.process_device import resolve_backend_process_device from unilab.ipc.async_runner import _SPAWN_CTX from unilab.ipc.inference_slot import SharedInferenceSlot from unilab.ipc.replay_buffer import DEFAULT_REPLAY_INGRESS_DEPTH, ReplayBuffer @@ -159,6 +160,10 @@ def __init__( **kwargs, ): kwargs["device"] = require_offpolicy_replay_device(kwargs.get("device")) + collector_backend_device = resolve_backend_process_device( + str(kwargs.get("sim_backend", "mujoco")), + kwargs["device"], + ) super().__init__(**kwargs) if replay_prefetch_mode != "one_tick": raise ValueError( @@ -168,6 +173,7 @@ def __init__( # Per-rank CPU block owned by this rank's collector (multi-GPU DP); # merged into the collector-only env override at collector startup. self.collector_cpu_ids = list(collector_cpu_ids) if collector_cpu_ids is not None else None + self.collector_backend_device = collector_backend_device # Multi-GPU synchronous data parallelism (None = the bit-identical # single-rank path): startup model broadcast, then gradient averaging # before every actor/critic/temperature optimizer step. @@ -181,7 +187,8 @@ def __init__( self.runtime_manifest = { "inference_owner": "learner", "collector_actor": False, - "collector_accelerator_context": False, + "collector_accelerator_context": self.collector_backend_device is not None, + "collector_backend_device": self.collector_backend_device, "collector_torch_inference": False, "learner_actor_reused": True, "logger_owner_rank": 0, @@ -918,7 +925,9 @@ def learn( f"({self.replay_transfer_backend.get('device_family')})" ) logger.log_status(f"Inference owner: learner.actor ({self.device})") - logger.log_status("Collector model/accelerator ownership: none") + if self.collector_backend_device is not None: + logger.log_status(f"Collector backend device: {self.collector_backend_device}") + logger.log_status("Collector actor/inference ownership: none") logger.log_status("Replay learner lightweight: fixed (log_interval=1)") self._active_logger = logger logger.start() @@ -943,6 +952,7 @@ def learn( "inference_request_queue": inference_request_queue, "inference_response_queue": inference_response_queue, "sim_backend": self.sim_backend, + "backend_device": self.collector_backend_device, "env_cfg_override": self._collector_env_cfg_override(), "inference_slot": inference_slot, "seed": derive_worker_seed(self.seed, worker_index=0), diff --git a/src/unilab/algos/offpolicy/worker.py b/src/unilab/algos/offpolicy/worker.py index 9c5312b6b..7bb8e748b 100644 --- a/src/unilab/algos/offpolicy/worker.py +++ b/src/unilab/algos/offpolicy/worker.py @@ -10,6 +10,7 @@ from unilab.algos.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.algos.offpolicy.thread_budget import apply_torch_thread_runtime +from unilab.base.backend.process_device import configure_backend_process_device from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.observations import split_obs_dict from unilab.base.registry import ensure_registries @@ -156,6 +157,7 @@ def off_policy_collector_fn( algo_type: str = "sac", metrics_queue=None, sim_backend: str = "mujoco", + backend_device: str | None = None, env_cfg_override: dict | None = None, seed: int | None = None, trace_enabled: bool = False, @@ -179,6 +181,7 @@ def off_policy_collector_fn( algo_type=algo_type, metrics_queue=metrics_queue, sim_backend=sim_backend, + backend_device=backend_device, env_cfg_override=env_cfg_override, seed=seed, trace_enabled=trace_enabled, @@ -199,6 +202,7 @@ def _run_collector( algo_type, metrics_queue, sim_backend, + backend_device, env_cfg_override, seed, trace_enabled, @@ -209,6 +213,7 @@ def _run_collector( from unilab.base import registry apply_torch_thread_runtime(torch_thread_runtime, role="collector", torch_module=torch) + configured_backend_device = configure_backend_process_device(sim_backend, backend_device) ensure_registries() apply_training_seed(seed, torch_runtime=False, cuda=False) @@ -265,6 +270,8 @@ def _run_collector( "actor_owned": False, "weight_sync_attached": False, "torch_inference": False, + "collector_accelerator_context": configured_backend_device is not None, + "collector_backend_device": configured_backend_device, "cuda_context_initialized": bool(torch.cuda.is_initialized()), } if trace_recorder: diff --git a/src/unilab/base/backend/mjwarp/runtime.py b/src/unilab/base/backend/mjwarp/runtime.py new file mode 100644 index 000000000..ea9685cad --- /dev/null +++ b/src/unilab/base/backend/mjwarp/runtime.py @@ -0,0 +1,18 @@ +"""Process-global runtime setup owned by the ``mjwarp`` backend.""" + +from __future__ import annotations + +from .dependencies import load_mjwarp_dependencies + + +def bind_mjwarp_process_device(device: str) -> str: + """Make one CUDA device Warp's default/current device for this process.""" + dependencies = load_mjwarp_dependencies() + dependencies.warp.set_device(device) + selected = dependencies.warp.get_device() + if not bool(selected.is_cuda): + raise RuntimeError( + "mjwarp backend requires an active CUDA Warp device; " + f"resolved {selected!s} from {device!r}" + ) + return str(selected) diff --git a/src/unilab/base/backend/process_device.py b/src/unilab/base/backend/process_device.py new file mode 100644 index 000000000..6729e8c18 --- /dev/null +++ b/src/unilab/base/backend/process_device.py @@ -0,0 +1,46 @@ +"""Cold-path process-device routing for simulation backends. + +The collector process must select any process-global accelerator device before +materializing its environment. Backend-specific selection stays behind this +adapter so algorithm workers do not import or probe optional backend runtimes. +""" + +from __future__ import annotations + + +def resolve_backend_process_device( + backend_type: str, + learner_device: str | None, +) -> str | None: + """Return the backend device that must follow one rank's learner device. + + Host MuJoCo and the other adapters do not need a runner-owned process + binding. ``mjwarp`` allocates through Warp's process-global default + device, so its collector must receive the exact CUDA device already + assigned to the rank's learner. + """ + if backend_type != "mjwarp": + return None + if learner_device is None: + raise ValueError("mjwarp requires an explicit CUDA process device") + + resolved = str(learner_device).strip() + if resolved.split(":", 1)[0].lower() != "cuda": + raise ValueError( + f"mjwarp requires a CUDA process device shared with its learner; got {resolved!r}" + ) + return resolved + + +def configure_backend_process_device( + backend_type: str, + learner_device: str | None, +) -> str | None: + """Bind a backend runtime to its resolved device before env materialization.""" + resolved = resolve_backend_process_device(backend_type, learner_device) + if resolved is None: + return None + + from .mjwarp.runtime import bind_mjwarp_process_device + + return bind_mjwarp_process_device(resolved) diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index e6c0ca4a2..595435e9a 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -365,6 +365,7 @@ def _build_sac_runner_with_fakes( overrides: list[str], *, cpu_count: int = 128, + backend_binding_calls: list[tuple[str, str]] | None = None, ): """build_runner("sac", ...) with learner/env/runner fakes; returns captured state.""" module = _offpolicy() @@ -377,6 +378,12 @@ def fake_create_env(*args, **kwargs): return _FakeEnv() monkeypatch.setattr(module.os, "cpu_count", lambda: cpu_count) + if backend_binding_calls is not None: + monkeypatch.setattr( + module, + "configure_backend_process_device", + lambda backend, device: backend_binding_calls.append((str(backend), str(device))), + ) import unilab.algos.fast_sac.double_buffer as owner_module @@ -389,6 +396,23 @@ def fake_create_env(*args, **kwargs): return runner, probe_env_calls +def test_build_runner_binds_mjwarp_rank_process_to_learner_device( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(UNILAB_DP_RANK, "1") + monkeypatch.setenv(UNILAB_DP_LOG_DIR, "/tmp/offpolicy_test_run") + bindings: list[tuple[str, str]] = [] + + runner, _ = _build_sac_runner_with_fakes( + monkeypatch, + ["task=g1_walk_flat/mjwarp", "training.devices=[0,1]"], + backend_binding_calls=bindings, + ) + + assert bindings == [("mjwarp", "cuda:1")] + assert runner.kwargs["device"] == "cuda:1" + + def test_build_runner_partitions_collector_cpus_per_rank(monkeypatch: pytest.MonkeyPatch): # Spawned rank: rank comes from the env, world_size from training.devices. monkeypatch.setenv(UNILAB_DP_RANK, "1") @@ -399,6 +423,7 @@ def test_build_runner_partitions_collector_cpus_per_rank(monkeypatch: pytest.Mon cpu_count=128, ) assert runner.kwargs["collector_cpu_ids"] == list(range(64, 128)) + assert runner.kwargs["device"] == "cuda:1" # The thread budget is resolved against the rank's CPU share, not the host. assert runner.kwargs["torch_thread_runtime"]["cpu_count"] == 64 # The num_envs=1 probe env must never see cpu_ids (it would size its diff --git a/tests/algos/test_offpolicy_runner_unit.py b/tests/algos/test_offpolicy_runner_unit.py index df9e65e09..d5b70e641 100644 --- a/tests/algos/test_offpolicy_runner_unit.py +++ b/tests/algos/test_offpolicy_runner_unit.py @@ -305,7 +305,13 @@ def _get_iter_wall_time(self): return 0.0 -def _make_device_runner(monkeypatch: pytest.MonkeyPatch, learner=None): +def _make_device_runner( + monkeypatch: pytest.MonkeyPatch, + learner=None, + *, + device: str = "cuda", + sim_backend: str = "mujoco", +): monkeypatch.setattr( device_runner_module, "require_offpolicy_replay_device", lambda value: value ) @@ -321,8 +327,61 @@ def _make_device_runner(monkeypatch: pytest.MonkeyPatch, learner=None): updates_per_step=2, policy_frequency=1, env_steps_per_sync=1, - device="cuda", + device=device, + sim_backend=sim_backend, + ) + + +def test_mjwarp_collector_backend_device_follows_learner_device( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = _make_device_runner( + monkeypatch, + device="cuda:3", + sim_backend="mjwarp", + ) + + assert runner.device == "cuda:3" + assert runner.collector_backend_device == "cuda:3" + assert runner.runtime_manifest["collector_accelerator_context"] is True + assert runner.runtime_manifest["collector_backend_device"] == "cuda:3" + + +def test_mjwarp_collector_start_forwards_learner_device( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + _FakePipeline.close_calls = 0 + monkeypatch.setattr(device_runner_module, "ReplayBuffer", _FakeReplayBuffer) + monkeypatch.setattr(device_runner_module, "GPUResidentReplayPipeline", _FakePipeline) + monkeypatch.setattr(device_runner_module, "OffPolicyLogger", _FakeLogger) + monkeypatch.setattr(device_runner_module.torch, "save", lambda *args, **kwargs: None) + monkeypatch.setattr(device_runner_module.time, "sleep", lambda seconds: None) + + real_empty = torch.empty + + def empty_without_cuda(*args, **kwargs): + if str(kwargs.get("device", "")).startswith("cuda"): + kwargs["device"] = "cpu" + return real_empty(*args, **kwargs) + + monkeypatch.setattr(device_runner_module.torch, "empty", empty_without_cuda) + runner = _make_device_runner( + monkeypatch, + device="cuda:3", + sim_backend="mjwarp", ) + collector_kwargs = {} + + def capture_collector(*, target_fn, kwargs): + del target_fn + collector_kwargs.update(kwargs) + + monkeypatch.setattr(runner, "_start_collector", capture_collector) + runner.learn(max_iterations=0, save_interval=0, log_dir=str(tmp_path)) + + assert collector_kwargs["sim_backend"] == "mjwarp" + assert collector_kwargs["backend_device"] == "cuda:3" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") diff --git a/tests/algos/test_offpolicy_worker.py b/tests/algos/test_offpolicy_worker.py index 17b27352f..e2293856e 100644 --- a/tests/algos/test_offpolicy_worker.py +++ b/tests/algos/test_offpolicy_worker.py @@ -6,6 +6,7 @@ import pytest import torch +import unilab.algos.offpolicy.worker as worker_module from unilab.algos.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.algos.offpolicy.worker import ( _publish_inference_tick, @@ -16,6 +17,55 @@ ) +def test_collector_binds_backend_device_before_env_materialization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[tuple[str, ...]] = [] + + class _EnvMaterializedError(RuntimeError): + pass + + def configure(backend_type: str, device: str | None) -> str: + events.append(("bind", backend_type, str(device))) + return str(device) + + def make_env(name: str, **kwargs): + events.append(("make", name, str(kwargs["sim_backend"]))) + raise _EnvMaterializedError + + from unilab.base import registry + + monkeypatch.setattr(worker_module, "configure_backend_process_device", configure) + monkeypatch.setattr(worker_module, "apply_torch_thread_runtime", lambda *args, **kwargs: None) + monkeypatch.setattr(worker_module, "ensure_registries", lambda: None) + monkeypatch.setattr(worker_module, "apply_training_seed", lambda *args, **kwargs: None) + monkeypatch.setattr(registry, "make", make_env) + + with pytest.raises(_EnvMaterializedError): + worker_module._run_collector( + stop_event=None, + env_name="DummyEnv", + num_envs=2, + replay_buffer=None, + inference_slot=None, + inference_request_queue=None, + inference_response_queue=None, + algo_type="sac", + metrics_queue=None, + sim_backend="mjwarp", + backend_device="cuda:3", + env_cfg_override=None, + seed=None, + trace_enabled=False, + trace_thread_time=False, + ) + + assert events == [ + ("bind", "mjwarp", "cuda:3"), + ("make", "DummyEnv", "mjwarp"), + ] + + def test_inference_request_publish_timeout_is_explicit() -> None: requests: queue.Queue[int] = queue.Queue(maxsize=1) requests.put_nowait(0) diff --git a/tests/base/backend/test_process_device.py b/tests/base/backend/test_process_device.py new file mode 100644 index 000000000..f575f45d1 --- /dev/null +++ b/tests/base/backend/test_process_device.py @@ -0,0 +1,73 @@ +"""Cold-path backend process-device routing tests.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from unilab.base.backend.mjwarp import runtime as mjwarp_runtime +from unilab.base.backend.process_device import ( + configure_backend_process_device, + resolve_backend_process_device, +) + + +class _FakeWarpDevice: + def __init__(self, name: str, *, is_cuda: bool = True) -> None: + self.name = name + self.is_cuda = is_cuda + + def __str__(self) -> str: + return self.name + + +class _FakeWarp: + def __init__(self, *, is_cuda: bool = True) -> None: + self.is_cuda = is_cuda + self.set_calls: list[str] = [] + self.selected = _FakeWarpDevice("cpu", is_cuda=False) + + def set_device(self, device: str) -> None: + self.set_calls.append(device) + self.selected = _FakeWarpDevice(device, is_cuda=self.is_cuda) + + def get_device(self) -> _FakeWarpDevice: + return self.selected + + +def test_mjwarp_process_device_follows_rank_learner_device(monkeypatch: pytest.MonkeyPatch) -> None: + warp = _FakeWarp() + monkeypatch.setattr( + mjwarp_runtime, + "load_mjwarp_dependencies", + lambda: SimpleNamespace(warp=warp), + ) + + assert configure_backend_process_device("mjwarp", "cuda:3") == "cuda:3" + assert warp.set_calls == ["cuda:3"] + + +@pytest.mark.parametrize("backend_type", ["mujoco", "motrix", "drake"]) +def test_host_or_backend_owned_devices_do_not_receive_runner_binding(backend_type: str) -> None: + assert resolve_backend_process_device(backend_type, "cuda:2") is None + + +@pytest.mark.parametrize("device", [None, "cpu", "mps", "xpu:1"]) +def test_mjwarp_process_device_fails_closed_without_cuda(device: str | None) -> None: + with pytest.raises(ValueError, match="CUDA process device"): + resolve_backend_process_device("mjwarp", device) + + +def test_mjwarp_binding_rejects_non_cuda_warp_resolution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + warp = _FakeWarp(is_cuda=False) + monkeypatch.setattr( + mjwarp_runtime, + "load_mjwarp_dependencies", + lambda: SimpleNamespace(warp=warp), + ) + + with pytest.raises(RuntimeError, match="active CUDA Warp device"): + mjwarp_runtime.bind_mjwarp_process_device("cuda:1")