However, there are some limitations and uncertainties. Jetstream is currently limited to supporting TPUs with GPU compatibility supposedly coming in the future. It's also unclear how Google arrived at the claim of 3x improvement in performance per dollar. Optimum TPU currently only works with the Gemma 7B model and doesn't yet support training generative models on TPUs, only running them, although Google promises improvements in the future.
Key takeaways:
- Google has debuted a number of open-source tools aimed at supporting generative AI projects and infrastructure, including MaxDiffusion, a collection of reference implementations of various diffusion models, and Jetstream, an engine to run generative AI models.
- Jetstream is currently limited to supporting TPUs with GPU compatibility supposedly coming in the future, and offers up to 3x higher “performance per dollar” for models like Google’s own Gemma 7B and Meta’s Llama 2.
- Google has also made new additions to MaxText, its collection of text-generating AI models targeting TPUs and Nvidia GPUs in the cloud, which now includes Gemma 7B, OpenAI’s GPT-3, Llama 2 and models from AI startup Mistral.
- Google has collaborated with AI startup Hugging Face to create Optimum TPU, a tool designed to reduce the barrier to entry for getting generative AI models onto TPU hardware, although currently it only works with the Gemma 7B model and doesn't yet support training generative models on TPUs.