Google’s Push to Improve PyTorch Performance on Its AI Chips
Google, part of Alphabet Inc., is working on a new initiative to enhance the performance of PyTorch, the world’s most widely-used AI software framework, on its artificial intelligence (AI) chips. This move aims to weaken Nvidia’s dominance in the AI computing market, according to sources familiar with the matter.
Google’s Aggresive Plan to Make TPU a Viable Alternative to Nvidia’s GPUs
This initiative is part of Google’s aggressive strategy to make its Tensor Processing Units (TPUs) a viable alternative to Nvidia’s Graphics Processing Units (GPUs), which currently lead the market. Sales of TPUs have become a crucial revenue growth driver in Google’s cloud business, as the company strives to demonstrate to investors that its AI investments are yielding returns.
Overcoming Hardware Adoption Barriers
However, hardware alone is not enough to stimulate adoption. The new initiative, known internally as “TorchTPU,” aims to eliminate a key barrier that has slowed TPU adoption by making them fully compatible and easy to develop for customers who have already built their technology infrastructure using PyTorch software, according to the sources.
Potential Open-Sourcing of Software Components
Google is also considering making parts of the software open-source to accelerate adoption among clients, some sources said.
Google’s Strategic Focus on TorchTPU
Compared to previous attempts to support PyTorch on TPUs, Google has devoted more organizational attention, resources, and strategic importance to TorchTPU as demand for chip adoption grows while companies see the software stack as a bottleneck, sources said.
PyTorch’s Relevance and Meta’s Involvement
PyTorch, an open-source project heavily supported by Meta Platforms, is one of the most popular tools used by developers creating AI models. In Silicon Valley, few developers write every line of code that will run on Nvidia’s chips, Advanced Micro Devices’ chips, or Google’s chips. Instead, they rely on tools like PyTorch, a collection of pre-written libraries and frameworks that automate many common tasks in AI software development.
Launched in 2016, PyTorch’s story has been closely tied to Nvidia’s CUDA development. Some Wall Street analysts consider CUDA Nvidia’s strongest shield against competitors.
Nvidia’s Software Optimization vs. Google’s Jax and XLA
Nvidia engineers have long ensured that software developed with PyTorch runs as fast and efficiently as possible on their chips. Google, on the other hand, has long made its internal software developers use a different code framework called Jax, and its TPU chips use a tool called XLA to execute that code efficiently.
Most of Google’s AI software stack and performance optimization have been built around Jax, widening the gap between how Google uses its chips and how clients want to use them.
Google’s TPU Offerings for Clients
Alphabet has long reserved most of its TPUs exclusively for internal use. This changed in 2022 when Google Cloud’s computing unit successfully lobbied to oversee the group selling TPUs, significantly increasing TPU allocation in Google Cloud.
As AI interest from clients has grown, Google has tried to capitalize by increasing TPU production and sales to external customers. However, the disconnect between PyTorch, used by most AI developers worldwide, and Jax, for which Google’s chips are currently optimized, means most developers find it difficult to easily adopt Google’s chips without significant additional engineering work.
Collaborative Efforts with Meta
To expedite development, Google is collaborating closely with Meta, PyTorch’s creator and manager, according to sources. The two tech giants have discussed agreements allowing Meta access to more TPUs, a move first reported by The Information.
Early offers to Meta were structured as managed services, where clients like Meta would install Google-designed chips for running their software and models, and Google would provide operational support.
Meta has a strategic interest in working on software that makes operating TPUs easier, an attempt to reduce inference costs and diversify its AI infrastructure away from Nvidia’s GPUs to gain more bargaining power, sources said.
Meta declined to comment.
This year, Google began selling TPUs directly to client data centers instead of limiting access through its own cloud. Amin Vahdat, a Google veteran, was appointed AI infrastructure chief this month, reporting directly to CEO Sundar Pichai.
Google needs this infrastructure both to run its own AI products, including the Gemini chatbot and AI-powered search, and to supply clients of Google Cloud, which sells TPU access to companies like Anthropic.
Key Questions and Answers
- What is Google’s new initiative about? Google is working on enhancing PyTorch performance on its AI chips to challenge Nvidia’s dominance in the AI computing market.
- Why is this important? Nvidia’s strong position in the AI computing market is due not only to its hardware but also to its CUDA software ecosystem, deeply integrated with PyTorch and becoming the default method for training and running large AI models.
- What challenges does Google face in promoting its TPU chips? The disconnect between PyTorch, widely used by AI developers, and Jax, Google’s preferred internal learning framework, creates a significant barrier to TPU adoption.
- How is Google addressing these challenges? Google’s TorchTPU initiative aims to make TPUs fully compatible with PyTorch, and the company is considering open-sourcing software components to accelerate adoption.
- Why is Meta involved? Meta has a strategic interest in working with Google on TPU-related software to reduce inference costs and diversify its AI infrastructure away from Nvidia’s GPUs.
