Elamur Research⌕
☰
home / papers / 2023 / TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

Paper · 2023 · Norman P. Jouppi; George Kurian; Sheng Li; Peter Ma; Rahul Nagarajan; Lifeng Nai; Nishant Patil; Suvinay Subramanian; Andy Swing; Brian Towles; Cliff Young; Xiang Zhou; Zongwei Zhou; and David Patterson

Summary

Google's TPU v4 paper tackles a hard problem: building a machine-learning supercomputer with thousands of chips that stays reliable and efficient while workloads keep changing. Over time, demand moved from MLPs and RNNs toward Transformers and LLMs, alongside embedding-heavy recommender models (DLRMs). The paper presents three novelties of a system in production since 2020.

  1. Optical interconnect. Electrically cabled 4×4×4 cubes of 64 chips are linked through 48 Palomar optical circuit switches (OCSes) built on MEMS mirrors. Together they form a 4096-chip 3D torus. Its topology can be reconfigured, including twisted tori, and it can bypass failed CPU hosts.
  2. SparseCore. This dataflow unit for embeddings speeds up DLRMs 5x–7x while using roughly 5% of die area and power.
  3. Their interaction. Embeddings create all-to-all traffic, which benefits from flexible topologies with higher bisection bandwidth.

Against TPU v3, TPU v4 is 2.1x faster with 2.7x better performance/Watt. The optics cost under 5% of the system and draw under 3% of power. PaLM trained at 57.8% of peak FLOPS. The lesson is that a flexible interconnect can be a cheap lever for scale, uptime and efficiency.

Key points

  • System size. A full machine has 4096 chips in 64 racks, with 4 TPUs per host (about 1K hosts).
    • Each 64-chip block exposes 96 optical links, which attach to 48 OCSes.
    • The Palomar OCS is 136×136: 128 working ports plus 8 spares for testing and repair.
  • Chip design. Each chip holds two TensorCores sharing a 128 MiB CMEM. Each TensorCore contains:
    • four 128×128 MXUs,
    • a 128-lane VPU,
    • 16 MiB of VMEM.
  • Workload shift. By October 2022, Transformers took 57% of TPU v4 training (LLM 31%, BERT 26%), while RNNs dropped to 2%.
  • Comparisons with similar-sized systems.
    • Versus Graphcore IPU Bow: about 4.3x–4.5x faster.
    • Versus Nvidia A100: 1.2x–1.7x faster and 1.3x–1.9x lower power.
    • In Google Cloud: about 2–6x less energy and about 20x less CO2e than typical on-premise DSAs.
  • Availability and deployment. Without OCSes, hosts would need 99.9% availability to give reasonable slice goodput. With OCSes, each rack could enter service on its own instead of waiting for the full system.

Interesting ideas

  • I1. Small blocks plus a reprogrammable switch layer. Building from small, independently testable blocks joined by a reprogrammable "patchboard" lets capacity come online gradually and lets failed parts be skipped. The same pattern suits compute clusters, data pipelines, or any system assembled from unreliable components.
  • I2. Goodput with spares as an evaluation metric. Measuring effective throughput while explicitly reserving spare capacity, rather than counting raw peak, gives an honest picture of what users actually get. This is worth reusing when planning shared GPU budgets.
  • I3. Logging the real workload mix over time. The authors tracked which model types actually consumed their hardware over several years and used that record to justify design choices. An NLP group could keep a similar record to decide where optimization effort pays off.
  • I4. Matching topology to communication pattern. Users pick the slice shape per job, including twisted tori, to fit how that job communicates (all-reduce versus all-to-all). The general lesson is to let the job's traffic pattern drive the infrastructure configuration.
  • I5. Layered carbon accounting. The CO2e estimate combines chip efficiency, data-center efficiency and the grid's carbon-free share. This layered approach gives a reusable template for reporting the footprint of model training.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢GoogleMountain View · United States430 works · 424,363 citations🎓University of California, BerkeleyBerkeley · United States175 works · 227,440 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksComputer Vision and Pattern Recognition · 2,869 worksHardware and Architecture · 81 worksComputer Networks and Communications · 64 works

🧭 Topics 1Research subjects and languages the work is about.GPU & parallel computing · 19 works

💡 Concepts 1Key ideas and methods explained on our concept pages.Machine learning · 51 works

Cite

Cite as Norman P. Jouppi; George Kurian; Sheng Li; Peter Ma; Rahul Nagarajan; Lifeng Nai; Nishant Patil; Suvinay Subramanian; Andy Swing; Brian Towles; Cliff Young; Xiang Zhou; Zongwei Zhou; and David Patterson. 2023. TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings. arXiv:2304.01433.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2304.01433 · arXiv 2304.01433

@misc{zhou2023tpu,
  title = {{TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings}},
  author = {Norman P. Jouppi; George Kurian; Sheng Li; Peter Ma; Rahul Nagarajan; Lifeng Nai; Nishant Patil; Suvinay Subramanian; Andy Swing; Brian Towles; Cliff Young; Xiang Zhou; Zongwei Zhou; and David Patterson},
  year = {2023},
  doi = {10.48550/arxiv.2304.01433},
  eprint = {2304.01433},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2304.01433}
}

References

62 references, 6 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Anil, R., Gadanho, S., Huang, D., Jacob, N., Li, Z., Lin, D., Phillips, T., Pop, C., Regan, K., Shamir, G.I. and Shivanna, R., 2022. On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models. 16th ACM Conference on Recommender Systems. Barroso, L.A., and Hölzle, U., 2009, F… find
  2. Barroso, L.A., Hölzle, U. and Ranganathan, P., 2018. The datacenter as a computer: Designing warehouse-scale machines, Third Edition. Synthesis Lectures on Computer Architecture, 13(3), pp.i-189. find
  3. Bloomberg October 26, 2015. Google turning its lucrative web search over to AI machines. find
  4. Bouknight, W.J., Denenberg, S.A., McIntyre, D.E., Randall, J.M., Sameh, A.H. and Slotnick, D.L., 1972. The ILLIAC IV system. Proceedings of the IEEE, 60(4), pp. 369-388. find
  5. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A. and Agarwal, S., 2020. Language models are few-shot learners. Advances in Neural Information Processing Systems 33 (NeurIPS 2020).
  6. Camara, J.M., Moreto, M., Vallejo, E., Beivide, R., Miguel-Alonso, J., Martínez, C. and Navaridas, J., 2010. Twisted torus topologies for enhanced interconnection networks. IEEE Transactions on Parallel and Distributed Systems, 21(12), pp. 1765-1778. find
  7. Camarero, C., Martinez, C. and Beivide, R., 2014. Lattice graphs for high-scale interconnection topologies. IEEE Transactions on Parallel and Distributed Systems, 26(9), pp. 2506-2519. find
  8. Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H.W., Sutton, C., Gehrmann, S. and Schuh, P., 2022. PaLM: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311.
  9. Colamco, 2022. Mellanox QM8790 - Quantum HDR Switch, https://www.colamco.com/product/mellanox-infiniband-switch-mq m8790-hs2f-1503410?utm_source=froogle&utm_medium=referral. ↗
  10. Covington, P., Adams, J. and Sargin, E., 2016. Deep neural networks for Youtube recommendations. In Proceedings of the 10th ACM conference on recommender systems, pp. 191-198. find
  11. Dally, W.J. and Towles, B.P., 2004. Principles and practices of interconnection networks. Elsevier. find
  12. Deepmind, Nov 18, 2019, Advanced machine learning helps Play Store users discover personalised apps. find
  13. Devlin, J., Chang, M.W., Lee, K. and Toutanova, K., 2018. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.
  14. Dungworth, M., Harrell, J., Levine, M., Nelson, S., Oberlin, S. and Reinhardt, S.P., 2011. CRAY T3E. find
  15. Fan, X., Weber, W.D. and Barroso, L.A., 2007. Power provisioning for a warehouse-sized computer. In 2007 ACM/IEEE 34th International Symposium on Computer Architecture (ISCA), pp.13-23. find
  16. Farrington, N. ., G. Porter, S. Radhakrishnan, H. H. Bazzaz, V. Subramanya, Y. Fainman, G. Papen, and A. Vahdat. Helios: A hybrid electrical/optical switch architecture for modular data centers. In Proceedings of ACM SIGCOMM, Aug. 2010. find
  17. Ghaemmaghami, B., Ozdal, M., Komuravelli, R., Korchev, D., Mudigere, D., Nair, K. and Naumov, M., 2022. Learning to Collide: Recommendation System Model Compression with Learned Hash Functions. arXiv preprint arXiv:2203.15837. — [arXiv:2203.15837](https://arxiv.org/abs/2203.15837) find
  18. Google, Tracking our carbon free energy progress, https://sustainability.google/progress/energy/. ↗
  19. Google, Google Data Center Efficiency, https://www.google.com/about/datacenters/efficiency/. ↗
  20. Graphcore, 2022. IPU-POD64 Reference Design Datasheet, docs.graphcore.ai/projects/ ipu-pod64-datasheet/en/latest/ find
  21. Horowitz, M., 2014. Computing's energy problem (and what we can do about it). In 2014 IEEE International Solid-State Circuits Conference Digest of Technical Papers (ISSCC), pp. 10-14. find
  22. Insight, 2022. Mellanox Quantum QM8790 - switch - 40 ports, https://www.insight.com/en_US/shop/ product/MQM8790-HS2F/Mellanox/MQM8790-HS2F/MellanoxQua ntumQM8790-swit/. ↗
  23. International Energy Agency, Global Energy & CO2 Status Report 2019, Report Extract Emissions. find
  24. Jouppi, N.P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., Boyle, R., Cantin, P., Chao, C., Clark, C., Coriell, J., Daley, M., Dau, M., Dean, J., Gelb, B., Ghaemmaghami, T., Gottipati, R., Gulland, W., Hagmann, R., Ho, C.R., Hogberg, D.… find
  25. Jouppi, N.P., Yoon, D.H., Kurian, G., Li, S., Patil, N., Laudon, J., Young, C. and Patterson, D., 2020. A domain-specific supercomputer for training deep neural networks. Communications of the ACM, 63(7), pp. 67-78. find
  26. Jouppi, N.P., Yoon, D.H., Ashcraft, M., Gottscho, M., Jablin, T.B., Kurian, G., Laudon, J., Li, S., Ma, P., Ma, X., Norrie, T., Patil, N., Prasad, S., Young, C., Zhou, Z., and Patterson, D. 2021. Ten lessons from three generations shaped Google’s TPUv4i. In 2021 ACM/IEEE 48th Annual International Sy… find
  27. Kamil, S., Pinar, A., Gunter, D., Lijewski, M., Oliker, L. and Shalf, J., 2007, May. Reconfigurable hybrid interconnection for static and dynamic scientific applications. In Proceedings of the 4th Int'l conference on computing frontiers, pp. 183-194. find
  28. Kamil, S., Oliker, L., Pinar, A. and Shalf, J., 2009. Communication requirements and interconnect optimization for high-end scientific applications. IEEE Transactions on Parallel and Distributed Systems, 21(2), pp. 188-202. find
  29. Karandikar, S., Mao, H., Kim, D., Biancolin, D., Amid, A., Lee, D., Pemberton, N., Amaro, E., Schmidt, C., Chopra, A. and Huang, Q., 2018, June. FireSim: FPGA-accelerated cycle-exact scale-out system simulation in the public cloud. In 2018 ACM/IEEE 45th Annual International Symposium on Computer Arc… find
  30. Khani, M., Ghobadi, M., Alizadeh, M., Zhu, Z., Glick, M., Bergman, K., Vahdat, A., Klenk, B. and Ebrahimi, E., 2021, August. SiP-ML: high-bandwidth optical network interconnects for machine learning training. In Proceedings of the 2021 ACM SIGCOMM 2021 Conference (pp. 657-675). find
  31. Li, S., Tan, M., Pang, R., Li, A., Cheng, L., Le, Q.V. and Jouppi, N.P., 2021. Searching for Fast Model Families on Datacenter Accelerators. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8085-8095. find
  32. Li, S., Andersen, G., Chen, T., Cheng, L., Grady, J., Huang, D., Le, Q., Li, A., Li, X., Li, Y., Liang, C., Lu, Y., Ni, Y., Pang, F., Ranganathan , P., Tan, M., Wicke, M., Wu, G., Zhu, S., and Jouppi, N., 2023. Hyperscale Hardware Optimized Neural Architecture Search, In Proceedings of the 28th Inte… find
  33. Masanet, E., Shehabi, A., Lei, N., Smith, S. and Koomey, J., 2020. Recalibrating global data center energy-use estimates. Science, 367(6481), pp.984-986. find
  34. Minkenberg, C., Rodriguez, G., Prisacari, B., Schares, L., Heidelberger, P., Chen, D. and Stunkel, C., 2016, March. Performance benefits of optical circuit switches for large-scale dragonfly networks. In 2016 Optical Fiber Communications Conference and Exhibition (OFC) (pp. 1-3). IEEE. find
  35. MLCommons, V 2.0 Results, June 29, 2022, https://mlcommons.org/en/training-normal-20/. ↗
  36. Mudigere, D., Hao, Y., Huang, J., Jia, Z., Tulloch, A., Sridharan, S., Liu, X., Ozdal, M., Nie, J., Park, J. and Luo, L., 2022, June. Software-hardware co-design for fast and scalable training of deep learning recommendation models. In Proceedings of the 49th Annual International Symposium on Comput… find
  37. Nayak, P., 2021. MUM: A new AI milestone for understanding information. Google, May, 18. find
  38. Norrie, T., Patil, N., Yoon, D.H., Kurian, G., Li, S., Laudon, J., Young, C., Jouppi, N. and Patterson, D., 2021. The design process for Google's training chips: TPUv2 and TPUv3. IEEE Micro, 41(2), pp. 56-63. find
  39. Nvidia, 2020. Nvidia A100 Tensor Core GPU Architecture. find
  40. Nvidia, 2021. Nvidia DGX SuperPOD: Scalable Infrastructure for AI Leadership Reference Architecture. find
  41. Patterson, D., Gonzalez, J., Le, Q., Liang, C., Munguia, L.M., Rothchild, D., So, D., Texier, M. and Dean, J., 2022. The carbon footprint of machine learning training will plateau then shrink, IEEE Computer, 55(7), pp. 18-28. find
  42. Poutievski, L., Mashayekhi, O., Ong, J., Singh, A., Tariq, M., Wang, R., Zhang, J., Beauregard, V., Conner, P., Gribble, S. and Kapoor, R., 2022. Jupiter evolving: transforming Google's datacenter network via optical circuit switches and software-defined networking. In Proceedings of the ACM SIGCOMM… find
  43. Raghavendra, R., Ranganathan, P., Talwar, V., Wang, Z. and Zhu, X., 2008, March. No "power" struggles: coordinated multi-level power management for the data center. In Proceedings of the 13th International Conference on architectural support for programming languages and operating systems (ASPLOS),… find
  44. Schwartz, R., Dodge, J., Smith, N.A. and Etzioni, O., 2020. Green AI. Communications of the ACM, 63(12), pp. 54-63.
  45. Scott, S.L., 1996. Synchronization and communication in the T3E multiprocessor. In Proceedings of the 7th International conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), pp. 26-36. find
  46. Sequin, C.H., 1981. Doubly twisted torus networks for VLSI processor arrays. In Proceedings of the 8th International annual Symposium on Computer Architecture (ISCA), pp. 471-480. find
  47. Sethi, G., Bhattacharya, P., Choudhary, D., Wu, C.-J., and Kozyrakis, C., 2022. FlexShard: Flexible Sharding for Industry-Scale Sequence Recommendation Models, arxiv preprint arXiv:2301.02959. — [arXiv:2301.02959](https://arxiv.org/abs/2301.02959) find
  48. Shalf, J., Kamil, S., Oliker, L. and Skinner, D., 2005, November. Analyzing ultra-scale application communication requirements for a reconfigurable hybrid interconnect. In SC'05: Proceedings of the 2005 ACM/IEEE Conference on Supercomputing (pp. 17-17). find
  49. Shaw, D.E., Deneroff, M.M., Dror, R.O., Kuskin, J.S., Larson, R.H., Salmon, J.K., Young, C., Batson, B., Bowers, K.J., Chao, J.C. and Eastwood, M.P., 2008. Anton, a special-purpose machine for molecular dynamics simulation. Communications of the ACM, 51(7), pp. 91-99. find
  50. Singla, A., Singh, K. Ramachandran, L. Xu, and Y. Zhang. Proteus: A topology malleable data center network. In Proceedings of ACM Hotnets, Oct. 2010. find
  51. Taylor, P., 2022. Data center average annual power usage effectiveness (PUE) worldwide 2007-2022, Nov 22. find
  52. Teh, M. Y., S. Zhao, P. Cao, and K. Bergman. Couder: Robust topology engineering for optical circuit switched data center networks. arXiv preprint arXiv:2010.00090, Sept. 2020. — [arXiv:2010.00090](https://arxiv.org/abs/2010.00090) find
  53. Thoppilan, R., De Freitas, D., Hall, J., Shazeer, N., Kulshreshtha, A., Cheng, H.T., Jin, A., Bos, T., Baker, L., Du, Y. and Li, Y., 2022. LaMDA: Language models for dialog applications. arXiv preprint arXiv:2201.08239.
  54. Urata, R., Liu, H., Yasumura, K., Mao, E., Berger, J., Zhou, X., Lam, C., Bannon, R., Hutchinson, D., Nelson, D. and Poutievski, L., 2022. Mission Apollo: Landing Optical Circuit Switching at Datacenter Scale. arXiv preprint arXiv:2208.10041. — [arXiv:2208.10041](https://arxiv.org/abs/2208.10041) find
  55. US Energy Information Agency, 2022. Oklahoma State Profile and Energy Estimates, https://www.eia.gov/state/ ?sid=OK\#:~:text=In%202021%2C%20wind%20supplied%2041,ele ctricity%20net%20generation%20from%20wind. ↗
  56. Vahdat, A., H. Liu, X. Zhao, and C. Johnson, 2011. The emerging optical data center. In Proceedings of the Optical Fiber Communications Conference. find
  57. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł. and Polosukhin, I., 2017. Attention is all you need. Advances in neural information processing systems, 30 (NeurIPS 2017).
  58. Wang, S., Wei, J., Sabne, A., David, A., Llbeyi, B., Hechtman, B., Chen, D., Murthy, K. S., Maggioni, M., Zhang, Q., Kumar, S., Guo, T., Xu, Y., and Zhou, Z., 2023. Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models, In Proceedings of the 28th Internatio… — [arXiv:2202.00433](https://arxiv.org/abs/2202.00433) find
  59. Williams, S., Waterman, A. and Patterson, D., 2009. Roofline: an insightful visual performance model for multicore architectures. Communications of the ACM, 52(4), pp.65-76. find
  60. Wu, Q., Deng, Q., Ganesh, L., Hsu, C.H., Jin, Y., Kumar, S., Li, B., Meza, J. and Song, Y.J., 2016. Dynamo: Facebook’s Data Center-Wide Power Management System. 2016 ACM/ IEEE 43rd Annual International Symposium on Computer Architecture (ISCA), Seoul, (2016), pp. 469–480. find
  61. Xu, Y., Lee, H., Chen, D., Hechtman, B., Huang, Y., Joshi, R., Krikun, M., Lepikhin, D., Ly, A., Maggioni, M. and Pang, R., 2021. GSPMD: general and scalable parallelization for ML computation graphs. arXiv preprint arXiv:2105.04663. find
  62. Zhao, M., Agarwal, N., Basant, A., Gedik, B., Pan, S., Ozdal, M., Komuravelli, R., Pan, J., Bao, T., Lu, H. and Narayanan, S., 2022, June. Understanding data storage and ingestion for large-scale deep recommendation model training: Industrial product. In Proceedings of the 49th Annual International… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/jouppi-2023-tpu-v4-optically-reconfigurable-supercomputer/