Overview

Welcome to the 🤗 Optimum-TPU tutorials! Here you’ll find comprehensive examples and guides to help you leverage Google Cloud TPUs effectively with Hugging Face models.

Available Examples

Text Generation

Learn how to perform efficient inference for text generation tasks:

• Basic Generation Script (examples/text-generation/generation.py)
  • Demonstrates text generation using models like Gemma and Mistral
  • Features greedy sampling implementation
  • Shows how to use static caching for improved performance
  • Includes performance measurement and timing analysis
  • Supports custom model loading and configuration

Language Model Fine-tuning

Explore how to fine-tune language models on TPU infrastructure:

1. Interactive Gemma Tutorial (examples/language-modeling/gemma_tuning.ipynb)

   • Complete notebook showing Gemma fine-tuning process
   • Covers environment setup and TPU configuration
   • Demonstrates FSDPv2 integration for efficient model sharding
   • Includes dataset preparation and PEFT/LoRA implementation
   • Provides step-by-step training workflow

2. LLaMA Fine-tuning Guide (examples/language-modeling/llama_tuning.ipynb)

   • Detailed guide for fine-tuning LLaMA-2 and LLaMA-3 models
   • Explains SPMD and FSDP concepts
   • Shows how to implement efficient data parallel training
   • Includes practical code examples and prerequisites

Additional Resources

• Visit the Optimum-TPU GitHub repository for more details
• Explore the Google Cloud TPU documentation for deeper understanding of TPU architecture

For the latest updates and to contribute to these examples, visit our GitHub repository.