Skip to content

Commit fb75534

Browse files
sirakiincopybara-github
sirakiin
authored and
copybara-github
committed
Add Gemma3-270M model export example colab.
PiperOrigin-RevId: 810140711
1 parent b56f946 commit fb75534

File tree

1 file changed

+233
-0
lines changed

1 file changed

+233
-0
lines changed

‎ai_edge_torch/generative/colabs/Gemma3_270M_convertion.ipynb‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {
6+
"id": "1wRAjEfxE-rV"
7+
},
8+
"source": [
9+
"##### Copyright 2025 The AI Edge Torch Authors."
10+
]
11+
},
12+
{
13+
"cell_type": "code",
14+
"execution_count": null,
15+
"metadata": {
16+
"cellView": "form",
17+
"id": "qG29JvSAGKht"
18+
},
19+
"outputs": [],
20+
"source": [
21+
"#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
22+
"# you may not use this file except in compliance with the License.\n",
23+
"# You may obtain a copy of the License at\n",
24+
"#\n",
25+
"# https://www.apache.org/licenses/LICENSE-2.0\n",
26+
"#\n",
27+
"# Unless required by applicable law or agreed to in writing, software\n",
28+
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
29+
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
30+
"# See the License for the specific language governing permissions and\n",
31+
"# limitations under the License."
32+
]
33+
},
34+
{
35+
"cell_type": "markdown",
36+
"metadata": {
37+
"id": "O4F5iVzh--fg"
38+
},
39+
"source": [
40+
"# Exporting Gemma3 270M with AI Edge Torch\n",
41+
"\n",
42+
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-ai-edge/ai_edge_torch/generative/colabs/Gemma3_270M_convertion.ipynb)\n",
43+
"\n",
44+
"In this colab, we will show you how to export a Gemma-3-270M model to LiteRT-LM format with AI Edge Torch.\n",
45+
"\n",
46+
"It works with the base Gemma-3-270M-it model and its fine-tuned models. For later, checkout the [Full Model Fine-Tune using Hugging Face Transformers](https://ai.google.dev/gemma/docs/core/huggingface_text_full_finetune) tutorial.\n",
47+
"\n",
48+
"#Prerequisite for exporting google/gemma-3-270m-it\n",
49+
"\n",
50+
"- Create HuggingFace token with permission access to\n",
51+
" - google/gemma-3-270m-it\n",
52+
"\n",
53+
" This is needed to download the checkpoint and tokenizer.\n",
54+
"\n",
55+
"- Open Colab Secrets: In your Google Colab notebook, locate the Secrets icon in the left-hand sidebar and click on it.\n",
56+
"- Add a new secret: Click the \"Add Secret\" button.\n",
57+
"- Name your secret: Enter \"HF_TOKEN\" for your token in the \"Name\" field.\n",
58+
"- Paste your token: In the \"Value\" field, paste the actual token you want to store.\n",
59+
"\n",
60+
"#Prerequisite for exporting fine-tuned model\n",
61+
"\n",
62+
"- Access to the finetuned repo in Hugging Face Hub, or\n",
63+
"\n",
64+
"- Access to the finetuned checkpoint\n",
65+
"\n",
66+
"\n",
67+
"## Note: When running notebooks in this repository with Google Colab, some users may see the following warning message:\n",
68+
"\n",
69+
"![Colab warning](https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/data/colab_warning.jpg?raw=true)\n",
70+
"\n",
71+
"Please click `Restart Session` and run again.\n",
72+
"\n",
73+
"\n",
74+
"This colab works with a free tier colab runtime.\n"
75+
]
76+
},
77+
{
78+
"cell_type": "code",
79+
"execution_count": null,
80+
"metadata": {
81+
"id": "Stdvqj8A-5sj"
82+
},
83+
"outputs": [],
84+
"source": [
85+
"# @title Install dependencies and environment setup\n",
86+
"\n",
87+
"!pip install ai-edge-litert-nightly==2.0.2.dev20250917\n",
88+
"!pip uninstall -y tensorflow\n",
89+
"!pip install ai-edge-torch-nightly==0.7.0.dev20250920 --force-reinstall"
90+
]
91+
},
92+
{
93+
"cell_type": "code",
94+
"execution_count": null,
95+
"metadata": {
96+
"id": "H-XUP1wA_oT5"
97+
},
98+
"outputs": [],
99+
"source": [
100+
"# Setup Hugging Face Hub credentials\n",
101+
"\n",
102+
"import os\n",
103+
"from google.colab import userdata\n",
104+
"os.environ[\"HF_TOKEN\"] = userdata.get('HF_TOKEN')"
105+
]
106+
},
107+
{
108+
"cell_type": "code",
109+
"execution_count": null,
110+
"metadata": {
111+
"id": "5gT64EZu_28k"
112+
},
113+
"outputs": [],
114+
"source": [
115+
"# @title Import needed packages.\n",
116+
"from huggingface_hub import snapshot_download\n",
117+
"from ai_edge_torch.generative.examples.gemma3 import gemma3\n",
118+
"from ai_edge_torch.generative.utilities import converter\n",
119+
"from ai_edge_torch.generative.utilities.export_config import ExportConfig\n",
120+
"from ai_edge_torch.generative.layers import kv_cache"
121+
]
122+
},
123+
{
124+
"cell_type": "markdown",
125+
"metadata": {
126+
"id": "7T3G3Uj6MVk0"
127+
},
128+
"source": [
129+
"# Exporting the checkpoint to LiteRT-LM format.\n",
130+
"\n",
131+
"In this example, we directly use the google/gemma-3-270m-it repo. But you can also replace it with your fine-tuned model directory or repo ID.\n",
132+
"\n",
133+
"If you are following the fine-tune colab and storing your checkpoint to Google Drive as the default setup, you can point to the checkpoint with the followings instead of downloading the base checkpoint.\n",
134+
"\n",
135+
"```\n",
136+
"from google.colab import drive\n",
137+
"drive.mount('/content/drive')\n",
138+
"checkpoint_dir = '/content/drive/MyDrive/MyGemmaNPC'\n",
139+
"```\n"
140+
]
141+
},
142+
{
143+
"cell_type": "code",
144+
"execution_count": null,
145+
"metadata": {
146+
"id": "l8DSwaSq_8Er"
147+
},
148+
"outputs": [],
149+
"source": [
150+
"# @title Download checkpoint\n",
151+
"\n",
152+
"checkpoint_dir = snapshot_download('google/gemma-3-270m-it')\n"
153+
]
154+
},
155+
{
156+
"cell_type": "markdown",
157+
"metadata": {
158+
"id": "trGuvI-bAM1j"
159+
},
160+
"source": [
161+
"# Convert to LiteRT-LM format\n",
162+
"\n",
163+
"After the following cell, you will be able to download the exported `.litertlm` file under `/contents/`, which will be accessible from the `Files` pannel."
164+
]
165+
},
166+
{
167+
"cell_type": "code",
168+
"execution_count": null,
169+
"metadata": {
170+
"id": "waCdEPbcADAm"
171+
},
172+
"outputs": [],
173+
"source": [
174+
"output_dir = '/content/'\n",
175+
"\n",
176+
"# Import the weights and build the PyTorch model\n",
177+
"pytorch_model = gemma3.build_model_270m(checkpoint_dir)\n",
178+
"\n",
179+
"# Setup the export configurations and parameters for text generation models.\n",
180+
"export_config = ExportConfig()\n",
181+
"export_config.kvcache_layout = kv_cache.KV_LAYOUT_TRANSPOSED\n",
182+
"export_config.mask_as_input = True\n",
183+
"\n",
184+
"# Configs specific for text generation models.\n",
185+
"litertlm_config = {\n",
186+
" \"tokenizer_model_path\": os.path.join(checkpoint_dir, 'tokenizer.model'),\n",
187+
" \"start_token_id\": 2, # \"\u003cbos\u003e\"\n",
188+
" \"stop_token_ids\": [1, 106], # [\"\u003ceos\u003e\", \"\u003cend_of_turn\u003e\"]\n",
189+
" \"prompt_prefix\": \"\u003cstart_of_turn\u003euser\\n\",\n",
190+
" \"prompt_suffix\": \"\u003cend_of_turn\u003e\\n\u003cstart_of_turn\u003emodel\\n\",\n",
191+
" \"model_prompt_prefix\": \"\u003cstart_of_turn\u003emodel\\n\",\n",
192+
" \"model_prompt_suffix\": \"\u003cend_of_turn\u003e\\n\",\n",
193+
" \"user_prompt_prefix\": \"\u003cstart_of_turn\u003euser\\n\",\n",
194+
" \"user_prompt_suffix\": \"\u003cend_of_turn\u003e\\n\",\n",
195+
" \"output_format\": \"litertlm\",\n",
196+
"}\n",
197+
"\n",
198+
"# Convert to LiteRT or LiteRT-LM Format\n",
199+
"converter.convert_to_litert(\n",
200+
" pytorch_model,\n",
201+
" output_path=output_dir,\n",
202+
" output_name_prefix=\"gemma\",\n",
203+
" prefill_seq_len=2048,\n",
204+
" kv_cache_max_len=4096,\n",
205+
" quantize=\"dynamic_int8\",\n",
206+
" export_config=export_config,\n",
207+
" **litertlm_config\n",
208+
")\n"
209+
]
210+
}
211+
],
212+
"metadata": {
213+
"colab": {
214+
"private_outputs": true,
215+
"provenance": [
216+
{
217+
"file_id": "1P33SZyxx2s_k8INd5dYEnKsBaJVTTGoF",
218+
"timestamp": 1758487391162
219+
}
220+
],
221+
"toc_visible": true
222+
},
223+
"kernelspec": {
224+
"display_name": "Python 3",
225+
"name": "python3"
226+
},
227+
"language_info": {
228+
"name": "python"
229+
}
230+
},
231+
"nbformat": 4,
232+
"nbformat_minor": 0
233+
}

0 commit comments

Comments
 (0)