PR #23347: [gpu] Allow explicitly setting slice_index in se_gpu_pjrt_client #24503
Conversation
This file contains 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
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
…client Imported from GitHub PR #23347 Allows overriding the slice index used by se_gpu_pjrt_client. More explicit control over which slice a device ends up in is desirable: - Various parts of the ecosystem equate slices with "devices communicating via fast interconnect". With the arrival of NVL72 we want devices managed by multiple hosts to form a single slice. - For debugging purposes it can be useful to allow devices on the same host (managed in separate processes) to be treated as different slices. For example, [Orbax](https://github.com/google/orbax)'s local checkpointing presumes the existence of at least two slices, so overriding the boot id will allow us to test local checkpointing on a single host. (Companion PR in JAX: jax-ml/jax#26906) Copybara import of the project: -- 8d16790 by Georg Stefan Schmid <[email protected]>: [gpu] Allow overriding XLA slice_index Merging this change closes #23347 COPYBARA_INTEGRATE_REVIEW=#23347 from gspschmid:gschmid/xla-override-boot-id 8d16790 PiperOrigin-RevId: 743085978
copybara-service
bot
force-pushed
the
test_743055829
branch
from
e5011db to
c3087e0
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
PR #23347: [gpu] Allow explicitly setting slice_index in se_gpu_pjrt_client
Imported from GitHub PR #23347
Allows overriding the slice index used by se_gpu_pjrt_client.
More explicit control over which slice a device ends up in is desirable:
(Companion PR in JAX: jax-ml/jax#26906)
Copybara import of the project:
--
8d16790 by Georg Stefan Schmid [email protected]:
[gpu] Allow overriding XLA slice_index
Merging this change closes #23347
FUTURE_COPYBARA_INTEGRATE_REVIEW=#23347 from gspschmid:gschmid/xla-override-boot-id 8d16790