[Newton] Fix multi-GPU initialization on non-default devices#6322
Conversation
Pin Torch and Warp to the simulation device before Newton allocates GPU state, and bind CUDA graph capture to that same device. This prevents nonzero distributed ranks from starting capture on Warp's default cuda:0 context.
github-actions
Bot
added
bug
Greptile SummaryThis PR fixes multi-GPU Newton initialization by pinning the Torch and Warp device contexts to
Confidence Score: 4/5Safe to merge; the change is a targeted, well-commented initialization fix with an accompanying regression test and a full Newton manager test run reported as passing. The core fix in wp.ScopedCapture(device=device) is straightforward and the device-pinning guards in start_simulation and initialize_solver are idempotent. The one gap is that _capture_or_defer_graph itself has no guard, so if it is called from set_decimation on the Newton-actuator path without a prior pinning call still in effect, allocations inside that recapture window could land on the wrong device. The new unit test validates the capture-device binding but does not cover the torch.cuda.set_device / wp.set_device side-effects directly. newton_manager.py around _capture_or_defer_graph — verify whether the set_decimation call path guarantees device context is already current before entering this method on Newton-actuator workloads. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Rank as Non-zero Rank (cuda:N)
participant SM as start_simulation()
participant IS as initialize_solver()
participant CG as _capture_or_defer_graph()
participant Torch as torch.cuda
participant Warp as wp (Warp)
Note over Rank: Before fix: default context = cuda:0
Rank->>SM: start_simulation()
SM->>Torch: set_device("cuda:N")
SM->>Warp: wp.set_device("cuda:N")
SM->>IS: initialize_solver()
IS->>Torch: set_device("cuda:N") [idempotent]
IS->>Warp: wp.set_device("cuda:N") [idempotent]
IS->>CG: _capture_or_defer_graph()
CG->>Warp: "wp.ScopedCapture(device="cuda:N")"
Note over Warp: Graph capture bound to cuda:N
Warp-->>CG: capture.graph
CG-->>IS: _graph set
IS-->>SM: solver + graph ready
SM-->>Rank: initialized on cuda:N
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Rank as Non-zero Rank (cuda:N)
participant SM as start_simulation()
participant IS as initialize_solver()
participant CG as _capture_or_defer_graph()
participant Torch as torch.cuda
participant Warp as wp (Warp)
Note over Rank: Before fix: default context = cuda:0
Rank->>SM: start_simulation()
SM->>Torch: set_device("cuda:N")
SM->>Warp: wp.set_device("cuda:N")
SM->>IS: initialize_solver()
IS->>Torch: set_device("cuda:N") [idempotent]
IS->>Warp: wp.set_device("cuda:N") [idempotent]
IS->>CG: _capture_or_defer_graph()
CG->>Warp: "wp.ScopedCapture(device="cuda:N")"
Note over Warp: Graph capture bound to cuda:N
Warp-->>CG: capture.graph
CG-->>IS: _graph set
IS-->>SM: solver + graph ready
SM-->>Rank: initialized on cuda:N
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def test_cuda_graph_capture_uses_simulation_device(monkeypatch): | ||
| """CUDA graph capture should use the simulation device instead of Warp's default device.""" | ||
| from isaaclab.physics import PhysicsManager | ||
|
|
||
| captured_devices = [] | ||
| captured_graph = object() | ||
|
|
||
| class FakeScopedCapture: | ||
| def __init__(self, device=None): | ||
| captured_devices.append(device) | ||
| self.graph = captured_graph | ||
|
|
||
| def __enter__(self): | ||
| return self | ||
|
|
||
| def __exit__(self, exc_type, exc_value, traceback): | ||
| return False | ||
|
|
||
| monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) | ||
| monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False) | ||
| monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) | ||
| monkeypatch.setattr(NewtonManager, "_solver", None, raising=False) | ||
| monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False)) | ||
| monkeypatch.setattr(NewtonManager, "_simulate_physics_only", classmethod(lambda cls: None)) | ||
| monkeypatch.setattr(wp, "ScopedCapture", FakeScopedCapture) | ||
|
|
||
| NewtonManager._capture_or_defer_graph() | ||
|
|
||
| assert captured_devices == ["cuda:1"] | ||
| assert NewtonManager._graph is captured_graph |
There was a problem hiding this comment.
initialize_solver / start_simulation
The new test verifies that wp.ScopedCapture receives device="cuda:1", but the two torch.cuda.set_device + wp.set_device guards added in start_simulation and initialize_solver are not exercised. If those guards were accidentally removed or made conditional on the wrong predicate, the regression would go undetected — the test would still pass. Consider adding a parallel test that monkeypatches torch.cuda.set_device and wp.set_device, sets PhysicsManager._device = "cuda:1", calls the relevant entry points, and asserts both were invoked with the correct device string.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Select the resolved CUDA device in both runtimes at shared launcher and simulation initialization boundaries. Keep Newton graph capture explicitly bound to the simulation device without backend-specific default-device setup.
…n-mgpu-capture-device-codex
Synchronize Torch and Warp in the base physics manager before backend-specific setup. Keep the helper internal and cover the launcher and backend initialization boundaries directly.
…n-mgpu-capture-device-codex # Conflicts: # source/isaaclab/test/app/test_kwarg_launch.py
Reconcile against the split-out changes that have since merged: - Drop the renderer-pinning ISAACLAB_PIN_KIT_GPU block and its changelog fragment; superseded by isaac-sim#5933 (default one-GPU + ISAACLAB_FABRIC_USE_GPU_INTEROP). - Drop the inline newton_manager.py torch/Warp device pins and take develop's version; superseded by isaac-sim#6322 (centralized set_cuda_device + ScopedCapture(device=device)). - Resolve test_views_xform_prim_fabric.py: keep develop's isaac-sim#5677/isaac-sim#5380 test bodies, retain the device_scope test_devices() parametrization.
The Newton non-default-device init fix moved to its own PR (isaac-sim#6322, merged), so this PR no longer makes a user-facing isaaclab_newton change -- only device_scope test migrations. Convert the fragment to .skip.
…overage (#5823) ## 1. Summary - **CI:** Adds a dedicated multi-GPU pytest workflow that runs one shard per non-default `cuda:N`, with a shared atomic work queue, per-file reports, and end-of-run reconciliation. - **isaaclab:** Adds composable `DeviceScope` flags and `test_devices()` for `scope ∩ runtime` parametrization while retaining exact custom string masks. - **isaaclab:** Uses `ISAACLAB_TEST_DEVICES` as the single source for both test parametrization and explicit Kit launch-device selection; `AppLauncher` has no implicit test-environment override. - **isaaclab:** Migrates 19 device-parametrized test modules across the core, Newton, OV PhysX, and PhysX suites while preserving the existing cpu + cuda:0 single-GPU behavior. ## 2. Device selection ```python from isaaclab.test.utils import DeviceScope, test_devices test_devices() # cpu + every GPU test_devices(DeviceScope.CUDA) # every GPU test_devices(DeviceScope.CPU | DeviceScope.NON_DEFAULT_CUDA) # composed scope test_devices("101X") # exact custom mask ``` - `ISAACLAB_TEST_DEVICES` limits which devices a run may use; unset preserves the historical cpu + cuda:0 runtime. - The multi-GPU lane discovers non-default-capable scopes and narrows each shard to one concrete device. ## 3. Dependencies All prerequisite changes have merged into `develop`, so this PR no longer carries them — it is now scoped to the device-selection infrastructure and the multi-GPU pytest workflow only: - #5695 — Cross-platform Part 1 (base pytest markers / shared scaffolding). - #5881 — Sim honors the device kwarg over `sim_cfg.device` in `build_simulation_context`. - #5933 — Kit renderer defaults to one GPU + adds `ISAACLAB_FABRIC_USE_GPU_INTEROP`; the multi-GPU lane sets `ISAACLAB_FABRIC_USE_GPU_INTEROP=0`. - #6322 — Newton multi-GPU initialization on non-default CUDA devices (`cuda:1` and higher). ## 4. Test plan - [x] Device-selection unit suite passes. - [x] Changed Python modules compile; multi-GPU shell scripts pass `bash -n`. - [x] Full documentation build (`./isaaclab.sh -d`). - [x] Full pre-commit suite (`./isaaclab.sh -f`). - [x] GitHub Actions green on latest `develop` (one `rendering-correctness-kitless` flake cleared on rerun; fixed upstream on `develop` by #6431).
2 participants
1. Summary
Fix CUDA initialization on non-default devices by selecting the resolved device in PyTorch before synchronizing Warp at shared launcher and physics-backend initialization boundaries. Newton binds CUDA graph capture explicitly to the simulation device without backend-specific default-device setup.
The exact three-rank camera workload failed 2/2 without device-scoped capture and has now passed 6/6 with the fix, including two runs on the final centralized implementation and one after merging current
develop.2. Root cause
Warp used its default
cuda:0context while nonzero distributed ranks used their assigned devices. Newton's collision pipeline then mixed allocations from different CUDA contexts, producing CUDA errors 900 and 101 during graph capture.3. Design
AppLauncher.sim_launcher.PhysicsManager.initializebefore backend-specific setup, covering directSimulationContextuse for every backend.wp.ScopedCapture(device=device)in Newton so graph ownership remains explicit.4. Related work
Extracted from #5823 so the product fix can be reviewed and landed independently of the multi-GPU CI infrastructure.
5. Test plan
develop./isaaclab.sh -f