[Pin Update] Version 20230826

[alanwaketan]

Summary:
This change bumps libTPU.so to 20230826 and update the corresponding openxla/xla pin.

Test Plan:
CI.

[alanwaketan]

For ResNet MP on V4-8:
After the change:

| Training Device=xla:0/1 Epoch=1 Step=2340 Loss=0.00135 Rate=1776.80 GlobalRate=993.25 Time=02:13:37

Before the change:

| Training Device=xla:0/0 Epoch=1 Step=260 Loss=0.01169 Rate=1773.55 GlobalRate=221.43 Time=02:31:15

For ResNet SPMD on V4-8:
After the change:

| Training Device=xla:0/0 Epoch=1 Step=4680 Loss=0.00135 Rate=1882.09 GlobalRate=1552.79 Time=02:22:06

Before the change:

| Training Device=xla:0/0 Epoch=1 Step=1020 Loss=0.00385 Rate=1881.47 GlobalRate=951.80 Time=02:27:09

No performance regression. Will do a LLaMA2 training test later.

[lsy323]

LGTM

[alanwaketan]

Looks like one of the GPU test has failed. There could be a regression on GPU. Will need to setup a GPU env to further debug.

[JackCaoG]

//test/cpp:test_tensor timed out, wondering if it is a real error or some flakeyness. I bump up the test timeout for GPU to 240 minutes in 76495f0 so I think it is not the workflow itself timedout but that specified test has some timeout threashold.

[alanwaketan]

//test/cpp:test_tensor timed out, wondering if it is a real error or some flakeyness. I bump up the test timeout for GPU to 240 minutes in 76495f0 so I think it is not the workflow itself timedout but that specified test has some timeout threashold.

Yea, I need to setup a GPU env to repro.

[wonjoolee95]

@ManfeiBai, wondering if you have any GPU that we can use to repro this?

[ManfeiBai]

@ManfeiBai, wondering if you have any GPU that we can use to repro this?

yes, we have some GPU to repro this, let me pull to these GPU instance so that we could repro this

[wonjoolee95]

Thanks @ManfeiBai, let me know if you need any help.

[ManfeiBai]

this might be the GPU CI issue log: https://source.cloud.google.com/results/invocations/a2d48ef1-784f-43d8-b4e5-69caa341e906/targets/%2F%2Ftest%2Fcpp:test_tensor/log

[alanwaketan]

76495f0

this might be the GPU CI issue log: https://source.cloud.google.com/results/invocations/a2d48ef1-784f-43d8-b4e5-69caa341e906/targets/%2F%2Ftest%2Fcpp:test_tensor/log

So you are suggesting the time out could be unrelated? and we should increase the test time out limit?

[ManfeiBai]

76495f0

yes, agree with the increase the test time out limit plan, and let's also trigger GPU test again to confirm issue consistantly too?

[alanwaketan]

76495f0

yes, agree with the increase the test time out limit plan, and let's also trigger GPU test again to confirm issue consistantly too?

The issue is consistent. I have ran it twice.

[ManfeiBai]

76495f0

yes, agree with the increase the test time out limit plan, and let's also trigger GPU test again to confirm issue consistantly too?

The issue is consistent. I have ran it twice.

then I would agree with the increase the test time out limit plan, it may help us for more info

[wonjoolee95]

Until we can debug the GPU CI for certain, we should probably not cherry-pick this into the r2.1 branch. That means we'll need to re-open your PR to fix the FSDP GPU unit tests, @ManfeiBai. Let me know if you think otherwise, @alanwaketan.

[alanwaketan]

Until we can debug the GPU CI for certain, we should probably not cherry-pick this into the r2.1 branch. That means we'll need to re-open your PR to fix the FSDP GPU unit tests, @ManfeiBai. Let me know if you think otherwise, @alanwaketan.

Let me set the time out now. If it doesn't work, then let's re-open Manfei's PR.

[alanwaketan]

It seems we need a longer time limit. @wonjoolee95 @JackCaoG @ManfeiBai Please take a look to see if the current way is preferable?

[ManfeiBai]

LGTM

[alanwaketan]

Thanks all for reviewing!