Erland / gemma-7b-JAX

google/gemma-7b - JAX/Flax

This repository contains the JAX/Flax version of the google/gemma-7b model, originally a PyTorch model from google. This conversion enables efficient inference and training on TPUs and GPUs using the JAX/Flax framework.

Model Description

google/gemma-7b is a transformer-based language model developed by google.

Conversion Details

This model was converted from the original PyTorch implementation to JAX/Flax. The conversion process involved the following steps:

1. Loading the PyTorch model and configuration: The pretrained PyTorch model and its configuration were loaded using the Hugging Face Transformers library.
2. Creating an equivalent Flax model architecture: A Flax model with the same architecture as the original PyTorch model was created.
3. Converting the PyTorch weights to Flax format: The weights from the PyTorch model were converted to the Flax format using the convert_pytorch_state_dict_to_flax utility function provided by Hugging Face.
4. Verifying the converted weights: The converted Flax weights were compared against the original PyTorch weights to ensure that the conversion process was performed accurately.

Important Note about max_position_embeddings

During the conversion process, it was necessary to modify the max_position_embeddings parameter in the model's configuration. The original value of {original_max_pos_embed} led to out-of-memory (OOM) errors on the hardware used for conversion. To resolve this, max_position_embeddings was adjusted to {new_max_pos_embed}.

Implications of this change:

• The model may not be able to handle sequences longer than 512 tokens without truncation or other modifications.
• If you fine-tune this model, keep in mind the revised max_position_embeddings when preparing your training data.

Weight Comparison

The following table summarizes the comparison between the weights of the original PyTorch model and the converted JAX/Flax model. This detailed verification confirms that the conversion was accurate and that both models should produce (approximately) the same outputs given the same inputs.

Note:

• Allclose indicates whether the weights are approximately equal within the specified relative ( rtol=1e-5 ) and absolute ( atol=1e-3 ) tolerances using jnp.allclose() .
• Max Diff , Mean Diff , and Std Diff provide further details on the differences between the weights if Allclose is False , which might be expected for some layers due to numerical precision differences between frameworks.

Hardware Used for Conversion

The conversion process was performed on the following hardware configuration:

• CPU: x86_64
• RAM: 334.56 GB
• OS: Linux-6.1.85+-x86_64-with-glibc2.35
• JAX version: 0.5.2
• Flax version: 0.10.5
• Transformers version: 4.50.3
• GPU: None

This conversion took approximately 553.40 seconds to complete.

Usage

Here's how you can use the converted model in JAX/Flax for text generation:

import jax
import jax.numpy as jnp
from transformers import FlaxAutoModelForCausalLM, AutoTokenizer

model_name = "Erland/gemma-7b-JAX"  # Replace with your repository name
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = FlaxAutoModelForCausalLM.from_pretrained(model_name, from_pt=False) # from_pt should be False since it's already flax

# Example prompt
prompt = "The quick brown fox"

# Tokenize the prompt
tokenized_prompt = tokenizer(prompt, return_tensors="np")

# Generate text
output_ids = model.generate(tokenized_prompt.input_ids, max_length=50)

# Decode the generated text
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

Limitations

Sequence Length: As mentioned earlier, the max_position_embeddings has been modified to 512. Be mindful of this limitation when working with long sequences.

Numerical Precision: Minor differences in outputs compared to the original PyTorch model might be observed due to numerical precision variations between PyTorch and JAX/Flax, particularly on different hardware.

Acknowledgements

We thank the original authors of google/gemma-7b at google for their groundbreaking work in developing this powerful language model.

We acknowledge the Hugging Face Transformers library for providing the essential tools and infrastructure that made this conversion possible.

Thanks to the JAX and Flax teams for developing such performant and flexible frameworks for numerical computation and deep learning.

License

This JAX/Flax model is released under the original model license.