OpenXLA pin update

[yeounoh]

This moves the pin to

• strip_prefix = "xla-b166243711f71b0a55daa1eda36b1dc745886784",

and libtpu build to

• _libtpu_version = '0.1.dev20240213'

Locally tested, and

>>> import torch_xla.core.xla_model as xm
>>> xm.xla_device()
WARNING:root:PJRT is now the default runtime. For more information, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1707861898.453343   82066 pjrt_api.cc:100] GetPjrtApi was found for tpu at /root/.local/lib/python3.8/site-packages/libtpu/libtpu.so
I0000 00:00:1707861898.453434   82066 pjrt_api.cc:79] PJRT_Api is set for device type tpu
I0000 00:00:1707861898.453443   82066 pjrt_api.cc:146] The PJRT plugin has PJRT API version 0.40. The framework PJRT API version is 0.40.
device(type='xla', index=0)

[lsy323]

Thanks @yeounoh! openxla_patches/stablehlo_quant_seralization.diff is still needed to workaround openxla/stablehlo#1812

[yeounoh]

Thanks @yeounoh! openxla_patches/stablehlo_quant_seralization.diff is still needed to workaround openxla/stablehlo#1812

We can drop it, based on our verification

root@0e0864097658:/workspace/pytorch/xla# python test/stablehlo/test_pt2e_qdq.py 
2024-02-13 22:31:03.922291: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
WARNING:root:PJRT is now the default runtime. For more information, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1707863465.033336   91383 pjrt_api.cc:100] GetPjrtApi was found for tpu at /workspace/_libtpu.so
I0000 00:00:1707863465.033427   91383 pjrt_api.cc:79] PJRT_Api is set for device type tpu
I0000 00:00:1707863465.033435   91383 pjrt_api.cc:146] The PJRT plugin has PJRT API version 0.40. The framework PJRT API version is 0.40.
..s
----------------------------------------------------------------------
Ran 3 tests in 3.115s

OK (skipped=1)

[lsy323]

Thanks @yeounoh!

[lsy323]

Not sure if we should check the resnet18 speed? Remembered it's verified in previous pin update. Here is script of previous pin update

[JackCaoG]

Let's do a quick run of resnet to make sure we don't have any visible regression.

[yeounoh]

With default batch size 128

| Training Device=xla:0/0 Epoch=2 Step=1200 Loss=0.00135 Rate=1784.07 GlobalRate=1781.73 Time=18:48:34
| Training Device=xla:0/1 Epoch=2 Step=1200 Loss=0.00135 Rate=1784.08 GlobalRate=1781.74 Time=18:48:34
| Training Device=xla:0/2 Epoch=2 Step=1200 Loss=0.00135 Rate=1783.65 GlobalRate=1781.69 Time=18:48:34

With batch size 256

| Training Device=xla:0/1 Epoch=2 Step=160 Loss=0.00135 Rate=1818.24 GlobalRate=1806.78 Time=18:54:50
| Training Device=xla:0/2 Epoch=2 Step=180 Loss=0.00135 Rate=1818.24 GlobalRate=1807.91 Time=18:54:53
| Training Device=xla:0/3 Epoch=2 Step=180 Loss=0.00135 Rate=1818.23 GlobalRate=1808.03 Time=18:54:53
| Training Device=xla:0/1 Epoch=2 Step=180 Loss=0.00135 Rate=1818.29 GlobalRate=1808.04 Time=18:54:53
| Training Device=xla:0/0 Epoch=2 Step=180 Loss=0.00135 Rate=1817.99 GlobalRate=1807.90 Time=18:54:53

No visible regression on ResNet..