Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Set --xla_latency_hiding_scheduler_rerun to 1 #5736

Merged
merged 4 commits into from
Oct 26, 2023
Merged

Set --xla_latency_hiding_scheduler_rerun to 1 #5736

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
90d8774
Set --xla_latency_hiding_scheduler_rerun to 1
alanwaketan Oct 25, 2023
a0e0d3d
Adds more comment
alanwaketan Oct 25, 2023
c7f477b
Introduce LIBTPU_INIT_ARGS
alanwaketan Oct 26, 2023
4afa5d5
Fix linters
alanwaketan Oct 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions torch_xla/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,26 @@ def _setup_xla_flags():
os.environ['XLA_FLAGS'] = ' '.join(flags)


def _setup_libtpu_flags():
flags = os.environ.get('LIBTPU_INIT_ARGS', '').split(' ')
# This flag will rerun the latency hidding scheduler if the default
# shared memory limit 95% leads to OOM. Each rerun will choose a value
# 0.9x of the previous run, and the number of rerun is set to 1 now.
# Shared memory limit refers to --xla_tpu_scheduler_percent_shared_memory_limit.
# Lower shared memory limit means less communiation and computation overlapping,
# and thus worse performance.
flags = _set_missing_flags(flags,
(('xla_latency_hiding_scheduler_rerun', '1'),))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! I was actually wondering about this, we set XLA_FLAGS here, but in other cases we pass the flags through LIBTPU_INIT_ARGS. Do you know if there's a difference? If you saw the appropriate log output from the test, it seems both work...

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh right.. sorry I missed this. so for everything in compiler/xla/xla.proto we used XLA_FLAGS. xla_latency_hiding_scheduler_rerun is one of those TPU specified flags that needs to be passed in with LIBTPU_INIT_ARGS.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. I miss this lol.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, makes sense that LIBTPU_INIT_ARGS would be tpu-specific lol, thanks @JackCaoG. Is there a rule of thumb to tell which flag goes where? I'm thinking in terms of hashing the compilation environment, I suppose we'll just need to ensure both env vars are included.

os.environ['LIBTPU_INIT_ARGS'] = ' '.join(flags)


def _setup_default_env():
os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '1')
os.environ.setdefault('GRPC_VERBOSITY', 'ERROR')

if tpu.num_available_chips() > 0:
_setup_libtpu_flags()

os.environ.setdefault('ALLOW_MULTIPLE_LIBTPU_LOAD', '1')
os.environ.setdefault('TPU_ML_PLATFORM', 'PyTorch/XLA')

Expand Down