fix checkpoint path

iankur · Sep 1, 2024 · 5ef7ad9 · 5ef7ad9
1 parent cc2f1a6
commit 5ef7ad9
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/recipes/run_vq_type_ablation.sh b/recipes/run_vq_type_ablation.sh
@@ -38,7 +38,7 @@ tune run recipes/eleuther_eval.py \
     --config recipes/config/eleuther_evaluation.yaml \
     $VQ_KEY=True $VQ_VALUE=False $FAST_QUANTIZER=True \
     $RESIDUAL_CODEBOOKS=$K $CODES=$C $CODE_DIM=$dhat \
-    $REORDER_CHANNEL=False checkpointer.output_dir=$CKPT_DIR \
+    $REORDER_CHANNEL=False checkpointer.checkpoint_dir=$CKPT_DIR \
     checkpointer.checkpoint_files=['meta_model_0.pt'] \
     $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="eval_key"
 
@@ -58,29 +58,30 @@ tune run recipes/eleuther_eval.py \
     --config recipes/config/eleuther_evaluation.yaml \
     $VQ_KEY=True $VQ_VALUE=False $FAST_QUANTIZER=True \
     $RESIDUAL_CODEBOOKS=$K $CODES=$C $CODE_DIM=$dhat \
-    $REORDER_CHANNEL=True checkpointer.output_dir=$CKPT_DIR \
+    $REORDER_CHANNEL=True checkpointer.checkpoint_dir=$CKPT_DIR \
     checkpointer.checkpoint_files=['meta_model_0.pt'] \
     $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="eval_key_reorder"
 
 # train and eval with only vq value
-CKPT_DIR="/home/ubuntu/vqllm/recipes/ckpts/key_reorder"
+# Note: reordering is not applicable for value
+CKPT_DIR="/home/ubuntu/vqllm/recipes/ckpts/value"
 
 # train
 tune run recipes/full_finetune_single_device.py \
     --config recipes/config/llama3_8b_single_device.yaml \
     $VQ_KEY=False $VQ_VALUE=True $FAST_QUANTIZER=True \
     $RESIDUAL_CODEBOOKS=$K $CODES=$C $CODE_DIM=$dhat \
     $REORDER_CHANNEL=True checkpointer.output_dir=$CKPT_DIR \
-    $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="train_key_reorder"
+    $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="train_value"
 
 # eval
 tune run recipes/eleuther_eval.py \
     --config recipes/config/eleuther_evaluation.yaml \
     $VQ_KEY=False $VQ_VALUE=True $FAST_QUANTIZER=True \
     $RESIDUAL_CODEBOOKS=$K $CODES=$C $CODE_DIM=$dhat \
-    $REORDER_CHANNEL=True checkpointer.output_dir=$CKPT_DIR \
+    $REORDER_CHANNEL=True checkpointer.checkpoint_dir=$CKPT_DIR \
     checkpointer.checkpoint_files=['meta_model_0.pt'] \
-    $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="eval_key_reorder"
+    $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="eval_value"
 
 # train and eval with vq for both key (with channel reordering) and value
 CKPT_DIR="/home/ubuntu/vqllm/recipes/ckpts/key_reorder_value"
@@ -98,6 +99,6 @@ tune run recipes/eleuther_eval.py \
     --config recipes/config/eleuther_evaluation.yaml \
     $VQ_KEY=True $VQ_VALUE=True $FAST_QUANTIZER=True \
     $RESIDUAL_CODEBOOKS=$K $CODES=$C $CODE_DIM=$dhat \
-    $REORDER_CHANNEL=True checkpointer.output_dir=$CKPT_DIR \
+    $REORDER_CHANNEL=True checkpointer.checkpoint_dir=$CKPT_DIR \
     checkpointer.checkpoint_files=['meta_model_0.pt'] \
     $WANDB_PROJECT $WANDB_GROUP $WANDB_NAME="eval_key_reorder_value"