Skip to content

[PJRT] Use PJRT by default if XRT is not configured - #4599

Merged
will-cromar merged 3 commits into
masterfrom
wcromar/auto-pjrt
Feb 10, 2023
Merged

will-cromar merged 3 commits into
masterfrom
wcromar/auto-pjrt

Conversation

@will-cromar

Copy link
Copy Markdown
Collaborator

If you don't set XRT configuration today, we print a warning that there is no XLA configuration. Instead, let's set the default device to PJRT:CPU or PJRT:TPU, depending on what is available.

Existing workloads that already set up XRT configuration will be unaffected by this change. However, we'll still print a warning for now to let the user know that they're using PJRT instead of XRT.

Example output:

# With libtpu installed:
$ python -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices('TPU'))"
WARNING:torch_xla:XRT configuration not detected. Defaulting to preview PJRT runtime. To silence this warning and continue using PJRT, explicitly set PJRT_DEVICE to a supported device. To use XRT, set any of the following environment variables: ['XRT_TPU_CONFIG', 'XRT_DEVICE_MAP', 'XRT_WORKERS']
WARNING:torch_xla:For more information about the status of PJRT, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:torch_xla:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
['xla:0', 'xla:1', 'xla:2', 'xla:3']
$ PJRT_DEVICE=TPU python -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices('TPU'))"
['xla:0', 'xla:1', 'xla:2', 'xla:3']

# With libtpu uninstalled
$ python -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices('CPU'))"
WARNING:torch_xla:XRT configuration not detected. Defaulting to preview PJRT runtime. To silence this warning and continue using PJRT, explicitly set PJRT_DEVICE to a supported device. To use XRT, set any of the following environment variables: ['XRT_TPU_CONFIG', 'XRT_DEVICE_MAP', 'XRT_WORKERS']
WARNING:torch_xla:For more information about the status of PJRT, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:torch_xla:Defaulting to PJRT_DEVICE=CPU
['xla:0']
Comment thread torch_xla/__init__.py Outdated

XRT_RUN_SERVER_PROCESS = 'torch_xla.core._xrt_run_server'
XRT_SERVER_REGEX = '^python3 -m {} [0-9]+$'.format(XRT_RUN_SERVER_PROCESS)
XRT_CONFIG_ENV_VARS = ['XRT_TPU_CONFIG', 'XRT_DEVICE_MAP', 'XRT_WORKERS']

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

currently if user sets GPU_NUM_DEVICES it will default to xrt:gpu, user have to overwrite that with additional PJRT_DEVICE=GPU. I guess we intentionally don't change that behavior there?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Right. I left GPU out intentionally because we don't have the same easy way of knowing that the user intended XRT, and there are still some feature gaps on PJRT GPU.

@will-cromar
will-cromar merged commit 43ee2eb into master Feb 10, 2023
chandrasekhard2 pushed a commit that referenced this pull request Feb 22, 2023
* Use PJRT runtime by default

* formatting

* Don't set the runtime when `GPU_NUM_DEVICES` is present
chandrasekhard2 pushed a commit that referenced this pull request Feb 22, 2023
* Use PJRT runtime by default

* formatting

* Don't set the runtime when `GPU_NUM_DEVICES` is present
mateuszlewko pushed a commit that referenced this pull request Mar 15, 2023
* Use PJRT runtime by default

* formatting

* Don't set the runtime when `GPU_NUM_DEVICES` is present
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

2 participants