#!/bin/bash

NNODES=1
NPROC_PER_NODE=4
LR=1e-5
NUM_EPOCHS=1
BATCH_SIZE_TRAINING=2
MODEL_NAME="/xxx/models--meta-llama--Llama-3.2-11B-Vision-Instruct/snapshots/075e8feb24b6a50981f6fdc161622f741a8760b1"
DIST_CHECKPOINT_ROOT_FOLDER="./finetuned_model"
DIST_CHECKPOINT_FOLDER="fine-tuned"
DATASET="custom_dataset"
CUSTOM_DATASET_TEST_SPLIT="test"
CUSTOM_DATASET_FILE="recipes/quickstart/finetuning/datasets/xxx_dataset.py"
RUN_VALIDATION=True
BATCHING_STRATEGY="padding"
OUTPUT_DIR="finetune/output"

torchrun --master_port 12412 \
         --nnodes $NNODES \
         --nproc_per_node $NPROC_PER_NODE \
         recipes/quickstart/finetuning/finetuning.py \
         --enable_fsdp \
         --lr $LR \
         --num_epochs $NUM_EPOCHS \
         --batch_size_training $BATCH_SIZE_TRAINING \
         --model_name $MODEL_NAME \
         --dist_checkpoint_root_folder $DIST_CHECKPOINT_ROOT_FOLDER \
         --dist_checkpoint_folder $DIST_CHECKPOINT_FOLDER \
         --use_fast_kernels \
         --dataset $DATASET \
         --custom_dataset.test_split $CUSTOM_DATASET_TEST_SPLIT \
         --custom_dataset.file $CUSTOM_DATASET_FILE \
         --run_validation $RUN_VALIDATION \
         --batching_strategy $BATCHING_STRATEGY \
         --output_dir $OUTPUT_DIR

ls
__0_0.distcp  __1_0.distcp  __2_0.distcp  __3_0.distcp  train_params.yaml