Upgrade Flax NNX Gemma Sampling Inference doc #4325

8bitmp3 · 2024-10-23T21:25:11Z

Preview: https://flax--4325.org.readthedocs.build/en/4325/guides/gemma.html

Also fixes broken code after:

! git clone https://github.com/google/flax.git flax_examples

...
- sys.path.append("./flax_examples/flax/nnx/examples/gemma")
+ sys.path.append("./flax_examples/examples/gemma")
...

review-notebook-app · 2024-10-23T21:25:17Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

8bitmp3 · 2024-10-23T21:45:35Z

docs_nnx/guides/gemma.ipynb

-    "You will find in this colab a detailed tutorial explaining how to use NNX to load a Gemma checkpoint and sample from it."
+    "In this tutorial, you will learn step-by-step how to use Flax NNX to load the [Gemma](https://ai.google.dev/gemma) open model files and use them to perform sampling/inference for generating text. You will use the [Flax NNX `gemma` code](https://github.com/google/flax.git) that was written with Flax and JAX.\n",
+    "\n",
+    "> Gemma is a family of lightweight, state-of-the-art open models based on Google DeepMind’s [Gemini](https://deepmind.google/technologies/gemini/#introduction). Read more about [Gemma](https://blog.google/technology/developers/gemma-open-models/) and [Gemma 2](https://blog.google/technology/developers/google-gemma-2/).\n",


Similar to what we did in other Gemma docs - added some background.

8bitmp3 · 2024-10-23T21:46:28Z

docs_nnx/guides/gemma.ipynb

    "\n",
-    "You will find in this colab a detailed tutorial explaining how to use NNX to load a Gemma checkpoint and sample from it."


"guide" in the title, "tutorial" in the first paragraph -> let's use "tutorial".

8bitmp3 · 2024-10-23T21:47:46Z

docs_nnx/guides/gemma.ipynb

    "\n",
-    "Now select and download the checkpoint you want to try. Note that you will need an A100 runtime for the 7b models."


@cgarciae Since there are checkpoints and tokenizer files, changed to "model" instead of "checkpoint".

Checking if free TPU v2-8 is sufficient.

8bitmp3 · 2024-10-23T21:48:46Z

docs_nnx/guides/gemma.ipynb

@@ -19,16 +19,24 @@
   "cell_type": "markdown",
   "metadata": {},
   "source": [
-    "# Getting Started with Gemma Sampling using NNX: A Step-by-Step Guide\n",


@cgarciae Adding "inference" next to "sampling" for search.

8bitmp3 · 2024-10-23T21:51:15Z

docs_nnx/guides/gemma.ipynb

    "\n",
-    "1. Visit https://www.kaggle.com/ and create an account.\n",
-    "2. Go to your account settings, then the 'API' section.\n",
-    "3. Click 'Create new token' to download your key.\n",


@cgarciae Adding Step 3 "OPTIONAL" and removing "OPTIONAL". since Colab asks for access here after running the code below, so users won't have to manually entering the details if they are stored in Colab:

import kagglehub kagglehub.login()

"1. To create an account, visit Kaggle and click on 'Register'."
"2. If/once you have an account, you need to sign in, go to your 'Settings', and under 'API' click on 'Create New Token' to generate and download your Kaggle API key."
"3. OPTIONAL: In Google Colab, under 'Secrets' add your Kaggle username and API key, storing the username as KAGGLE_USERNAME and the key as KAGGLE_KEY. If you are using a Kaggle Notebook for free TPU or other hardware acceleration, it has a key storage feature under 'Add-ons' > 'Secrets', along with instructions for accessing stored keys."

8bitmp3 · 2024-10-23T21:52:19Z

docs_nnx/guides/gemma.ipynb

@@ -82,13 +90,21 @@
   "cell_type": "markdown",
   "metadata": {},
   "source": [
-    "If everything went well, you should see:\n",


Adding an extra optional step here, similar to what we have in the Gemma docs.

"Note: In Google Colab, you can instead authenticate into Kaggle using the code below after following the optional step 3 from above...."

8bitmp3 · 2024-10-23T21:53:09Z

docs_nnx/guides/gemma.ipynb

@@ -124,7 +140,7 @@
   "cell_type": "markdown",
   "metadata": {},
   "source": [
-    "Flax examples are not exposed as packages so you need to use the workaround in the next cells to import from NNX's Gemma example."


@cgarciae Edited:

"To interact with the Gemma model, you will use the Flax NNX Gemma code from google/flax examples on GitHub. Since it is not exposed as packages, you need to use the following workaround in the next cells to import from the Flax NNX Gemma example."

8bitmp3 · 2024-10-23T21:54:12Z

docs_nnx/guides/gemma.ipynb

@@ -195,7 +218,9 @@
   "cell_type": "markdown",
   "metadata": {},
   "source": [
-    "Use the `transformer_lib.TransformerConfig.from_params` function to automatically load the correct configuration from a checkpoint. Note that the vocabulary size is smaller than the number of input embeddings due to unused tokens in this release."


Added the source code that has more docstring(s) since transformer_lib is an alias for Flax NNX examples -> gemma.transformer:

"Then, use the Flax NNX transformer_lib.TransformerConfig.from_params function to automatically load the correct configuration from a checkpoint."

8bitmp3 · 2024-10-23T21:54:36Z

docs_nnx/guides/gemma.ipynb

@@ -212,7 +237,9 @@
   "cell_type": "markdown",
   "metadata": {},
   "source": [
-    "Finally, build a sampler on top of your model and your tokenizer."


Added the source code with the docstring.

"Build a Flax NNX Sampler on top of your model and tokenizer with the right parameter shapes."

8bitmp3 · 2024-10-23T21:55:07Z

docs_nnx/guides/gemma.ipynb

@@ -235,7 +261,11 @@
   "cell_type": "markdown",
   "metadata": {},
   "source": [
-    "You're ready to start sampling ! This sampler uses just-in-time compilation, so changing the input shape triggers recompilation, which can slow things down. For the fastest and most efficient results, keep your batch size consistent."


Added some background on JAX JIT after studying the source code (it's not NNX JIT).

"Note: This Flax NNX gemma.Sampler uses JAX’s just-in-time (JIT) compilation, so changing the input shape triggers recompilation, which can slow things down. For the fastest and most efficient results, keep your batch size consistent."

8bitmp3 · 2024-10-23T22:12:32Z

docs_nnx/guides/gemma.ipynb

@@ -136,6 +152,14 @@
    "! git clone https://github.com/google/flax.git flax_examples"


Fixing

! git clone https://github.com/google/flax.git flax_examples

... - sys.path.append("./flax_examples/flax/nnx/examples/gemma") + sys.path.append("./flax_examples/examples/gemma") ...

8bitmp3 · 2024-10-23T22:35:19Z

Sampler configuration (https://github.com/google/flax/blob/main/examples/gemma/sampler.py):

    transformer=transformer,
    vocab=vocab,
    params=params['transformer'],
)

Throws an error

TypeError                                 Traceback (most recent call last)
...
in <cell line: 1>()
----> 1 sampler = sampler_lib.Sampler(
      2     transformer=transformer,
      3     vocab=vocab,
      4     params=params['transformer'],
      5 )

TypeError: Sampler.__init__() got an unexpected keyword argument 'params'

@cgarciae PTAL

8bitmp3 · 2024-10-23T22:38:54Z

docs_nnx/guides/gemma.ipynb

-    "3. Click 'Create new token' to download your key.\n",
+    "1. To create an account, visit [Kaggle](https://www.kaggle.com/) and click on 'Register'.\n",
+    "2. If/once you have an account, you need to sign in, go to your ['Settings'](https://www.kaggle.com/settings), and under 'API' click on 'Create New Token' to generate and download your Kaggle API key.\n",
+    "3. OPTIONAL: In [Google Colab](https://colab.research.google.com/), under 'Secrets' add your Kaggle username and API key, storing the username as `KAGGLE_USERNAME` and the key as `KAGGLE_KEY`. If you are using a [Kaggle Notebook](https://www.kaggle.com/code) for free TPU or other hardware acceleration, it has a key storage feature under 'Add-ons' > 'Secrets', along with instructions for accessing stored keys.\n",


TODO: Should remove Optional for Colab users?

cgarciae · 2024-10-30T18:35:53Z

Hey @8bitmp3! I cleaned up this guide a little bit. Can you take a look at the new version?

8bitmp3 · 2024-10-31T09:38:14Z

thanks @cgarciae 👍
on it

8bitmp3 · 2024-11-04T22:20:52Z

Reopening after #4334 fixes

8bitmp3 requested review from cgarciae and IvyZX October 23, 2024 21:25

8bitmp3 self-assigned this Oct 23, 2024

8bitmp3 commented Oct 23, 2024

View reviewed changes

8bitmp3 force-pushed the update-nnx-gemma branch from e7e7195 to 243439f Compare October 23, 2024 22:03

8bitmp3 commented Oct 23, 2024

View reviewed changes

8bitmp3 force-pushed the update-nnx-gemma branch from 243439f to 52dc69b Compare October 23, 2024 22:19

8bitmp3 commented Oct 23, 2024

View reviewed changes

8bitmp3 marked this pull request as ready for review October 28, 2024 21:43

8bitmp3 closed this Nov 4, 2024

8bitmp3 force-pushed the update-nnx-gemma branch from 52dc69b to d8b1a92 Compare November 4, 2024 22:11

Upgrade Flax NNX Gemma

7eb2405

8bitmp3 reopened this Nov 4, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Upgrade Flax NNX Gemma Sampling Inference doc #4325

Upgrade Flax NNX Gemma Sampling Inference doc #4325

8bitmp3 commented Oct 23, 2024 •

edited

Loading

review-notebook-app bot commented Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024

8bitmp3 Oct 23, 2024 •

edited

Loading

8bitmp3 commented Oct 23, 2024

8bitmp3 Oct 23, 2024

cgarciae commented Oct 30, 2024

8bitmp3 commented Oct 31, 2024

8bitmp3 commented Nov 4, 2024

		"\n",
		"You will find in this colab a detailed tutorial explaining how to use NNX to load a Gemma checkpoint and sample from it."

		"\n",
		"Now select and download the checkpoint you want to try. Note that you will need an A100 runtime for the 7b models."

		@@ -136,6 +152,14 @@
		"! git clone https://github.com/google/flax.git flax_examples"

Upgrade Flax NNX Gemma Sampling Inference doc #4325

Are you sure you want to change the base?

Upgrade Flax NNX Gemma Sampling Inference doc #4325

Conversation

8bitmp3 commented Oct 23, 2024 • edited Loading

review-notebook-app bot commented Oct 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

8bitmp3 Oct 23, 2024 • edited Loading

Choose a reason for hiding this comment

8bitmp3 commented Oct 23, 2024

Choose a reason for hiding this comment

cgarciae commented Oct 30, 2024

8bitmp3 commented Oct 31, 2024

8bitmp3 commented Nov 4, 2024

8bitmp3 commented Oct 23, 2024 •

edited

Loading

8bitmp3 Oct 23, 2024 •

edited

Loading