Add a new resnet50+convergence+spmd test

[yeounoh]

Description

Add a new test for Resnet50+SPMD for convergence. Use batch sharding.

[jonb377]

@yeounoh Do you have a oneshot run for this?

[jonb377]

I tried a local run, it's very slow:

Epoch 1 train begin 22:36:40
| Training Device=xla:0/0 Epoch=1 Step=0 Loss=7.11059 Rate=9.68 GlobalRate=9.68 Time=22:38:26
| Training Device=xla:0/0 Epoch=1 Step=200 Loss=6.48232 Rate=67.85 GlobalRate=101.57 Time=23:10:26

We might need to optimize our resnet training script some more to avoid timeouts.

EDIT: It seems like dataloading is the bottleneck, PjRt convergence tests are also pretty slow: http://shortn/_whmMQ4flXg

[jonb377]

LGTM, thanks Yeounoh!

[jonb377]

Let's make the timeout 14 hours to match the PjRt tests. I've kicked off a oneshot here, let's merge if it passes: http://shortn/_b8rAgZ0mkY

cloudtop [~/d/x/1/ml-testing-accelerators] % git diff
diff --git a/tests/pytorch/nightly/resnet50-mp.libsonnet b/tests/pytorch/nightly/resnet50-mp.libsonnet
index b24e5bc5..ba86abab 100644
--- a/tests/pytorch/nightly/resnet50-mp.libsonnet
+++ b/tests/pytorch/nightly/resnet50-mp.libsonnet
@@ -209,5 +209,6 @@ local tpus = import 'templates/tpus.libsonnet';
     // SPMD
     resnet50 + functional + v4_8 + timeouts.Hours(2) + spmd(['batch']),
     resnet50 + functional + v4_8 + timeouts.Hours(2) + spmd(['spatial']),
+    resnet50 + convergence + v4_8 + timeouts.Hours(14) + spmd(['batch']),
   ],
 }

[yeounoh]

Thanks, it failed due to RESOURCE_EXHAUSTED after an epoch. Is this the same error you were running into for the convergence test?

[jonb377]

I would guess this is related to the eval loop's dataloader sharding. I also noticed that our rate is still ~1/8 that of the regular PjRt tests - probably data processing for all devices in a single thread.

[alanwaketan]

I tried a local run, it's very slow:

Epoch 1 train begin 22:36:40
| Training Device=xla:0/0 Epoch=1 Step=0 Loss=7.11059 Rate=9.68 GlobalRate=9.68 Time=22:38:26
| Training Device=xla:0/0 Epoch=1 Step=200 Loss=6.48232 Rate=67.85 GlobalRate=101.57 Time=23:10:26

We might need to optimize our resnet training script some more to avoid timeouts.

EDIT: It seems like dataloading is the bottleneck, PjRt convergence tests are also pretty slow: http://shortn/_whmMQ4flXg

If that's the case, can we use bert or some other language models?

[yeounoh]

I tried a local run, it's very slow:

Epoch 1 train begin 22:36:40
| Training Device=xla:0/0 Epoch=1 Step=0 Loss=7.11059 Rate=9.68 GlobalRate=9.68 Time=22:38:26
| Training Device=xla:0/0 Epoch=1 Step=200 Loss=6.48232 Rate=67.85 GlobalRate=101.57 Time=23:10:26

We might need to optimize our resnet training script some more to avoid timeouts.
EDIT: It seems like dataloading is the bottleneck, PjRt convergence tests are also pretty slow: http://shortn/_whmMQ4flXg

If that's the case, can we use bert or some other language models?

Yea, we can just use llama, but it would require us to emit the correct tb metrics. Will follow up with @will-cromar @jonb377

Convergence too slow

[alanwaketan]

Curious on are we going to use resnet or not for the convergence test?

[jonb377]

Curious on are we going to use resnet or not for the convergence test?

I chatted with @will-cromar - there's still hope for resnet. We can increase the number of worker threads to better match the MP PjRt performance without reinventing the SPMD dataloader.

I'll test this out later today (run with --num_workers 64, it's still pretty slow)