fix: bind mjwarp collectors to rank-local GPUs

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.

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 不可靠，

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:

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),

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:

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)

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)

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

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")