Xla: Evalution part too slow on a single device

[vfdev-5]

❓ Questions and Help

Could you, please, hint why evaluation part takes 4 minutes (training part ~20 seconds) for the code based almost the same as test_train_mp_imagenet.py (the only replacement is dataset and lr scheduler).

Code is executed on a single device without multi-processing (however, metrics mesh reduce ops are kept).

Here is a part of logs:

==> Preparing data..
Epoch 1 train begin 11:11:18
| Training Device=xla:1 Epoch=1 Step=0 Loss=7.34335 Rate=259.95 GlobalRate=259.94 Time=11:11:20
| Training Device=xla:1 Epoch=1 Step=10 Loss=4.03379 Rate=3121.23 GlobalRate=1885.00 Time=11:11:21
| Training Device=xla:1 Epoch=1 Step=20 Loss=2.15459 Rate=2725.02 GlobalRate=2121.40 Time=11:11:23
| Training Device=xla:1 Epoch=1 Step=30 Loss=1.91277 Rate=3358.41 GlobalRate=2471.27 Time=11:11:24
| Training Device=xla:1 Epoch=1 Step=40 Loss=1.84981 Rate=2995.99 GlobalRate=2534.82 Time=11:11:26
| Training Device=xla:1 Epoch=1 Step=50 Loss=1.63390 Rate=3028.22 GlobalRate=2621.60 Time=11:11:28
| Training Device=xla:1 Epoch=1 Step=60 Loss=1.71373 Rate=2849.56 GlobalRate=2638.85 Time=11:11:30
| Training Device=xla:1 Epoch=1 Step=70 Loss=1.70562 Rate=3236.19 GlobalRate=2733.06 Time=11:11:31
| Training Device=xla:1 Epoch=1 Step=80 Loss=1.81418 Rate=2807.94 GlobalRate=2705.17 Time=11:11:33
| Training Device=xla:1 Epoch=1 Step=90 Loss=1.65861 Rate=3147.11 GlobalRate=2765.35 Time=11:11:35
Epoch 1 train end 11:11:35 - elapsed: 17.101221084594727

| Test Device=xla:1 Step=0 Epoch=1 Time=11:11:37
| Test Device=xla:1 Step=10 Epoch=1 Time=11:13:36
Epoch 1 test end 11:15:23, Accuracy=41.09    <<<<<<<< 11:15:23  <--- 11:11:35

Epoch 2 train begin 11:15:23
| Training Device=xla:1 Epoch=2 Step=0 Loss=1.56741 Rate=322.85 GlobalRate=322.85 Time=11:15:25
| Training Device=xla:1 Epoch=2 Step=10 Loss=1.50487 Rate=2304.56 GlobalRate=1878.57 Time=11:15:26
| Training Device=xla:1 Epoch=2 Step=20 Loss=1.59329 Rate=2402.65 GlobalRate=2119.65 Time=11:15:28
| Training Device=xla:1 Epoch=2 Step=30 Loss=1.39935 Rate=2864.46 GlobalRate=2373.74 Time=11:15:30
| Training Device=xla:1 Epoch=2 Step=40 Loss=1.51320 Rate=2731.99 GlobalRate=2434.36 Time=11:15:32
| Training Device=xla:1 Epoch=2 Step=50 Loss=1.44603 Rate=3000.97 GlobalRate=2551.72 Time=11:15:34
| Training Device=xla:1 Epoch=2 Step=60 Loss=1.41598 Rate=2638.81 GlobalRate=2525.06 Time=11:15:36
| Training Device=xla:1 Epoch=2 Step=70 Loss=1.33224 Rate=3002.12 GlobalRate=2606.45 Time=11:15:37
| Training Device=xla:1 Epoch=2 Step=80 Loss=1.37978 Rate=2600.16 GlobalRate=2569.15 Time=11:15:39
| Training Device=xla:1 Epoch=2 Step=90 Loss=1.46465 Rate=2473.87 GlobalRate=2548.12 Time=11:15:42
Epoch 2 train end 11:15:42 - elapsed: 19.041923761367798

| Test Device=xla:1 Step=0 Epoch=2 Time=11:15:45
| Test Device=xla:1 Step=10 Epoch=2 Time=11:17:44
Epoch 2 test end 11:19:31, Accuracy=50.25      <<<<<<<<  11:19:31 <--- 11:15:42

Same code executed with xmp.spawn on 8 devices works as expected. Evaluation time is less than second.

https://colab.research.google.com/drive/16ERnqxd4_3wXCBN64fmrxDjQ910NoKhZ?usp=sharing

Thanks

[zcain117]

I don't it's normal for the eval to be so slow compared to the train. I'm looking into this now.

If you have a chance, could you see if the problem still happens when you use an older image?

Date nightly from 20200607 or earlier would be a good test.

[dlibenzi]

I am able to repro.
Something bad is happening on the service side:

Metric: ExecuteTime
  TotalSamples: 326
  Accumulator: 07m39s959ms809.543us
  ValueRate: 567ms126.112us / second
  Rate: 0.463414 / second
  Percentiles: 1%=018ms818.023us; 5%=018ms018.115us; 10%=018ms194.676us; 20%=018ms390.280us; 50%=019ms850.736us; 80%=027ms017.621us; 90%=12s887ms409.969us; 95%=12s892ms057.878us; 99%=12s911ms203.690us

Metric: XrtExecute
  TotalSamples: 326
  Accumulator: 07m37s381ms650.074us
  Mean: 01s219ms959.049us
  StdDev: 04s583ms555.374us
  Rate: 0.463372 / second
  Percentiles: 25%=016ms338.370us; 50%=017ms622.506us; 80%=017ms156.066us; 90%=12s884ms193.336us; 95%=12s890ms490.679us; 99%=12s907ms904.407us

Execute time takes many seconds sometimes, though the XRT metrics signaling swap/compaction do not trigger.
Something is wrong on the TF service side.

[vfdev-5]

@zcain117 I tried previously 1.5 and even on training there was something weird. I can retry with 20200607 as you said. I do not if this changes something if

Something is wrong on the TF service side.

[zcain117]

@vfdev-5 you are right that updating the VM alone will not fix if the problem is TF service side. But as far as I can tell, the pytorch-xla-env-setup.py script at the top of your colab will update both the VM and the TPU to the same dated version: https://github.com/pytorch/xla/blob/master/contrib/scripts/env-setup.py#L101

[vfdev-5]

@zcain117 good news about VERSION = "20200607", seems like evaluation part is again quick:

Epoch 1 train end 15:18:22 - elapsed: 24.438180923461914
| Test Device=xla:1 Step=0 Epoch=1 Time=15:18:26
| Test Device=xla:1 Step=10 Epoch=1 Time=15:18:28
Epoch 1 test end 15:18:29, Accuracy=41.09
...
Epoch 2 train end 15:18:46 - elapsed: 16.995563507080078
| Test Device=xla:1 Step=0 Epoch=2 Time=15:18:48
| Test Device=xla:1 Step=10 Epoch=2 Time=15:18:50
Epoch 2 test end 15:18:51, Accuracy=50.25

[dlibenzi]

Look like a regression TF service side.
We will look into it.

[zcain117]

I think we've narrowed down the issue to the maxpool operation changing on the Pytorch side. We're working on a fix but I think the only option is to avoid maxpool (it happens here in the torchvision resnet) or continue using the June 7 build for now

[zcain117]

Ailing submitted the maxpool fix earlier this week and our torchvision models have gone back to their normal speed.

It should be safe to use nightly again - let us know if you see any other issues

[vfdev-5]

@zcain117 thanks for the update, I can retry my code and give you the feedback.

[zcain117]

I'll mark this as fixed for now but do let me know if you see a big speed difference between nightly and VERSION = "20200607"