Tuesday, April 16, 2024
HomeSoftware DevelopmentIntroducing Gemma fashions in Keras

Introducing Gemma fashions in Keras

Posted by Martin Görner – Product Supervisor, Keras

The Keras workforce is completely satisfied to announce that Gemma, a household of light-weight, state-of-the artwork open fashions constructed from the identical analysis and expertise that we used to create the Gemini fashions, is now out there within the KerasNLP assortment. Due to Keras 3, Gemma runs on JAX, PyTorch and TensorFlow. With this launch, Keras can also be introducing a number of new options particularly designed for big language fashions: a brand new LoRA API (Low Rank Adaptation) and huge scale model-parallel coaching capabilities.

If you wish to dive immediately into code samples, head right here:

Get began

Gemma fashions are available in moveable 2B and 7B parameter sizes, and ship important advances in opposition to comparable open fashions, and even some bigger ones. For instance:

  • Gemma 7B scores a brand new best-in class 64.3% of appropriate solutions within the MMLU language understanding benchmark (vs. 62.5% for Mistral-7B and 54.8% for Llama2-13B)
  • Gemma provides +11 proportion factors to the GSM8K benchmark rating for grade-school math issues (46.4% for Gemma 7B vs. Mistral-7B 35.4%, Llama2-13B 28.7%)
  • and +6.1 proportion factors of appropriate solutions in HumanEval, a coding problem (32.3% for Gemma 7B, vs. Mistral 7B 26.2%, Llama2 13B 18.3%).

Gemma fashions are supplied with a well-known KerasNLP API and a super-readable Keras implementation. You may instantiate the mannequin with a single line of code:

gemma_lm = keras_nlp.fashions.GemmaCausalLM.from_preset("gemma_2b_en")

And run it immediately on a textual content immediate – sure, tokenization is built-in, though you may simply break up it out if wanted – learn the Keras NLP information to see how.

gemma_lm.generate("Keras is a", max_length=32)
> "Keras is a well-liked deep studying framework for neural networks..."

Attempt it out right here: Get began with Gemma fashions

Fantastic-tuning Gemma Fashions with LoRA

Due to Keras 3, you may select the backend on which you run the mannequin. Right here is find out how to swap:

os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch".
import keras # import keras after having chosen the backend

Keras 3 comes with a number of new options particularly for big language fashions. Chief amongst them is a brand new LoRA API (Low Rank Adaptation) for parameter-efficient fine-tuning. Right here is find out how to activate it:

# Notice: rank=4 replaces the weights matrix of related layers with the 
# product AxB of two matrices of rank 4, which reduces the quantity of 
# trainable parameters.

This single line drops the variety of trainable parameters from 2.5 billion to 1.3 million!

Attempt it out right here: Fantastic-tune Gemma fashions with LoRA.

Fantastic-tuning Gemma fashions on a number of GPU/TPUs

Keras 3 additionally helps large-scale mannequin coaching and Gemma is the proper mannequin to strive it out. The brand new Keras distribution API gives data-parallel and model-parallel distributed coaching choices. The brand new API is supposed to be multi-backend however in the meanwhile, it’s carried out for the JAX backend solely, due to its confirmed scalability (Gemma fashions had been educated with JAX).

To fine-tune the bigger Gemma 7B, a distributed setup is beneficial, for instance a TPUv3 with 8 TPU cores that you may get without cost on Kaggle, or an 8-GPU machine from Google Cloud. Right here is find out how to configure the mannequin for distributed coaching, utilizing mannequin parallelism:

device_mesh = keras.distribution.DeviceMesh(
   (1, 8), # Mesh topology
   ["batch", "model"], # named mesh axes
   gadgets=keras.distribution.list_devices() # precise accelerators

# Mannequin config
layout_map = keras.distribution.LayoutMap(device_mesh)
layout_map["token_embedding/embeddings"] = (None, "mannequin")
layout_map["decoder_block.*attention.*(query|key|value).*kernel"] = (
   None, "mannequin", None)
layout_map["decoder_block.*attention_output.*kernel"] = (
   None, None, "mannequin")
layout_map["decoder_block.*ffw_gating.*kernel"] = ("mannequin", None)
layout_map["decoder_block.*ffw_linear.*kernel"] = (None, "mannequin")

# Set the mannequin config and load the mannequin
model_parallel = keras.distribution.ModelParallel(
   device_mesh, layout_map, batch_dim_name="batch")
gemma_lm = keras_nlp.fashions.GemmaCausalLM.from_preset("gemma_7b_en")
# Prepared: now you can prepare with mannequin.match() or generate textual content with generate()

What this code snippet does is about up the 8 accelerators right into a 1 x 8 matrix the place the 2 dimensions are referred to as “batch” and “mannequin”. Mannequin weights are sharded on the “mannequin” dimension, right here break up between the 8 accelerators, whereas knowledge batches will not be partitioned because the “batch” dimension is 1.

Attempt it out right here: Fantastic-tune Gemma fashions on a number of GPUs/TPUs.

What’s Subsequent

We are going to quickly be publishing a information exhibiting you find out how to appropriately partition a Transformer mannequin and write the 6 traces of partitioning setup above. It’s not very lengthy however it might not match on this submit.

You should have observed that layer partitionings are outlined by regexes on layer names. You may examine layer names with this code snippet. We ran this to assemble the LayoutMap above.

# That is for the primary Transformer block solely,
# however all of them have the identical construction
tlayer = gemma_lm.spine.get_layer('decoder_block_0')
for variable in tlayer.weights:
 print(f'{variable.path:<58}  {str(variable.form):<16}')

Full GSPMD mannequin parallelism works right here with just some partitioning hints as a result of Keras passes these settings to the highly effective XLA compiler which figures out all the opposite particulars of the distributed computation.

We hope you’ll take pleasure in enjoying with Gemma fashions. Right here can also be an instruction-tuning tutorial that you simply may discover helpful. And by the best way, if you wish to share your fine-tuned weights with the group, the Kaggle Mannequin Hub now helps user-tuned weights uploads. Head to the mannequin web page for Gemma fashions on Kaggle and see what others have already created!



Please enter your comment!
Please enter your name here

Most Popular

Recent Comments