4 ms·

+1, almost all (if not all) Google training runs on TPU. They don't use NVIDIA GPUs at all.

dekhn 2y ago
at some point some researchers were begging for GPUs... mainly for sparse work. I think that's why sparsecore was added to TPU (https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#sparsecore https://cloud.google.com/tpu/docs/system-architecture-tpu-vm...) in v4. I think at this point with their tech turnaround time they can catch up as competitors add new features and researchers want to use them.
- felarof 2y ago
  dumb question: wdym by sparse work? Is it embedding lookups? (TPUs have had BarnaCore for efficient embedding lookups since TPU v3)
  - dekhn 2y ago
    Mostly embedding, but IIRC DeepMind RL made use of sparsity- basically, huge matrices with only a few non-zero elements. BarnaCore existed and was used, but was tailored mostly for embeddings. BTW, IIRC they were called that because they were added "like a barnacle hanging off the side". The evolution of TPU has been interesting to watch; I came from the HPC and supercomputing space, and seeing Google as mostly-CPU for the longest time, and then finally learning how to build "supercomputers" over a decade+ (gradually adding many features that classical supercomputers have long had), was a very interesting process. Some very expensive mistakes along the way. But now they've paid down almost all the expensive up-front costs and can now ride on the margins, adding new bits and pieces while increasing the clocks and capacities on a cadence.
amelius 2y ago
Do they have the equivalent of CUDA, and what is it called?
- dekhn 2y ago
  Not exactly, although CUDA is a huge topic. See https://gist.github.com/shawwn/0e524d4a7a5d8fb152a86616559cc02a https://gist.github.com/shawwn/0e524d4a7a5d8fb152a86616559cc... for some description of the process. Basically, jax or other program is converted to XLA, https://opensource.googleblog.com/2023/05/pjrt-simplifying-ml-hardware-and-framework-integration.html https://opensource.googleblog.com/2023/05/pjrt-simplifying-m... then lowered to the specific architecture (which coudl be CPU, GPU, or TPU). Last time I looked, it was a horribly complicated stack with many parts changing rapidly, although with the switch to jax, things got cleaned up a bit. My personal favorite bits are the lower levels of jax, xla, and pjrt.

at some point some researchers were begging for GPUs... mainly for sparse work. I think that's why sparsecore was added to TPU (https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#sparsecore https://cloud.google.com/tpu/docs/system-architecture-tpu-vm...) in v4. I think at this point with their tech turnaround time they can catch up as competitors add new features and researchers want to use them.

felarof 2y ago
dumb question: wdym by sparse work? Is it embedding lookups? (TPUs have had BarnaCore for efficient embedding lookups since TPU v3)
- dekhn 2y ago
  Mostly embedding, but IIRC DeepMind RL made use of sparsity- basically, huge matrices with only a few non-zero elements. BarnaCore existed and was used, but was tailored mostly for embeddings. BTW, IIRC they were called that because they were added "like a barnacle hanging off the side". The evolution of TPU has been interesting to watch; I came from the HPC and supercomputing space, and seeing Google as mostly-CPU for the longest time, and then finally learning how to build "supercomputers" over a decade+ (gradually adding many features that classical supercomputers have long had), was a very interesting process. Some very expensive mistakes along the way. But now they've paid down almost all the expensive up-front costs and can now ride on the margins, adding new bits and pieces while increasing the clocks and capacities on a cadence.

amelius 2y ago

Do they have the equivalent of CUDA, and what is it called?

dekhn 2y ago
Not exactly, although CUDA is a huge topic. See https://gist.github.com/shawwn/0e524d4a7a5d8fb152a86616559cc02a https://gist.github.com/shawwn/0e524d4a7a5d8fb152a86616559cc... for some description of the process. Basically, jax or other program is converted to XLA, https://opensource.googleblog.com/2023/05/pjrt-simplifying-ml-hardware-and-framework-integration.html https://opensource.googleblog.com/2023/05/pjrt-simplifying-m... then lowered to the specific architecture (which coudl be CPU, GPU, or TPU). Last time I looked, it was a horribly complicated stack with many parts changing rapidly, although with the switch to jax, things got cleaned up a bit. My personal favorite bits are the lower levels of jax, xla, and pjrt.