1) Adding more comments at applying LoRA on Prefill params path.

2) Fixing model_ckpt_conversion.sh after refactoring and merging from main.

Author: aman2930

jetstream/core/orchestrator.py

@@ -666,6 +666,18 @@ def _prefill_thread(self, idx: int):
 
       adapter_id = request.adapter_id
 
+      # As prefill is happening one prompt at a time, for each prefill, we are
+      # applying the LoRA params on base to create a copy of params (equivalent
+      # to the size of base params) and use that for generating kv-cache. This
+      # copy is called the final_prefill_params, which is deleted soon after the
+      # generation of kv-cache.
+      # We can have memory-optimizations by updating the original copy of
+      # base params at the cost of extra computations to revert it back to original
+      # base params after kv-cache computation of each prompt, so that it can
+      # be used by the next prompt. But this optimization could also be tricky
+      # because as of now same params are being shared by prefill and generate,
+      # where generate always expect the base_params. So some race conditions need
+      # to be avoided.
       final_prefill_params = None
       if adapter_id == "":
         final_prefill_params = prefill_params

jetstream/tools/maxtext/model_ckpt_conversion.sh

@@ -38,10 +38,17 @@ export MODEL_BUCKET=$4
 # Point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you created, this bucket will store all the files generated by MaxText during a run, specifically the unscanned checkpoint.
 export BASE_OUTPUT_DIRECTORY=$5
 
-export LORA_LOCAL_PATH=$6
+export HUGGING_FACE_CHECKPOINT=$6
+
+export LORA_INPUT_ADAPTERS_PATH=$7
 
 export BUCKET_LOCATION=US
 
+if [[ -z "HUGGING_FACE_CHECKPOINT" ]]; then
+    echo "HUGGING_FACE_CHECKPOINT is required."
+    exit 1
+fi
+
 # Create three GCS buckets for the demo.
 gcloud storage buckets create ${MODEL_BUCKET} --location=${BUCKET_LOCATION} || true
 gcloud storage buckets create ${BASE_OUTPUT_DIRECTORY} --location=${BUCKET_LOCATION} || true
@@ -59,40 +66,56 @@ else
     # llama_or_mistral_ckpt.py requires local path, so we need to copy the checkpoint from CHKPT_BUCKET to local.
     tmp_ckpt_path="/tmp/"
     gcloud storage cp -r ${CHKPT_BUCKET} ${tmp_ckpt_path}
+
     path_parts=(${CHKPT_BUCKET//\// })
     directory_substring=${path_parts[-1]}
     CONVERT_CKPT_SCRIPT="llama_or_mistral_ckpt.py"
-    if [[ -x "${LORA_LOCAL_PATH}" ]]; then
+
+    if [[ ! -z "${LORA_INPUT_ADAPTERS_PATH}" ]]; then
+	lora_local_path="/tmp/"
+
+	if [[ "${LORA_INPUT_ADAPTERS_PATH}" =~ ^gs:// ]]; then
+	    path_parts=(${LORA_INPUT_ADAPTERS_PATH//\// })
+	    lora_dir_substring=${path_parts[-1]}
+
+	    lora_local_path="${tmp_ckpt_path}${lora_dir_substring}"
+	    if [[ ! -d ${lora_local_path} ]]; then
+		mkdir ${lora_local_path}
+	    fi
+	    gcloud storage cp -r ${LORA_INPUT_ADAPTERS_PATH} ${tmp_ckpt_path}
+	else
+	    lora_local_path=${LORA_INPUT_ADAPTERS_PATH}
+	fi
+
 	JAX_PLATFORMS=cpu python MaxText/${CONVERT_CKPT_SCRIPT} \
 	--base-model-path ${tmp_ckpt_path}${directory_substring} \
 	--maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
 	--model-size ${MODEL_NAME} \
-	--lora-config-path ${LORA_LOCAL_PATH}/adapter_config.json \
-	--lora-model-path ${LORA_LOCAL_PATH}/adapter_model.bin
+	--lora-input-adapters-path ${lora_local_path} \
+	--huggingface-checkpoint ${HUGGING_FACE_CHECKPOINT}
     else
 	JAX_PLATFORMS=cpu python MaxText/${CONVERT_CKPT_SCRIPT} \
 	--base-model-path ${tmp_ckpt_path}${directory_substring} \
 	--maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
-	--model-size ${MODEL_NAME}
+	--model-size ${MODEL_NAME} \
+	--huggingface-checkpoint ${HUGGING_FACE_CHECKPOINT}
     fi
 fi
 echo "Written MaxText compatible checkpoint to ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}"
 
 # We define `SCANNED_CKPT_PATH` to refer to the checkpoint subdirectory.
-# export SCANNED_CKPT_PATH=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}/0/items
 export SCANNED_CKPT_PATH=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}
 
 # Convert MaxText compatible checkpoints to unscanned checkpoints.
 # Note that the `SCANNED_CKPT_PATH` is in a `scanned` format which is great for training but for efficient decoding performance we want the checkpoint in an `unscanned` format.
 export RUN_NAME=${MODEL_NAME}_unscanned_chkpt_${idx}
 
-if [[ -x "${LORA_LOCAL_PATH}" ]]; then
+if [[ ! -z "${LORA_INPUT_ADAPTERS_PATH}" ]]; then
     JAX_PLATFORMS=cpu python MaxText/generate_param_only_checkpoint.py \
     MaxText/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIRECTORY} \
-    load_parameters_path=${SCANNED_CKPT_PATH}/base_weights/0/items \
-    lora_parameters_base_path=${SCANNED_CKPT_PATH}/lora_weights/0/items \
-    lora_config_path=${LORA_LOCAL_PATH}/adapter_config.json \
+    load_parameters_path=${SCANNED_CKPT_PATH}/base/0/items \
+    lora_input_adapters_path=${SCANNED_CKPT_PATH}/LoRAs \
     run_name=${RUN_NAME} \
     model_name=${MODEL_NAME} \
     force_unroll=true
@@ -101,7 +124,7 @@ else
     JAX_PLATFORMS=cpu python MaxText/generate_param_only_checkpoint.py \
     MaxText/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIRECTORY} \
-    load_parameters_path=${SCANNED_CKPT_PATH}/base_weights/0/items \
+    load_parameters_path=${SCANNED_CKPT_PATH}/0/items \
     run_name=${RUN_NAME} \
     model_name=${MODEL_NAME} \
     force_unroll=true