👋 Welcome to Kithara!#

Note

This project is under active development.

Introduction#

What is Kithara?#

Kithara is a lightweight library offering building blocks and recipes for tuning popular open source LLMs like Llama 3 and Gemma 2 on Google TPUs.

Kithara makes post-training easier to do on TPUs. Post-training algorithms include continued pretraining, fine-tuning, PEFT, RLHF, and quantization. These algorithms adapt a pre-trained language model to a specific task or domain by training it on a custom dataset. Using TPUs provides significant advantages in terms of performance, cost-effectiveness, and scalability, enabling faster training times and the ability to work with larger models and datasets.

Choose Kithara for the following benefits:

→ Hugging Face checkpoints accepted and generated by default

→ Simple pip install

→ User-friendly documentation

→ Library-based UX

→ Wide variety of models through Keras Hub

→ Performant models through Maxtext

→ Custom trainer code in Keras or Jax

→ Ray data loading

→ Version Control

→ Reuse code on GPUs

→ Wide variety of tuning algorithms: SFT, LoRA, Instruction tuning.

→ Coming soon: RL (DPO, PPO, ORPO, GRPO), Distillation, Quantization

→ Coming soon: Ray Tune for automatic hyperparameter selection

Target Audience#

This guide is intended for developers and researchers familiar with machine learning and large language models who want to leverage the power of TPUs for fine-tuning and other post-training workloads. It is structured to guide you through the process, from setup to deployment, with practical examples and explanations of key concepts.

Key Features#

This product supports a variety of tuning algorithms (SFT, DPO, LoRA, Continued pre-training), data formats (JSONL, Parquet, CSV, Text, HuggingFace Dataset, and more), and models (including Gemma 2 2B, 9B, 27B; Llama 3.1 8B, 70B). It also includes performance optimizations (FlashAttention, Scanning, Rematerialization) and parallelism options (FSDP, FSDP+DDP).

Design Principles#

Easy to use: → Frictionless onboarding
Hardware Flexibility: → Use the same code on GPU and TPU
Developer Centric: → Designed for customization instead of black-box recipes