How to Profile TPU Programs | How To Scale Your Model

[jacobaustin123]

TPU profiling, discussed!

[rengoog]

Nice :)

[jtvhk]

The linked Perfetto trace unfortunately fails to restore because of a version mismatch

[jacobaustin123]

Secretly it looks like if you ignore the error, it does actually load properly (I think). But I'll make an updated version which will go up tonight

[ghubnerr]

Props to the team for putting this up! Will probably be coming back here a lot :)

[batterseapower]

I think the description of question 2 might be slightly out of date. If I just run the example notebook on TPU v2 I get about 89ms/layer rather than 184ms/layer as the book says. If I'm careful to do an apples-to-apples comparison the code in the "answer" does seem to do a bit better - it achieves 82ms/layer.

[jacobaustin123]

Sorry it's taken me a while to get to this. Yes it looks like the XLA compiler unfortunately got better. I agree it's still a bit better but not nearly as much as before. I'll add a caveat to the problem that the speed-up isn't that significant anymore and try to add a new problem at some point. Thanks for calling this out.

[abhishm]

The colab example is broken. Could it be possible to make it run on the default TPU available (v2-8).

[blais]

Here we’ve zoomed into the FFW block. You’ll see the up-projection Op is a fusion (matmul) with inputs bf16[8, 1024, 8192] and

I think you mean [32, 1024, 8192] here.

[blais]

Actually never mind. I think that's right, that's the 4-way DP slice.