[ROCm]---Theory--- force ROCm xdist worker GPU env override#749
Open
magaonka-amd wants to merge 8 commits into
Open
[ROCm]---Theory--- force ROCm xdist worker GPU env override#749magaonka-amd wants to merge 8 commits into
magaonka-amd wants to merge 8 commits into
Conversation
magaonka-amd
force-pushed
the
fix/conftest-rocm-xdist-env-override
branch
4 times, most recently
from
2f7fd82 to
3a6d174
Compare
Use os.environ assignment for ROCR_VISIBLE_DEVICES (not setdefault) so CI runner pre-injected visibility (e.g. gha-gpu-isolation-settings) cannot leave every xdist worker seeing all GPUs. ROCm pinning uses ROCR only; do not set HIP_VISIBLE_DEVICES in conftest. Allow ROCm plugin/PJRT wheel builds from this repo again (remove build.py guard) so fork/PR wheel continuous workflows match jax-ml behavior. Plugin releases remain documented at ROCm/rocm-jax.
Instead of pinning each worker to a single GPU (which skips all multi-GPU tests), rotate the device ordering so every worker sees all GPUs but with a different primary device. Example with 4 GPUs: gw0: ROCR_VISIBLE_DEVICES=0,1,2,3 gw1: ROCR_VISIBLE_DEVICES=1,2,3,0 gw2: ROCR_VISIBLE_DEVICES=2,3,0,1 gw3: ROCR_VISIBLE_DEVICES=3,0,1,2 This spreads single-GPU test load across physical devices while keeping multi-GPU tests runnable. Reads existing ROCR_VISIBLE_DEVICES from CI runner isolation settings and rotates within that subset.
Ensures import jax loads from the checkout (jax/ under the repo root) before any jax in site-packages, matching tests/ to the tree that contains conftest.py and internal modules like jax._src.*.
The PallasCallMultimemTest and PallasCallMultimemThreadUnsafeTest classes access cuda_versions.cuda_supports_multicast() in setUp without first checking that cuda_versions is not None. On ROCm (and other non-CUDA platforms) this raises an AttributeError. Add an early skip when cuda_versions is None since multicast memory is a CUDA-specific hardware feature.
magaonka-amd
force-pushed
the
fix/conftest-rocm-xdist-env-override
branch
from
53448ad to
5084c96
Compare
Set --xla_force_host_platform_device_count=$gpu_count so the CPU backend creates enough logical devices for colocated_python tests that need one CPU device per GPU.
This reverts commit 6d67077.
mminutoli
force-pushed
the
amd-main
branch
4 times, most recently
from
40c94c1 to
b5d7662
Compare
mminutoli
force-pushed
the
amd-main
branch
3 times, most recently
from
28ad382 to
e3a4195
Compare
psanal35
force-pushed
the
amd-main
branch
2 times, most recently
from
729ed31 to
88e42dc
Compare
mminutoli
force-pushed
the
amd-main
branch
2 times, most recently
from
cb097f7 to
d347a4f
Compare
mminutoli
force-pushed
the
amd-main
branch
2 times, most recently
from
3f9e8b4 to
f25f2b4
Compare
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
3 participants
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.
Use os.environ assignment for ROCR_VISIBLE_DEVICES (not setdefault) so CI runner pre-injected visibility (e.g. gha-gpu-isolation-settings) cannot leave every xdist worker seeing all GPUs.
Set HIP_VISIBLE_DEVICES=0 after narrowing ROCR so workers do not inherit a multi-device HIP filter from the parent environment.