|
| 1 | +<!-- |
| 2 | + Copyright 2023–2026 Google LLC |
| 3 | +
|
| 4 | + Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + you may not use this file except in compliance with the License. |
| 6 | + You may obtain a copy of the License at |
| 7 | +
|
| 8 | + https://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | + Unless required by applicable law or agreed to in writing, software |
| 11 | + distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + See the License for the specific language governing permissions and |
| 14 | + limitations under the License. |
| 15 | + --> |
| 16 | + |
| 17 | +# LoRA Fine-tuning on single-host TPUs |
| 18 | + |
| 19 | +**Low-Rank Adaptation (LoRA)** is a Parameter-Efficient Fine-Tuning (PEFT) technique designed to optimize large language models while minimizing resource consumption. |
| 20 | + |
| 21 | +Unlike traditional full-parameter fine-tuning, LoRA: |
| 22 | + |
| 23 | +- **Freezes the pre-trained model weights**, preserving the original knowledge. |
| 24 | +- **Injects trainable rank decomposition matrices** into the Transformer layers. |
| 25 | + |
| 26 | +This approach **greatly reduces the number of trainable parameters** required for downstream tasks, making the process faster and more memory-efficient. |
| 27 | + |
| 28 | +This tutorial provides step-by-step instructions for setting up the environment and performing LoRA fine-tuning on a Hugging Face dataset using MaxText. |
| 29 | + |
| 30 | +We use [Tunix](https://github.com/google/tunix), a JAX-based library, to power these post-training tasks. |
| 31 | + |
| 32 | +In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get started! |
| 33 | + |
| 34 | +## Setup environment variables |
| 35 | + |
| 36 | +Login to Hugging Face. Provide your access token when prompted: |
| 37 | + |
| 38 | +```bash |
| 39 | +hf auth login |
| 40 | +``` |
| 41 | + |
| 42 | +Set the following environment variables before running LoRA Fine-tuning. |
| 43 | + |
| 44 | +```sh |
| 45 | +# -- Model configuration -- |
| 46 | +export MODEL_NAME=<MODEL_NAME> # e.g., 'gemma3-4b' |
| 47 | + |
| 48 | +# -- MaxText configuration -- |
| 49 | +export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # e.g., gs://my-bucket/my-output-directory or /path/to/my-output-directory |
| 50 | +export RUN_NAME=<RUN_NAME> # e.g., $(date +%Y-%m-%d-%H-%M-%S) |
| 51 | +export STEPS=<STEPS> # e.g., 1000 |
| 52 | +export PER_DEVICE_BATCH_SIZE=<BATCH_SIZE_PER_DEVICE> # e.g., 1 |
| 53 | +export LORA_RANK=<LORA_RANK> # e.g., 16 |
| 54 | +export LORA_ALPHA=<LORA_ALPHA> # e.g., 32.0 |
| 55 | +export LEARNING_RATE=<LEARNING_RATE> # e.g., 3e-6 |
| 56 | +export MAX_TARGET_LENGTH=<MAX_TARGET_LENGTH> # e.g., 1024 |
| 57 | + |
| 58 | +# -- Dataset configuration -- |
| 59 | +export DATASET_NAME=<DATASET_NAME> # e.g., openai/gsm8k |
| 60 | +export TRAIN_SPLIT=<TRAIN_SPLIT> # e.g., train |
| 61 | +export HF_DATA_DIR=<DATASET_PATH> # e.g., main |
| 62 | +export TRAIN_DATA_COLUMNS=<DATA_COLUMNS> # e.g., ['question','answer'] |
| 63 | +export CHAT_TEMPLATE_PATH=<TEMPLATE_PATH> # e.g., maxtext/examples/chat_templates/math_qa.json |
| 64 | + |
| 65 | +# -- LoRA Conversion configuration (Optional) -- |
| 66 | +export HF_LORA_ADAPTER_PATH=<HF_LORA_ADAPTER_PATH> # e.g., 'username/adapter-name' |
| 67 | +``` |
| 68 | + |
| 69 | +## Customizing Trainable Layers (Optional) |
| 70 | + |
| 71 | +By default, MaxText determines which layers to apply LoRA to based on the model's architecture by reading `src/maxtext/configs/post_train/lora_module_path.yml`. |
| 72 | + |
| 73 | +If you need to fine-tune specific components (e.g., targeting only Attention layers to optimize memory usage), you can override these defaults through the following hierarchy: |
| 74 | + |
| 75 | +### Configuration Hierarchy |
| 76 | + |
| 77 | +1. **Command Line Argument**: Pass the `lora_module_path` argument directly in your training command. This is the most flexible way for experimental iterations. |
| 78 | +2. **Task-Specific Config (`sft.yml`)**: Define the `lora_module_path` parameter in `src/maxtext/configs/post_train/sft.yml` to set a persistent configuration for your SFT runs. |
| 79 | +3. **Global Defaults**: Automatic detection via the model-to-regex mapping defined in `lora_module_path.yml`. |
| 80 | + |
| 81 | +## Get your model checkpoint |
| 82 | + |
| 83 | +This section explains how to prepare your model checkpoint for use with MaxText. You have two options: using an existing MaxText checkpoint or converting a Hugging Face checkpoint. |
| 84 | + |
| 85 | +### Option 1: Using an existing MaxText checkpoint |
| 86 | + |
| 87 | +If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section. |
| 88 | + |
| 89 | +```sh |
| 90 | +export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items or /path/to/my-model-checkpoint/0/items |
| 91 | +``` |
| 92 | + |
| 93 | +### Option 2: Converting a Hugging Face checkpoint |
| 94 | + |
| 95 | +Refer to the steps in [Hugging Face to MaxText](https://maxtext.readthedocs.io/en/maxtext-v0.2.1/guides/checkpointing_solutions/convert_checkpoint.html#hugging-face-to-maxtext) to convert a hugging face checkpoint to MaxText. Make sure you have the correct checkpoint files converted and saved. Similar as Option 1, you can set the following environment and move on. |
| 96 | + |
| 97 | +```sh |
| 98 | +export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items or /path/to/my-model-checkpoint/0/items |
| 99 | +``` |
| 100 | + |
| 101 | +## Run a Fresh LoRA Fine-Tuning on Hugging Face Dataset |
| 102 | + |
| 103 | +Once your environment variables and checkpoints are ready, you can start the LoRA fine-tuning process. |
| 104 | + |
| 105 | +Execute the following command to begin training: |
| 106 | + |
| 107 | +```sh |
| 108 | +python3 -m maxtext.trainers.post_train.sft.train_sft \ |
| 109 | + run_name="${RUN_NAME?}" \ |
| 110 | + base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \ |
| 111 | + model_name="${MODEL_NAME?}" \ |
| 112 | + load_parameters_path="${MAXTEXT_CKPT_PATH?}" \ |
| 113 | + hf_path="${DATASET_NAME?}" \ |
| 114 | + train_split="${TRAIN_SPLIT?}" \ |
| 115 | + hf_data_dir="${HF_DATA_DIR?}" \ |
| 116 | + train_data_columns="${TRAIN_DATA_COLUMNS?}" \ |
| 117 | + steps="${STEPS?}" \ |
| 118 | + per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \ |
| 119 | + max_target_length="${MAX_TARGET_LENGTH?}" \ |
| 120 | + learning_rate="${LEARNING_RATE?}" \ |
| 121 | + chat_template_path="${CHAT_TEMPLATE_PATH?}" \ |
| 122 | + enable_nnx=True \ |
| 123 | + pure_nnx_decoder=True \ |
| 124 | + lora.enable_lora=True \ |
| 125 | + lora.lora_rank="${LORA_RANK?}" \ |
| 126 | + lora.lora_alpha="${LORA_ALPHA?}" |
| 127 | +``` |
| 128 | + |
| 129 | +Your fine-tuned model checkpoints will be saved here: `$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints`. |
| 130 | + |
| 131 | +## (Optional) Resume from a previous LoRA checkpoint |
| 132 | + |
| 133 | +If you want to resume training from a previous run or further fine-tune an existing LoRA adapter, you can specify the LoRA checkpoint path. |
| 134 | + |
| 135 | +### Step 1: Convert HF LoRA adapter to MaxText format |
| 136 | + |
| 137 | +If your LoRA adapter is currently in Hugging Face format, you must convert it to MaxText format before it can be loaded. Use the integrated conversion utility: |
| 138 | + |
| 139 | +```sh |
| 140 | +python3 -m maxtext.checkpoint_conversion.to_maxtext \ |
| 141 | + model_name="${MODEL_NAME?}" \ |
| 142 | + hf_lora_adapter_path="${HF_LORA_ADAPTER_PATH?}" \ |
| 143 | + base_output_directory="${BASE_OUTPUT_DIRECTORY?}/converted_adapter" \ |
| 144 | + hardware=cpu skip_jax_distributed_system=True |
| 145 | +``` |
| 146 | + |
| 147 | +### Step 2: Set the restore path |
| 148 | + |
| 149 | +Point `LORA_RESTORE_PATH` to the converted MaxText adapter directory (the directory containing the `0/items` or Orbax files). |
| 150 | + |
| 151 | +- **load_parameters_path**: Points to the frozen base model weights (the original model). |
| 152 | +- **lora_restore_path**: Points to the previous LoRA adapter weights you wish to load. |
| 153 | + |
| 154 | +```sh |
| 155 | +export LORA_RESTORE_PATH=<LORA_RESTORE_PATH> # e.g., gs://my-bucket/run-1/checkpoints/0/items or /path/to/run-1/checkpoints/0/items |
| 156 | +``` |
| 157 | + |
| 158 | +### Step 3: Run LoRA Fine-Tuning with the Restore Path |
| 159 | + |
| 160 | +Once your environment variables and checkpoints are ready, you can start the LoRA fine-tuning process. |
| 161 | + |
| 162 | +Execute the following command to begin training: |
| 163 | + |
| 164 | +```sh |
| 165 | +python3 -m maxtext.trainers.post_train.sft.train_sft \ |
| 166 | + run_name="${RUN_NAME?}" \ |
| 167 | + base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \ |
| 168 | + model_name="${MODEL_NAME?}" \ |
| 169 | + load_parameters_path="${MAXTEXT_CKPT_PATH?}" \ |
| 170 | + lora.lora_restore_path="${LORA_RESTORE_PATH?}" \ |
| 171 | + hf_path="${DATASET_NAME?}" \ |
| 172 | + train_split="${TRAIN_SPLIT?}" \ |
| 173 | + hf_data_dir="${HF_DATA_DIR?}" \ |
| 174 | + train_data_columns="${TRAIN_DATA_COLUMNS?}" \ |
| 175 | + steps="${STEPS?}" \ |
| 176 | + per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \ |
| 177 | + max_target_length="${MAX_TARGET_LENGTH?}" \ |
| 178 | + learning_rate="${LEARNING_RATE?}" \ |
| 179 | + chat_template_path="${CHAT_TEMPLATE_PATH?}" \ |
| 180 | + enable_nnx=True \ |
| 181 | + pure_nnx_decoder=True \ |
| 182 | + lora.enable_lora=True \ |
| 183 | + lora.lora_rank="${LORA_RANK?}" \ |
| 184 | + lora.lora_alpha="${LORA_ALPHA?}" |
| 185 | +``` |
| 186 | + |
| 187 | +Your fine-tuned model checkpoints will be saved here: `$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints`. |
| 188 | + |
| 189 | +## (Optional) Convert Fine-tuned LoRA to Hugging Face Format |
| 190 | + |
| 191 | +After completing the fine-tuning process, your LoRA weights are stored in MaxText/Orbax format. To use these weights with the Hugging Face ecosystem (e.g., for inference or sharing), convert them back using the `to_huggingface.py` script. |
| 192 | + |
| 193 | +```sh |
| 194 | +python3 -m maxtext.checkpoint_conversion.to_huggingface \ |
| 195 | + model_name="${MODEL_NAME?}" \ |
| 196 | + lora.lora_restore_path="${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/checkpoints/<STEPS>/model_params" \ |
| 197 | + base_output_directory="${BASE_OUTPUT_DIRECTORY?}/hf_lora_adapter" |
| 198 | +``` |
| 199 | + |
| 200 | +- `lora.lora_restore_path`: Point this to the specific checkpoint directory (e.g., `.../checkpoints/1000/items`) that you want to export. |
| 201 | +- `base_output_directory`: The local or GCS directory where the Hugging Face `adapter_model.safetensors` and `adapter_config.json` will be saved. |
| 202 | +- `lora.lora_rank` / `lora.lora_alpha`: Must match the values used during the training phase to ensure the `adapter_config.json` is generated correctly. |
0 commit comments