fix: route GPU backend environments by DP rank - #1510
Merged
Conversation
…DA init Genesis (Quadrants) ignores torch.cuda.set_device and only honors the first CUDA_VISIBLE_DEVICES entry, crashing on non-zero physical GPUs. Pin the visible-device set at the earliest entry-point boundary (pure config-topology resolution, no torch queries), remap the process-local device to cuda:0 (rewriting LOCAL_RANK for rsl_rl DP and genesis_device_id overrides so unisim does not re-translate), and keep binding idempotent via a module-level latch. Capture the bound device returned by bind_backend_process_device in all train/play entry points.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #1508 by routing IsaacGym, IsaacSim, and Genesis environment device IDs from the configured data-parallel topology.
training.devicesindices.LOCAL_RANKin the remappedCUDA_VISIBLE_DEVICESnamespace.gs.init, including spawn collectors.unisim-coreversions that rejectgenesis_device_id.Genesis pinning update (a404d2d)
Genesis (Quadrants 1.3.0) ignores
torch.cuda.set_deviceand only honors the firstCUDA_VISIBLE_DEVICESentry, so routing alone crashed on any non-zero physical GPU withCUDA_ERROR_ILLEGAL_ADDRESS. Entry points now pinCUDA_VISIBLE_DEVICESto the routed physical GPU before any CUDA query (torch.cuda.is_available()itself latches the visible set, so pinning must happen first), remap the process-local device tocuda:0(rewritingLOCAL_RANKfor rsl_rl DP andgenesis_device_idoverrides), and keep binding idempotent via a module-level latch.Validation
246 passed; new Genesis pin tests included.ruff,mypy, andpyright: passed (one pre-existing optionaldrake_uni.runtimewarning).make test-allon final head a404d2d: passed —1595 passed, 20 skipped, benchmark smoke passed (env synced with--extra mujoco --extra motrixper CI).g1_walk_flat/genesisboth exit 0;nvidia-smiconfirms each rank's learner and collector sim processes land on their own physical GPU with no cross-GPU leakage. Single-GPU Genesis PPO/SAC regressions pass. IsaacSim PPO/SAC DP=2 verified earlier on the same host; IsaacGym blocked by Preview 4 lacking Blackwell kernels (single-GPU also fails; routing delivery to workers confirmed).No package release is performed in this PR.
Dependency PR
The Genesis public contract is implemented in the companion UniSim PR: unilabsim/unisim#29.