Conditionally set default TPU settings in __init__.py
#5696
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
JackCaoG
approved these changes
Oct 10, 2023
alanwaketan
reviewed
Oct 10, 2023
torch_xla/_internal/tpu.py
Outdated
| @@ -266,6 +266,9 @@ def configure_topology(local_rank: int, | |||
| os.environ.setdefault(xenv.TPU_VISIBLE_CHIPS, str(local_rank)) | |||
| os.environ.setdefault(xenv.TPU_PROCESS_PORT, str(ports[local_rank])) | |||
|
|
|||
| if version() == 4: | |||
There was a problem hiding this comment.
Is it guaranteed that this will be called all the time and will be before any modeling code involved tpu?
There was a problem hiding this comment.
Actually, this will only get called for sure when using torchrun or xmp.spawn.. Good catch.
It makes the most logical sense to me to set these up during client initialization, but that happens in C++ and all of the TPU utilities are in Python. I'll make this setting conditional in __init__.py for now.
will-cromar
changed the title
__init__.py
alanwaketan
approved these changes
Oct 10, 2023
zpcore
pushed a commit
that referenced
this pull request
Oct 19, 2023
* Set TPU_MEGACORE in configure_topology * remove * Move back to __init__.py
ghpvnist
pushed a commit
to ghpvnist/xla
that referenced
this pull request
Oct 31, 2023
* Set TPU_MEGACORE in configure_topology * remove * Move back to __init__.py
mbzomowski
pushed a commit
to mbzomowski-test-org/xla
that referenced
this pull request
Nov 16, 2023
* Set TPU_MEGACORE in configure_topology * remove * Move back to __init__.py
chunnienc
pushed a commit
to chunnienc/xla
that referenced
this pull request
Dec 14, 2023
* Set TPU_MEGACORE in configure_topology * remove * Move back to __init__.py
golechwierowicz
pushed a commit
that referenced
this pull request
Jan 12, 2024
* Set TPU_MEGACORE in configure_topology * remove * Move back to __init__.py
bhavya01
pushed a commit
that referenced
this pull request
Apr 22, 2024
* Set TPU_MEGACORE in configure_topology * remove * Move back to __init__.py
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.
Replace
_set_missing_envwith standard library call.Only set
TPU_MEGACOREon TPU v4. Otherwise, use the default value.