Welcome Gemma 2 - Google’s new open LLM
philschmid, osanseviero, pcuenq, lewtun, tomaarsen, reach-vb
Google released Gemma 2, the latest addition to its family of state-of-the-art open LLMs, and we are excited to collaborate with Google to ensure the best integration in the Hugging Face ecosystem. You can find the 4 open-weight models (2 base models & 2 fine-tuned ones) on the Hub. Among the features and integrations being released, we have:

- [Models on the Hub](https://huggingface.co/collections/google/g-667d6600fd5220e7b967f315)
- Hugging Face [Transformers integration](https://github.com/huggingface/transformers/releases/tag/v4.42.0)
- Integration with Google Cloud & Inference Endpoints

## Table of contents

- [What is Gemma 2?](#what-is-gemma-2)
- [Technical advances in Gemma 2](#technical-advances-in-gemma-2)
  - [Sliding window attention](#sliding-window-attention)
  - [Soft-capping and attention implementations](#soft-capping-and-attention-implementations)
  - [Knowledge Distillation](#knowledge-distillation)
  - [Model Merging](#model-merging)
- [Gemma 2 evaluation](#gemma-2-evaluation)
  - [Technical Report results](#technical-report-results)
  - [Open LLM Leaderboard results](#open-llm-leaderboard-results)
- [How to prompt Gemma 2](#how-to-prompt-gemma-2)
- [Demo](#demo)
- [Using Hugging Face Transformers](#using-hugging-facetransformers)
- [Integration with Google Cloud](#integration-with-google-cloud)
- [Integration with Inference Endpoints](#integration-with-inference-endpoints)
- [Fine-tuning with 🤗 TRL](#fine-tuning-with-trl)
- [Additional Resources](#additional-resources)
- [Acknowledgments](#acknowledgments)

## What is Gemma 2?

Gemma 2 is Google's latest iteration of open LLMs. It comes in two sizes, 9 billion and 27 billion parameters with base (pre-trained) and instruction-tuned versions. Gemma is based on Google Deepmind Gemini and has a context length of 8K tokens:

- [gemma-2-9b](https://huggingface.co/google/gemma-2-9b): Base 9B model.
- [gemma-2-9b-it](https://huggingface.co/google/gemma-2-9b-it): Instruction fine-tuned version of the base 9B model.
- [gemma-2-27b](https://huggingface.co/google/gemma-2-27b): Base 27B model.
- [gemma-2-27b-it](https://huggingface.co/google/gemma-2-27b-it): Instruction fine-tuned version of the base 27B model.

The Gemma 2 models were trained on ~2x more data than their first iteration, totaling 13 trillion tokens for the 27B version and 8 trillion tokens for the 9B version of web data (primarily English), code, and math. We don’t know the exact details of the training mix, and we can only guess that bigger and more careful data curation was a big factor in the improved performance.

Gemma 2 comes with the [same license](https://ai.google.dev/gemma/terms) as the first iteration, which is a permissive license that allows redistribution, fine-tuning, commercial use, and derivative works. 

## Technical advances in Gemma 2

Gemma 2 has many similarities with the first iteration. It has a context length of 8192 tokens and uses Rotary Position Embedding (RoPE). There are four main advances in Gemma 2 compared to the original Gemma: 

- [Sliding window attention](#sliding-window-attention): Interleave sliding window and full-quadratic attention for quality generation.
- [Logit soft-capping](#soft-capping-and-attention-implementations): Prevents logits from growing excessively by scaling them to a fixed range, improving training.
- [Knowledge Distillation](#knowledge-distillation): Leverage a larger teacher model to train a smaller model (for the 9B model).
- [Model Merging](#model-merging): Combines two or more LLMs into a single new model

Gemma 2 was trained on [Google Cloud TPU (27B on v5p](https://cloud.google.com/blog/products/ai-machine-learning/introducing-cloud-tpu-v5p-and-ai-hypercomputer?hl=en), [9B on TPU v4)](https://cloud.google.com/tpu/docs/v4) using [JAX](https://jax.readthedocs.io/en/latest/quickstart.html) and [ML Pathways](https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/). Gemma 2 Instruct has been optimized for dialogue applications and trained on a mix of synthetic and human-generated prompt-response pairs using Supervised Fine-Tuning (SFT), Distillation from a larger model, Reinforcement Learning from Human Feedback (RLHF) using a reward model oriented more towards conversational capabilities, and model merging using WARP to improve overall performance.

Similar to the pre-training mix, no details about the fine-tuning datasets or the hyperparameters associated with SFT and [RLHF](https://huggingface.co/blog/rlhf) have been shared.

### Sliding window attention

[Sliding window attention](https://huggingface.co/papers/2004.05150) is a method to reduce the memory and time requirements of the attention computations in transformer models and has been used in models such as [Mistral](https://huggingface.co/papers/2310.06825). The novelty of Gemma 2 is that a sliding window is applied to every other layer (local - 4096 tokens), while the layers in between still use full quadratic global attention (8192 tokens). We suppose this is a way to increase quality in long context situations (half of the layers still attend to all tokens) while partially benefiting from the advantages of sliding attention. 

### Soft-capping and attention implementations

Soft capping is a technique that prevents logits from growing excessively large without truncating them. It works by dividing the logits by a maximum value threshold (soft_cap), then passing them through a `tanh` layer (ensuring they are in the `(-1, 1)` range), and finally multiplying by the threshold again. This guarantees that the final values will be in the `(-soft_cap, +soft_cap)` interval without losing much information but stabilizing the training.

Putting it all together, the logits are calculated by: `logits ← soft_cap ∗ tanh(logits/soft_cap)`

Gemma 2 employs soft capping for the final layer and for every attention layer. The attention logits are capped at 50.0, and the final logits at 30.0.

At the time of release, soft-capping is incompatible with Flash Attention / SDPA, but they can still be used in inference for maximum efficiency. The Gemma 2 team observed very minor differences when soft-capping is removed during inference.

*Note: For stable fine-tuning runs, you still need to enable soft-capping and hence, we recommend fine-tuning with `eager` attention instead of SDPA.*

### Knowledge Distillation

Knowledge distillation is a popular technique for training a smaller *student* model to mimic the behavior of a larger but better-performing *teacher.* This works by augmenting the next-token prediction task of LLMs with a distribution of token probabilities from the teacher (e.g., GPT-4, Claude, or Gemini), which provides a richer signal for the student to learn from. 

According to the Gemma 2 tech report, knowledge distillation was used to pre-train the 9B model, while the 27B model was pre-trained from scratch.

For post-training, the Gemma 2 team generated a diverse set of completions from a teacher (unspecified in the report, but presumably Gemini Ultra), and then trained the student models on this synthetic data with SFT. This is the basis of many open models, such as [Zephyr](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta) and [OpenHermes](https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B), which are trained entirely on synthetic data from larger LLMs.

Although effective, this method has drawbacks since the model capacity mismatch between the student and teacher can lead to a *train-inference mismatch*, where the text generated by the student during inference is out-of-distribution compared to that seen during training.

To handle this issue, the Gemma 2 team used [“on-policy distillation”](https://arxiv.org/pdf/2306.13649), where the student generates completions from the SFT prompts. These completions are then used to compute the KL divergence between the teacher’s and student’s logits. By minimizing the KL divergence throughout training, the student learns to model the behavior of the teacher accurately while also minimizing the train-inference mismatch.

This approach is quite interesting, as we’ve seen in the community that on-policy methods like online DPO produce stronger models, and one advantage of on-policy distillation is that you only need the logits from the teacher, so you don’t need to rely on reward models or LLM-as-a-judge to improve the model. It will be exciting to see if this method becomes more popular among fine-tuners in the coming months!

### Model Merging

[Model merging](https://huggingface.co/blog/mlabonne/merge-models) is a technique that combines two or more LLMs into a single new model. It's relatively new and experimental and can be used without accelerators. [Mergekit](https://github.com/arcee-ai/mergekit) is a popular open-source toolkit for merging LLMs. It implements linear, SLERP, TIES, DARE, and other merging techniques.

According to the Technical Report, Gemma 2 used [Warp](https://arxiv.org/abs/2406.16768), a new merging technique that merges models in three distinct stages: 

1. Exponential Moving Average (EMA): This is applied during the reinforcement learning (RL) fine-tuning process.
2. Spherical Linear intERPolation (SLERP): This is applied after the RL fine-tuning of multiple policies.
3. Linear Interpolation Towards Initialization (LITI): This stage is applied after the SLERP stage.

## Gemma 2 evaluation

How good are the Gemma models? Below are performance comparisons to other open models based on the Technical Report and the new version of the [open LLM Leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard).

### Technical Report results

This Technical Report of Gemma 2 compares the performance of different open LLMs on the previous Open LLM Leaderboard benchmarks.

|            | Llama 3 (70B) | Qwen 1.5 (32B) | Gemma 2 (27B) |
|