Gemma-2-9b-en-tpu-fine-tuning.ipynb

{"metadata":{"colab":{"provenance":[]},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":5385487,"sourceType":"datasetVersion","datasetId":3122881},{"sourceId":82367,"sourceType":"modelInstanceVersion","modelInstanceId":56633,"modelId":78150},{"sourceId":85984,"sourceType":"modelInstanceVersion","modelInstanceId":72244,"modelId":78150}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"papermill":{"default_parameters":{},"duration":672.253242,"end_time":"2024-02-21T10:06:37.294427","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-02-21T09:55:25.041185","version":"2.5.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Overview\n\nGemma is a family of lightweight, state-of-the-art open models built from research and technology used to create Google Gemini models. Gemma can be further finetuned to suit specific needs. But Large Language Models, such as Gemma, can be very large in size and some of them may not fit on a sing accelerator for finetuning.\n\nIn the case of Gemma 2, the 9 billion model is too large to fit on a single Kaggle accelerator. This can make both fine-tuning and inference in Kaggle Notebooks difficult.\n\nOne option is quantizing a model. Quantization can allow efficient inference for larger models with less GPU or TPU memory.\n\nAnother option is model parallelism, which splits the model's parameters across a number of accelerators, and can be used easy for both inference and fine-tuning. This guide will show how to fine-tune a Gemma 2 model on 8 TPU V3 cores through the Kaggle TPU runtime. You can find out more about distributed training in this [Keras guide](https://keras.io/guides/distribution/).","metadata":{"id":"Tdlq6K0znh3O","papermill":{"duration":0.005633,"end_time":"2024-02-21T09:55:26.83679","exception":false,"start_time":"2024-02-21T09:55:26.831157","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Using accelerators\n\nTechnically you can use either TPU or GPU for this tutorial.\n\n### Notes on TPU environments\n\nGoogle has 3 products that provide TPUs:\n* [Colab](https://colab.sandbox.google.com/) provides TPU v2, which is not sufficient for this tutorial.\n* [Kaggle](https://www.kaggle.com/) offers TPU v3 for free and they work for this tutorial.\n* [Cloud TPU](https://cloud.google.com/tpu?hl=en) offers TPU v3 and newer generations. One way to set it up is:\n  1. Create a new [TPU VM](https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm#tpu-vms)\n  2. Set up [SSH port forwarding](https://cloud.google.com/solutions/connecting-securely#port-forwarding-over-ssh) for your intended Jupyter server port\n  3. Install Jupyter and start it on the TPU VM, then connect to Colab through \"Connect to a local runtime\"\n\n### Notes on multi-GPU setup\n\nAlthough this tutorial focuses on the TPU use case, you can easily adapt it for your own needs if you have a multi-GPU machine.\n\nIf you prefer to work through Colab, it's also possible to provision a multi-GPU VM for Colab directly through \"Connect to a custom GCE VM\" in the Colab Connect menu.\n\n\nWe will focus on using the **free TPU from Kaggle** here.","metadata":{"id":"z-jBO5hmDwrc","papermill":{"duration":0.005425,"end_time":"2024-02-21T09:55:26.847611","exception":false,"start_time":"2024-02-21T09:55:26.842186","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Before you begin","metadata":{"papermill":{"duration":0.005257,"end_time":"2024-02-21T09:55:26.85841","exception":false,"start_time":"2024-02-21T09:55:26.853153","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### Gemma setup\n\nTo complete this tutorial, you first need to accept the Gemma Terms of Use. You can navigate to the [Keras Gemma 2 Page](https://www.kaggle.com/models/keras/gemma2) to do this. You will see a banner at the top of the page with a button to \"Request Access\" if you have not already done this for your Kaggle user.","metadata":{"id":"aKvTsIkL98BG","papermill":{"duration":0.005259,"end_time":"2024-02-21T09:55:26.869282","exception":false,"start_time":"2024-02-21T09:55:26.864023","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Installation\n\nInstall Keras and KerasNLP with the Gemma model.","metadata":{"id":"AO7a1Q4Yyc9Z","papermill":{"duration":0.005175,"end_time":"2024-02-21T09:55:26.880069","exception":false,"start_time":"2024-02-21T09:55:26.874894","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!pip install -q -U keras-nlp tensorflow-text\n# Install tensorflow-cpu so tensorflow does not attempt to access the TPU.\n!pip install -q -U tensorflow-cpu","metadata":{"id":"WWEzVJR4Fx9g","papermill":{"duration":37.05282,"end_time":"2024-02-21T09:56:03.93859","exception":false,"start_time":"2024-02-21T09:55:26.88577","status":"completed"},"tags":[],"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-08-31T13:17:53.634774Z","iopub.execute_input":"2024-08-31T13:17:53.635093Z","iopub.status.idle":"2024-08-31T13:19:43.245273Z","shell.execute_reply.started":"2024-08-31T13:17:53.635055Z","shell.execute_reply":"2024-08-31T13:19:43.244252Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stdout","text":"\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\ntf-keras 2.16.0 requires tensorflow<2.17,>=2.16, but you have tensorflow 2.17.0 which is incompatible.\u001b[0m\u001b[31m\n\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.2\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.2\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n","output_type":"stream"}]},{"cell_type":"markdown","source":"### Set up Keras JAX backend","metadata":{"id":"fr9VnPm7FoMf","papermill":{"duration":0.006513,"end_time":"2024-02-21T09:56:03.951855","exception":false,"start_time":"2024-02-21T09:56:03.945342","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Import JAX and run a sanity check on TPU. Kaggle offers TPUv3-8 devices which have 8 TPU cores with 16GB of memory each.","metadata":{"id":"lbZsUvfhwL2D","papermill":{"duration":0.00599,"end_time":"2024-02-21T09:56:03.964473","exception":false,"start_time":"2024-02-21T09:56:03.958483","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import jax\n\njax.devices()","metadata":{"id":"BK4MpHLKGujb","outputId":"a60376b8-0937-45fc-809b-33eaa92cbc6c","papermill":{"duration":8.126711,"end_time":"2024-02-21T09:56:12.097265","exception":false,"start_time":"2024-02-21T09:56:03.970554","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:19:43.246862Z","iopub.execute_input":"2024-08-31T13:19:43.247138Z","iopub.status.idle":"2024-08-31T13:19:52.380549Z","shell.execute_reply.started":"2024-08-31T13:19:43.247104Z","shell.execute_reply":"2024-08-31T13:19:52.379698Z"},"trusted":true},"execution_count":2,"outputs":[{"name":"stderr","text":"WARNING: Logging before InitGoogle() is written to STDERR\nE0000 00:00:1725110388.292273      13 common_lib.cc:798] Could not set metric server port: INVALID_ARGUMENT: Could not find SliceBuilder port 8471 in any of the 0 ports provided in `tpu_process_addresses`=\"local\"\n=== Source Location Trace: ===\nlearning/45eac/tfrc/runtime/common_lib.cc:479\nE0831 13:19:48.338083746      13 oauth2_credentials.cc:238]            oauth_fetch: UNKNOWN:C-ares status is not ARES_SUCCESS qtype=A name=metadata.google.internal. is_balancer=0: Domain name not found {created_time:\"2024-08-31T13:19:48.338067947+00:00\", grpc_status:2}\n","output_type":"stream"},{"execution_count":2,"output_type":"execute_result","data":{"text/plain":"[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"},"metadata":{}}]},{"cell_type":"code","source":"import os\n\n# The Keras 3 distribution API is only implemented for the JAX backend for now\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n# Pre-allocate all TPU memory to minimize memory fragmentation and allocation overhead.\nos.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"1.0\"","metadata":{"id":"WEgg_OVIL2HY","papermill":{"duration":0.01342,"end_time":"2024-02-21T09:56:12.117529","exception":false,"start_time":"2024-02-21T09:56:12.104109","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:19:52.381605Z","iopub.execute_input":"2024-08-31T13:19:52.381972Z","iopub.status.idle":"2024-08-31T13:19:52.385894Z","shell.execute_reply.started":"2024-08-31T13:19:52.381942Z","shell.execute_reply":"2024-08-31T13:19:52.385200Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"markdown","source":"## Load model","metadata":{"id":"wo1xkzr62hXN","papermill":{"duration":0.00603,"end_time":"2024-02-21T09:56:12.129531","exception":false,"start_time":"2024-02-21T09:56:12.123501","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import keras\nimport keras_nlp","metadata":{"id":"kFCmWEKdMA_Y","outputId":"57359fb1-ea3e-482a-aaf3-dcffc265a5ec","papermill":{"duration":8.153944,"end_time":"2024-02-21T09:56:20.315559","exception":false,"start_time":"2024-02-21T09:56:12.161615","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:19:52.387545Z","iopub.execute_input":"2024-08-31T13:19:52.387807Z","iopub.status.idle":"2024-08-31T13:20:01.282788Z","shell.execute_reply.started":"2024-08-31T13:19:52.387780Z","shell.execute_reply":"2024-08-31T13:20:01.281784Z"},"trusted":true},"execution_count":4,"outputs":[{"name":"stderr","text":"/usr/local/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n  from .autonotebook import tqdm as notebook_tqdm\n","output_type":"stream"}]},{"cell_type":"markdown","source":"To load the model with the weights and tensors distributed across TPUs, first create a new `DeviceMesh`. `DeviceMesh` represents a collection of hardware devices configured for distributed computation and was introduced in Keras 3 as part of the unified distribution API.\n\nThe distribution API enables data and model parallelism, allowing for efficient scaling of deep learning models on multiple accelerators and hosts. It leverages the underlying framework (e.g. JAX) to distribute the program and tensors according to the sharding directives through a procedure called single program, multiple data (SPMD) expansion. Check out more details in the new [Keras 3 distribution API guide](https://keras.io/guides/distribution/).","metadata":{"id":"xrR8TpVS6uPs","papermill":{"duration":0.005852,"end_time":"2024-02-21T09:56:20.357594","exception":false,"start_time":"2024-02-21T09:56:20.351742","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Create a device mesh with (1, 8) shape so that the weights are sharded across\n# all 8 TPUs.\ndevice_mesh = keras.distribution.DeviceMesh(\n    (1, 8),\n    [\"batch\", \"model\"],\n    devices=keras.distribution.list_devices(),\n)","metadata":{"id":"7gxEkpUiP1Qf","papermill":{"duration":0.012454,"end_time":"2024-02-21T09:56:20.375812","exception":false,"start_time":"2024-02-21T09:56:20.363358","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:20:01.283916Z","iopub.execute_input":"2024-08-31T13:20:01.284395Z","iopub.status.idle":"2024-08-31T13:20:01.288316Z","shell.execute_reply.started":"2024-08-31T13:20:01.284364Z","shell.execute_reply":"2024-08-31T13:20:01.287649Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"markdown","source":"`LayoutMap` from the distribution API specifies how the weights and tensors should be sharded or replicated, using the string keys, for example, `token_embedding/embeddings` below, which are treated like regex to match tensor paths. Matched tensors are sharded with model dimensions (8 TPUs); others will be fully replicated.","metadata":{"id":"gTSJUwkC-7c6","papermill":{"duration":0.00582,"end_time":"2024-02-21T09:56:20.387661","exception":false,"start_time":"2024-02-21T09:56:20.381841","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model_dim = \"model\"\n\nlayout_map = keras.distribution.LayoutMap(device_mesh)\n\n# Weights that match 'token_embedding/embeddings' will be sharded on 8 TPUs\nlayout_map[\"token_embedding/embeddings\"] = (model_dim, None)\n# Regex to match against the query, key and value matrices in attention layers\nlayout_map[\"decoder_block.*attention.*(query|key|value)/kernel\"] = (model_dim, None, None)\nlayout_map[\"decoder_block.*attention_output/kernel\"] = (model_dim, None, None)\nlayout_map[\"decoder_block.*ffw_gating.*/kernel\"] = (None, model_dim)\nlayout_map[\"decoder_block.*ffw_linear/kernel\"] = (model_dim, None)","metadata":{"id":"8Wgh8h0qQCcu","papermill":{"duration":0.013198,"end_time":"2024-02-21T09:56:20.406598","exception":false,"start_time":"2024-02-21T09:56:20.3934","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:20:01.289106Z","iopub.execute_input":"2024-08-31T13:20:01.289351Z","iopub.status.idle":"2024-08-31T13:20:01.307071Z","shell.execute_reply.started":"2024-08-31T13:20:01.289326Z","shell.execute_reply":"2024-08-31T13:20:01.306440Z"},"trusted":true},"execution_count":6,"outputs":[]},{"cell_type":"markdown","source":"`ModelParallel` allows you to shard model weights or activation tensors across all devcies on the `DeviceMesh`. In this case, some of the Gemma 7B model weights are sharded across 8 TPU chips according the `layout_map` defined above. Now load the model in the distributed way.","metadata":{"id":"6n4Zlvk9ALhZ","papermill":{"duration":0.005938,"end_time":"2024-02-21T09:56:20.418485","exception":false,"start_time":"2024-02-21T09:56:20.412547","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model_parallel = keras.distribution.ModelParallel(\n    layout_map=layout_map,\n    batch_dim_name=\"batch\",\n)\n\nkeras.distribution.set_distribution(model_parallel)\ngemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma2_9b_en\")\ngemma_lm.summary()","metadata":{"id":"bu48vUnbQj0p","outputId":"fd216acb-852c-46f4-e8ac-2d8f91362d24","papermill":{"duration":145.668669,"end_time":"2024-02-21T09:58:46.092826","exception":false,"start_time":"2024-02-21T09:56:20.424157","status":"completed"},"tags":[],"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-08-31T13:20:01.307863Z","iopub.execute_input":"2024-08-31T13:20:01.308078Z","iopub.status.idle":"2024-08-31T13:23:21.053949Z","shell.execute_reply.started":"2024-08-31T13:20:01.308056Z","shell.execute_reply":"2024-08-31T13:23:21.052756Z"},"trusted":true},"execution_count":7,"outputs":[{"name":"stderr","text":"normalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m3584\u001b[0m)        │   \u001b[38;5;34m9,241,705,984\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m917,504,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">3584</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">9,241,705,984</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">917,504,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Total params: \u001b[0m\u001b[38;5;34m9,241,705,984\u001b[0m (34.43 GB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">9,241,705,984</span> (34.43 GB)\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m9,241,705,984\u001b[0m (34.43 GB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">9,241,705,984</span> (34.43 GB)\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> (0.00 B)\n</pre>\n"},"metadata":{}}]},{"cell_type":"markdown","source":"Let's verify that the model has been partitioned correctly. Let's take `decoder_block_1` as an example.","metadata":{"id":"ORCOIawAvpZ1","papermill":{"duration":0.006357,"end_time":"2024-02-21T09:58:46.105898","exception":false,"start_time":"2024-02-21T09:58:46.099541","status":"completed"},"tags":[]}},{"cell_type":"code","source":"decoder_block_1 = gemma_lm.backbone.get_layer('decoder_block_1')\nprint(type(decoder_block_1))\nfor variable in decoder_block_1.weights:\n  print(f'{variable.path:<48}  {str(variable.shape):<14}  {str(variable.value.sharding.spec)}')","metadata":{"id":"DqT7TRHKvoMK","papermill":{"duration":0.014281,"end_time":"2024-02-21T09:58:46.126669","exception":false,"start_time":"2024-02-21T09:58:46.112388","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:23:21.055168Z","iopub.execute_input":"2024-08-31T13:23:21.055819Z","iopub.status.idle":"2024-08-31T13:23:21.060571Z","shell.execute_reply.started":"2024-08-31T13:23:21.055786Z","shell.execute_reply":"2024-08-31T13:23:21.059777Z"},"trusted":true},"execution_count":8,"outputs":[{"name":"stdout","text":"<class 'keras_nlp.src.models.gemma.gemma_decoder_block.GemmaDecoderBlock'>\ndecoder_block_1/pre_attention_norm/scale          (3584,)         PartitionSpec(None,)\ndecoder_block_1/pre_attention_norm/scale          (3584,)         PartitionSpec(None,)\ndecoder_block_1/attention/query/kernel            (16, 3584, 256)  PartitionSpec('model', None, None)\ndecoder_block_1/attention/key/kernel              (8, 3584, 256)  PartitionSpec('model', None, None)\ndecoder_block_1/attention/value/kernel            (8, 3584, 256)  PartitionSpec('model', None, None)\ndecoder_block_1/attention/attention_output/kernel  (16, 256, 3584)  PartitionSpec('model', None, None)\ndecoder_block_1/pre_ffw_norm/scale                (3584,)         PartitionSpec(None,)\ndecoder_block_1/post_ffw_norm/scale               (3584,)         PartitionSpec(None,)\ndecoder_block_1/ffw_gating/kernel                 (3584, 14336)   PartitionSpec(None, 'model')\ndecoder_block_1/ffw_gating_2/kernel               (3584, 14336)   PartitionSpec(None, 'model')\ndecoder_block_1/ffw_linear/kernel                 (14336, 3584)   PartitionSpec('model', None)\n","output_type":"stream"}]},{"cell_type":"markdown","source":"## Inference before finetuning","metadata":{"id":"jc0ZzYIW0TSN","papermill":{"duration":0.005944,"end_time":"2024-02-21T09:58:46.139044","exception":false,"start_time":"2024-02-21T09:58:46.1331","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Let's try asking a model a question.","metadata":{}},{"cell_type":"code","source":"print(gemma_lm.generate(\"How can I plan a trip to Europe?\", max_length=512))","metadata":{"id":"ClaTyBp3Tgr4","outputId":"f2cf1ac7-469a-4c91-adff-74c2e3a1de89","papermill":{"duration":26.22201,"end_time":"2024-02-21T09:59:12.367352","exception":false,"start_time":"2024-02-21T09:58:46.145342","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:23:21.061698Z","iopub.execute_input":"2024-08-31T13:23:21.062004Z","iopub.status.idle":"2024-08-31T13:24:17.274743Z","shell.execute_reply.started":"2024-08-31T13:23:21.061976Z","shell.execute_reply":"2024-08-31T13:24:17.273445Z"},"trusted":true},"execution_count":9,"outputs":[{"name":"stdout","text":"How can I plan a trip to Europe?\n\n[User 0001]\n\nI'm planning a trip to Europe for the first time. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2 weeks. I'm going to be in Europe for 2\n","output_type":"stream"}]},{"cell_type":"markdown","source":"We are using the base Gemma 2 model, which means it has not been fine-tuned for any particular task. It has been trained to simply guess the next work on a vast amount of source documents.\n\nSuch a model is not yet a good fit for question answering. It will tend to continue predicting likely words, often continuing the question itself instead as if it was a random snippet of a random document on the web. It can easily get stuck in loops of high probability sequences.\n\nTo make it more useful, we can fine-tune on a question answering dataset. In this tutorial, we will use the Databricks Dolly dataset. This dataset contains 15,000 high-quality human-generated prompt / response pairs specifically designed for fine-tuning LLMs to follow instructions. Such fine-tuning is often called instruction fine-tuning, or IFT for short.","metadata":{}},{"cell_type":"markdown","source":"## Instruction fine-tuning","metadata":{"id":"IcPCXCwvXC7t","papermill":{"duration":0.006625,"end_time":"2024-02-21T09:59:12.393943","exception":false,"start_time":"2024-02-21T09:59:12.387318","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import json\ndata = []\nwith open('/kaggle/input/databricks-dolly-15k/databricks-dolly-15k.jsonl') as file:\n    for line in file:\n        features = json.loads(line)\n        # Filter out examples with context, to keep it simple.\n        if features[\"context\"]:\n            continue\n        # Format the entire example as a single string.\n        template = \"Instruction:\\n{instruction}\\n\\nResponse:\\n{response}\"\n        data.append(template.format(**features))\n\n# Truncate our data to speed up training.\ndata = data[:1000]","metadata":{"_kg_hide-output":true,"id":"6MVJlsuSXCcf","outputId":"fcfda7bb-899e-4606-88e1-42eed9e70fc0","papermill":{"duration":43.816972,"end_time":"2024-02-21T09:59:56.217135","exception":false,"start_time":"2024-02-21T09:59:12.400163","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:24:17.277804Z","iopub.execute_input":"2024-08-31T13:24:17.278167Z","iopub.status.idle":"2024-08-31T13:24:17.563233Z","shell.execute_reply.started":"2024-08-31T13:24:17.278133Z","shell.execute_reply":"2024-08-31T13:24:17.562149Z"},"trusted":true},"execution_count":10,"outputs":[]},{"cell_type":"markdown","source":"Let's look at a single training example.","metadata":{}},{"cell_type":"code","source":"data[0]","metadata":{"execution":{"iopub.status.busy":"2024-08-31T13:24:17.564708Z","iopub.execute_input":"2024-08-31T13:24:17.565004Z","iopub.status.idle":"2024-08-31T13:24:17.570579Z","shell.execute_reply.started":"2024-08-31T13:24:17.564977Z","shell.execute_reply":"2024-08-31T13:24:17.569645Z"},"trusted":true},"execution_count":11,"outputs":[{"execution_count":11,"output_type":"execute_result","data":{"text/plain":"'Instruction:\\nWhich is a species of fish? Tope or Rope\\n\\nResponse:\\nTope'"},"metadata":{}}]},{"cell_type":"markdown","source":"We will perform finetuning using [Low Rank Adaptation](https://arxiv.org/abs/2106.09685) (LoRA). LoRA is a fine-tuning technique which greatly reduces the number of trainable parameters for downstream tasks by freezing the full weights of the model and inserting a smaller number of new trainable weights into the model. Basically LoRA reparameterizes the larger full weight matrices by 2 smaller low-rank matrices AxB to train and this technique makes training much faster and more memory-efficient.","metadata":{"id":"xiLW0SpI1PfC","papermill":{"duration":0.032799,"end_time":"2024-02-21T09:59:56.345251","exception":false,"start_time":"2024-02-21T09:59:56.312452","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Enable LoRA for the model and set the LoRA rank to 4.\ngemma_lm.backbone.enable_lora(rank=4)","metadata":{"id":"3o_Gi3v_jp7s","papermill":{"duration":0.527031,"end_time":"2024-02-21T09:59:56.902475","exception":false,"start_time":"2024-02-21T09:59:56.375444","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:24:17.571749Z","iopub.execute_input":"2024-08-31T13:24:17.571981Z","iopub.status.idle":"2024-08-31T13:24:18.462940Z","shell.execute_reply.started":"2024-08-31T13:24:17.571957Z","shell.execute_reply":"2024-08-31T13:24:18.461830Z"},"trusted":true},"execution_count":12,"outputs":[]},{"cell_type":"code","source":"# Limit the input sequence length to 1024 to control memory usage.\ngemma_lm.preprocessor.sequence_length = 1024\ngemma_lm.compile(\n    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n    optimizer=keras.optimizers.Adam(learning_rate=5e-5),\n    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n)\ngemma_lm.summary()","metadata":{"_kg_hide-output":false,"id":"1-hQFy7hXWRl","outputId":"eb21b839-1e23-4be5-afcb-90a7bb0e4167","papermill":{"duration":362.718929,"end_time":"2024-02-21T10:05:59.651432","exception":false,"start_time":"2024-02-21T09:59:56.932503","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T13:24:18.464158Z","iopub.execute_input":"2024-08-31T13:24:18.464491Z","iopub.status.idle":"2024-08-31T13:24:18.626241Z","shell.execute_reply.started":"2024-08-31T13:24:18.464461Z","shell.execute_reply":"2024-08-31T13:24:18.625274Z"},"trusted":true},"execution_count":13,"outputs":[{"output_type":"display_data","data":{"text/plain":"\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m3584\u001b[0m)        │   \u001b[38;5;34m9,256,242,688\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m917,504,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">3584</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">9,256,242,688</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">917,504,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Total params: \u001b[0m\u001b[38;5;34m9,256,242,688\u001b[0m (34.48 GB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">9,256,242,688</span> (34.48 GB)\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m14,536,704\u001b[0m (55.45 MB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">14,536,704</span> (55.45 MB)\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m9,241,705,984\u001b[0m (34.43 GB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">9,241,705,984</span> (34.43 GB)\n</pre>\n"},"metadata":{}}]},{"cell_type":"markdown","source":"Note that enabling LoRA reduces the number of trainable parameters significantly, from 9 billion to only ~30 million.\n\nLet's fine-tune our model!","metadata":{"id":"CnpeavB4fZ7Y","papermill":{"duration":0.137841,"end_time":"2024-02-21T10:05:59.925998","exception":false,"start_time":"2024-02-21T10:05:59.788157","status":"completed"},"tags":[]}},{"cell_type":"code","source":"gemma_lm.fit(data, epochs=20, batch_size=4)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-08-31T13:24:18.627448Z","iopub.execute_input":"2024-08-31T13:24:18.627751Z","iopub.status.idle":"2024-08-31T14:53:27.215794Z","shell.execute_reply.started":"2024-08-31T13:24:18.627719Z","shell.execute_reply":"2024-08-31T14:53:27.214934Z"},"trusted":true},"execution_count":14,"outputs":[{"name":"stdout","text":"Epoch 1/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m326s\u001b[0m 1s/step - loss: 0.2340 - sparse_categorical_accuracy: 0.5391\nEpoch 2/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1907 - sparse_categorical_accuracy: 0.5839\nEpoch 3/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1860 - sparse_categorical_accuracy: 0.5892\nEpoch 4/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1823 - sparse_categorical_accuracy: 0.5955\nEpoch 5/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1779 - sparse_categorical_accuracy: 0.6024\nEpoch 6/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1719 - sparse_categorical_accuracy: 0.6126\nEpoch 7/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1525 - sparse_categorical_accuracy: 0.6472\nEpoch 9/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1396 - sparse_categorical_accuracy: 0.6744\nEpoch 10/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1252 - sparse_categorical_accuracy: 0.7047\nEpoch 11/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.1112 - sparse_categorical_accuracy: 0.7371\nEpoch 12/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0967 - sparse_categorical_accuracy: 0.7698\nEpoch 13/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0836 - sparse_categorical_accuracy: 0.8008\nEpoch 14/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0716 - sparse_categorical_accuracy: 0.8280\nEpoch 15/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0620 - sparse_categorical_accuracy: 0.8505\nEpoch 16/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0540 - sparse_categorical_accuracy: 0.8717\nEpoch 17/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0461 - sparse_categorical_accuracy: 0.8916\nEpoch 18/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0382 - sparse_categorical_accuracy: 0.9105\nEpoch 19/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0325 - sparse_categorical_accuracy: 0.9247\nEpoch 20/20\n\u001b[1m250/250\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m264s\u001b[0m 1s/step - loss: 0.0283 - sparse_categorical_accuracy: 0.9366\n","output_type":"stream"},{"execution_count":14,"output_type":"execute_result","data":{"text/plain":"<keras.src.callbacks.history.History at 0x7cab38328310>"},"metadata":{}}]},{"cell_type":"markdown","source":"## Inference after finetuning","metadata":{"id":"lBiOKlAy2MAe","papermill":{"duration":0.137064,"end_time":"2024-02-21T10:06:00.201917","exception":false,"start_time":"2024-02-21T10:06:00.064853","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Now that we have fine-tuned our model, we can try prompting with a question again. This time, we will use the specific format we used to combine our prompts and responses from the Databricks Dolly dataset.","metadata":{}},{"cell_type":"code","source":"print(gemma_lm.generate(\"Instruction:\\nHow can I plan a trip to Europe?\\n\\nResponse:\\n\", max_length=512))","metadata":{"id":"9yNyJ8CLXfw0","outputId":"f0a1c1e9-2221-4e3a-83df-829a47488b1c","papermill":{"duration":31.150786,"end_time":"2024-02-21T10:06:31.488004","exception":false,"start_time":"2024-02-21T10:06:00.337218","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-08-31T14:53:27.216840Z","iopub.execute_input":"2024-08-31T14:53:27.217109Z","iopub.status.idle":"2024-08-31T14:54:24.951587Z","shell.execute_reply.started":"2024-08-31T14:53:27.217082Z","shell.execute_reply":"2024-08-31T14:54:24.950576Z"},"trusted":true},"execution_count":15,"outputs":[{"name":"stdout","text":"Instruction:\nHow can I plan a trip to Europe?\n\nResponse:\nEurope has many different countries and within each country, there are many different cities. Each cities has different popular activities. For example, you can go hiking in Amsterdam, shopping in Paris, eat currywurst in Berlin or drink ouzo in Athens. First, you need to decide which countries you want to visit and within each country, which cities you want to go. Then you can look up online or ask a local for help on the specific cities. After that, you can book the airplane tickets and start planning your trip!\n","output_type":"stream"}]},{"cell_type":"markdown","source":"Much better! We could improve this model even more by fine-tuning with more data and tuning our learning rate and lora rank.\n\nAlternately, the Gemma models come with pre-instruction tuned checkpoints that can be used for question answering and a chat like experience out of the box. See [Gemma 2 inference using KerasNLP](https://www.kaggle.com/code/nilaychauhan/gemma-2-inference-using-kerasnlp) as an example.","metadata":{}},{"cell_type":"markdown","source":"# Save LoRa Weight","metadata":{}},{"cell_type":"code","source":"gemma_lm.backbone.save_lora_weights(\"LoRa_Gemma2_9b_en.lora.h5\")","metadata":{"execution":{"iopub.status.busy":"2024-08-31T14:59:57.786725Z","iopub.execute_input":"2024-08-31T14:59:57.787759Z","iopub.status.idle":"2024-08-31T14:59:58.057711Z","shell.execute_reply.started":"2024-08-31T14:59:57.787717Z","shell.execute_reply":"2024-08-31T14:59:58.056579Z"},"trusted":true},"execution_count":19,"outputs":[]}]}