RecurrentGemma provides open-weights language models based on Google DeepMind's Griffin architecture, designed for efficient long-sequence generation through a hybrid attention-recurrence mechanism. It targets researchers and developers needing high-performance LLMs for tasks involving extended text, offering optimized Flax and reference PyTorch implementations.

How It Works

The Griffin architecture replaces global attention with a combination of local attention and linear recurrences. This approach significantly speeds up inference for long sequences by reducing the computational complexity associated with traditional self-attention mechanisms, making it more efficient for generating lengthy outputs.

Quick Start & Requirements

• Installation: Use Poetry (poetry install -E full) or pip (pip install .[full]). Library-specific installs are available (-E jax, -E torch, -E test).
• Model Weights: Download from Kaggle (http://kaggle.com/models/google/recurrentgemma).
• Running Examples: python examples/sampling_jax.py --path_checkpoint=/path/to/weights --path_tokenizer=/path/to/tokenizer.model
• Colab Notebooks: Available for sampling and fine-tuning (requires Kaggle account and license acceptance).
• Hardware: Supports CPU, GPU, and TPU. Flax implementation is optimized for TPUs with Pallas kernels. Sampling is supported on T4, P100, V100, A100, TPUv2, and TPUv3+. Fine-tuning is supported on T4, P100, V100, A100, and TPUv3+.

Highlighted Details

• Novel Griffin architecture with local attention and linear recurrences for fast long-sequence generation.
• Optimized Flax implementation with low-level Pallas kernels for TPU performance.
• Reference PyTorch implementation provided.
• Includes tutorials for sampling and fine-tuning via Colab notebooks.

Maintenance & Community

• Open to bug reports and issues; details for PRs are in CONTRIBUTING.md.

Licensing & Compatibility

• Licensed under the Apache License, Version 2.0. Permissive for commercial use and closed-source linking.

Limitations & Caveats

• Fine-tuning is not supported on TPUv2.
• Requires accepting Gemma license terms and conditions from Kaggle to use Colab notebooks.