Skip to content

fix: bind mjwarp collectors to rank-local GPUs - #1291

Merged
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
fix/mjwarp-dp-collector-device
Aug 24, 2026
Merged

fix: bind mjwarp collectors to rank-local GPUs#1291
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
fix/mjwarp-dp-collector-device

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

Summary

  • bind each mjwarp rank process and collector process to the CUDA device assigned to that rank's learner before environment materialization
  • keep backend-specific Warp setup behind a cold-path backend adapter and leave MuJoCo, Motrix, and Drake as no-ops
  • expose the local collector backend device in the runtime manifest and document the multi-GPU placement contract

Before this change, learners followed training.devices, but fresh collector subprocesses let Warp default to logical cuda:0, so nonzero-rank mjwarp collectors could concentrate on GPU 0.

Linked Work

  • Issue: N/A (direct maintainer request)
  • Milestone: N/A

Validation

  • make check
  • uv run pytest -m "not slow"
  • Additional task-specific validation listed below

Commands actually run:

make test-all
uv run --extra mjwarp train --algo sac --task g1_walk_flat --sim mjwarp 'training.devices=[0,1]' algo.num_envs=16 algo.batch_size=16 algo.replay_buffer_n=4 algo.updates_per_step=1 algo.learning_starts=1 algo.max_iterations=1 algo.save_interval=1 algo.algo_params.use_compile=false training.use_amp=false training.no_play=true training.log_dir=/tmp/unilab-mjwarp-dp-smoke.Bi1UyO
UNILAB_DOCS_SKIP_AUTODOC=1 uv run --no-project --with-requirements requirements.txt sphinx-build -b html -n source build/html

Results:

  • make test-all: 2239 passed, 28 skipped, 278 deselected, 1 xfailed; ruff, mypy, pyright, and benchmark smoke passed
  • two-rank RTX 6000D mjwarp SAC smoke: Warp materialized on both cuda:0 and cuda:1; run summary completed after one iteration
  • Sphinx HTML build completed successfully

Impact

  • Backend impact: mjwarp device placement; MuJoCo / Motrix / Drake remain no-op
  • Platform impact: Linux CUDA
  • Training effect expected: yes; each mjwarp collector now shares its learner rank's GPU

Artifacts

  • W&B: N/A
  • benchmark result: N/A; short correctness smoke only
  • video / screenshot: N/A
  • ONNX / checkpoint: temporary smoke checkpoint only

Checklist

  • Added or updated tests where needed
  • Updated docs if behavior or workflow changed
  • Linked the driving issue — N/A (direct maintainer request)
  • Noted follow-up work explicitly: no committed mjwarp multi-GPU throughput or convergence benchmark is claimed

@TATP-233
TATP-233 requested a review from caozx1110 as a code owner August 24, 2026 15:12
@TATP-233
TATP-233 merged commit 49c3edb into dev/issue-1042-manager-based-api Aug 24, 2026
7 checks passed
@TATP-233
TATP-233 deleted the fix/mjwarp-dp-collector-device branch August 24, 2026 15:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant