From 034b37809d0901fcf1162d0c8cf16f476b6e7960 Mon Sep 17 00:00:00 2001 From: Emma Lien Date: Wed, 8 Apr 2026 02:03:25 +0000 Subject: [PATCH 1/2] docs: add LoRA tutorial --- docs/tutorials/posttraining/lora.md | 128 ++++++++++++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 docs/tutorials/posttraining/lora.md diff --git a/docs/tutorials/posttraining/lora.md b/docs/tutorials/posttraining/lora.md new file mode 100644 index 0000000000..adacff227a --- /dev/null +++ b/docs/tutorials/posttraining/lora.md @@ -0,0 +1,128 @@ + + +# LoRA Fine-tuning on single-host TPUs + +**Low-Rank Adaptation (LoRA)** is a Parameter-Efficient Fine-Tuning (PEFT) technique designed to optimize large language models while minimizing resource consumption. + +Unlike traditional full-parameter fine-tuning, LoRA: +* **Freezes the pre-trained model weights**, preserving the original knowledge. +* **Injects trainable rank decomposition matrices** into the Transformer layers. + +This approach **greatly reduces the number of trainable parameters** required for downstream tasks, making the process faster and more memory-efficient. + +This tutorial provides step-by-step instructions for setting up the environment and performing LoRA fine-tuning on a Hugging Face dataset using MaxText. + +We use [Tunix](https://github.com/google/tunix), a JAX-based library, to power these post-training tasks. + +In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get started! + +## Install MaxText and Post-Training dependencies + +For instructions on installing MaxText with post-training dependencies on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) and use the `maxtext[tpu-post-train]` installation path to include all necessary post-training dependencies. + +## Setup environment variables + +Set the following environment variables before running LoRA Fine-tuning. + +```sh +# -- Model configuration -- +export MODEL_NAME= # e.g., 'llama3.1-8b-Instruct' + +# -- MaxText configuration -- +export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/my-output-directory +export RUN_NAME= # e.g., $(date +%Y-%m-%d-%H-%M-%S) +export STEPS= # e.g., 1000 +export PER_DEVICE_BATCH_SIZE= # e.g., 1 +export HF_TOKEN= +export LORA_RANK= # e.g., 16 +export LORA_ALPHA= # e.g., 32.0 +export LEARNING_RATE= # e.g., 3e-6 +export MAX_TARGET_LENGTH= # e.g., 1024 +export WEIGHT_DTYPE= # e.g., bfloat16 +export DTYPE= # e.g., bfloat16 + +# -- Dataset configuration -- +export DATASET_NAME= # e.g., openai/gsm8k +export TRAIN_SPLIT= # e.g., train +export HF_DATA_DIR= # e.g., main +export TRAIN_DATA_COLUMNS= # e.g., ['question','answer'] +``` + +## Get your model checkpoint + +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. + +### Option 1: Using an existing MaxText checkpoint + +If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section. + +```sh +export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/0/items +``` + +### Option 2: Converting a Hugging Face checkpoint + +Refer 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 correct checkpoint files converted and saved. Similar as Option 1, you can set the following environment and move on. + +```sh +export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/0/items +``` +## (Optional) Resume from a previous LoRA checkpoint + +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. +- **load_parameters_path**: Points to the frozen base model weights (the original model). +- **lora_restore_path**: Points to the previous LoRA adapter weights you wish to load. + +```sh +# If starting fresh, you can leave this empty or skip this variable +export LORA_RESTORE_PATH= # e.g., gs://my-bucket/run-1/checkpoints/0/items +``` +## Run LoRA Fine-Tuning on Hugging Face Dataset + +Once your environment variables and checkpoints are ready, you can start the LoRA fine-tuning process. + +Execute the following command to begin training: + +```sh +python3 -m maxtext.trainers.post_train.sft.train_sft \ + maxtext/configs/post_train/sft.yml \ + run_name="${RUN_NAME}" \ + base_output_directory="${BASE_OUTPUT_DIRECTORY}" \ + model_name="${MODEL_NAME}" \ + load_parameters_path="${MAXTEXT_CKPT_PATH}" \ + lora_restore_path="${LORA_RESTORE_PATH}" \ + hf_access_token="${HF_TOKEN}" \ + hf_path="${DATASET_NAME}" \ + train_split="${TRAIN_SPLIT}" \ + hf_data_dir="${HF_DATA_DIR}" \ + train_data_columns="${TRAIN_DATA_COLUMNS}" \ + steps="${STEPS}" \ + per_device_batch_size="${PER_DEVICE_BATCH_SIZE}" \ + max_target_length="${MAX_TARGET_LENGTH}" \ + learning_rate="${LEARNING_RATE}" \ + weight_dtype="${WEIGHT_DTYPE}" \ + dtype="${DTYPE}" \ + profiler=xplane \ + enable_nnx=True \ + pure_nnx_decoder=True \ + enable_lora=True \ + lora_rank="${LORA_RANK}" \ + lora_alpha="${LORA_ALPHA}" \ + scan_layers=True +``` + +Your fine-tuned model checkpoints will be saved here: `$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints`. \ No newline at end of file From 473f9cf9358eee3530086ebe6f655a6e2938b8bb Mon Sep 17 00:00:00 2001 From: Emma Lien Date: Thu, 9 Apr 2026 09:59:23 +0000 Subject: [PATCH 2/2] fix some naming and add the conversion instruction --- docs/tutorials/posttraining/lora.md | 86 ++++++++++++++++++++--------- 1 file changed, 61 insertions(+), 25 deletions(-) diff --git a/docs/tutorials/posttraining/lora.md b/docs/tutorials/posttraining/lora.md index adacff227a..3306e8748d 100644 --- a/docs/tutorials/posttraining/lora.md +++ b/docs/tutorials/posttraining/lora.md @@ -1,5 +1,5 @@