Merge pull request #2349 from AI-Hypercomputer:sft_e2e_update

PiperOrigin-RevId: 807416863

Author: Google-ML-Automation

docs/tutorials/sft.md

@@ -20,7 +20,7 @@ Supervised fine-tuning (SFT) is a process where a pre-trained large language mod
 This tutorial demonstrates step-by-step instructions for setting up the environment and then training the Llama3.1 8B model on the [HuggingFaceH4/ultrachat_200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) dataset using SFT. If you wish to use a different dataset, you can [update the dataset configurations](https://github.com/AI-Hypercomputer/maxtext/blob/main/MaxText/configs/sft.yml).
 
 We use [Tunix](https://github.com/google/tunix), a JAX-based library designed for post-training tasks, to perform SFT.
- 
+
 In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get started!
 
 ## Setup virtual environment
@@ -44,18 +44,17 @@ bash setup.sh
 
 ## Run SFT
 There are two scenarios supported for running SFT:
-1. **Run SFT on Hugging Face checkpoint**  
+1. **Run SFT on Hugging Face checkpoint**
     Download the checkpoint directly from Hugging Face and fine-tune it using SFT.
 
-2. **Run SFT on MaxText checkpoint**  
+2. **Run SFT on MaxText checkpoint**
     Use a checkpoint generated by MaxText and fine-tune it using SFT.
 
 Choose the scenario that matches your workflow and follow the corresponding instructions below.
 
 ### Run SFT on Hugging Face checkpoint
 * The script will first convert a Hugging Face checkpoint to a MaxText checkpoint.
 * It then runs SFT on this converted checkpoint.
-* After fine-tuning, the script converts the resulting checkpoint back to the Hugging Face format.
 
 #### Setup environment variables
 ```
@@ -75,7 +74,6 @@ bash ~/maxtext/end_to_end/tpu/llama3.1/8b/run_sft.sh
 
 ### Run SFT on MaxText checkpoint
 * The script directly runs SFT on MaxText checkpoint.
-* After fine-tuning, the script converts the resulting checkpoint back to the Hugging Face format.
 
 #### Setup environment variables
 ```

end_to_end/tpu/llama3.1/8b/run_sft.sh

@@ -78,13 +78,3 @@ sorted_dirs=($(printf '%s\n' "${integer_dirs[@]}" | sort -n))
 largest_dir="${sorted_dirs[-1]}"
 FINE_TUNED_MODEL_CKPT_PATH=${CHECKPOINTS_PATH}/${largest_dir}
 echo "Fine-tuned model checkpoint: ${FINE_TUNED_MODEL_CKPT_PATH}"
-
-# Convert the fine-tuned MaxText checkpoint to Hugging Face checkpoint
-export LOCAL_PATH=./tmp/hf/${PRE_TRAINED_MODEL}/${RUN_NAME}
-python3 -m MaxText.utils.ckpt_conversion.to_huggingface "${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText}"/configs/base.yml \
-    model_name=${PRE_TRAINED_MODEL} \
-    hf_access_token=${HF_TOKEN} \
-    load_parameters_path=${FINE_TUNED_MODEL_CKPT_PATH}/model_params \
-    base_output_directory=${LOCAL_PATH} \
-    scan_layers=False
-echo "Converted Hugging Face checkpoint saved to: ${LOCAL_PATH}"