#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091
#SBATCH -c 20
#SBATCH -w anon
#SBATCH --gres=gpu:2
#SBATCH --job-name=mt_dao_cen_example
#SBATCH --tasks-per-node=1
#SBATCH --output=%x-%j.out
#SBATCH --time=01:00:00
#SBATCH --dependency=afterany:4871

# Default project path
PROJECT_PATH="$HOME/projects/mt_dao"

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "cen_example.sh: Error parsing options" >&2
	exit 1
fi

eval set -- "$OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	--)
		shift
		break
		;;
	*)
		break
		;;
	esac
done
echo "cen_example.sh: PROJECT_PATH=$PROJECT_PATH"

#! Fix the number of total samples/tokens to train on
BATCH_SIZE=1024
TOTAL_STEPS=5120
# TOTAL_STEPS=10
WARMUP_STEPS=$(uv run python -c "print(int(0.05*$TOTAL_STEPS))")  # 5% of the total steps
COOLDOWN_STEPS=$(uv run python -c "print(int(0.2*$TOTAL_STEPS))") # 20% of the total steps
COOLDOWN_STEPS=0
# EVAL_INTERVAL=$((TOTAL_STEPS / 100))
EVAL_INTERVAL=$TOTAL_STEPS
SAVE_INTERVAL=$TOTAL_STEPS
# Get the current date and time
DATETIME=$(date '+%Y%m%d_%H%M%S')
#! Set the sharding type
# SHARDING="fsdp"
SHARDING="ddp"
export RUN_UUID="test-$SHARDING-$DATETIME"

#! Initialize the external configs
EXTERNAL_CONFIGS=""
EXTERNAL_CONFIGS="llm_config=deptoe"
# EXTERNAL_CONFIGS="llm_config=mpt-16m"

#! Dataset configs
DATASET_NAME="smollm-corpus-shared"
export DATASET_CACHE_DIR="/nfs-share/datasets/mt_dao/dataset_cache"
mkdir -p $DATASET_CACHE_DIR
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS dataset=$DATASET_NAME"                                                      # Dataset name
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS dataset.train.root_local=$DATASET_CACHE_DIR/$DATASET_NAME"                  # Path of the local cache for the training dataset
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS dataset.val.root_local=$DATASET_CACHE_DIR/$DATASET_NAME"                    # Path of the local cache for the evaluation dataset
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS dataset/streams@dataset.train.streams=smollm_corpus_4_clients_proportional" # Stream configuration for the training dataset
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS dataset/streams@dataset.val.streams=smollm_corpus_4_clients_proportional"   # Stream configuration for the training dataset
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS centralized.stream_id=null"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.eval_loader=null"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.tokenizer.name='HuggingFaceTB/SmolLM-1.7B'"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.model.vocab_size=50368"

#! Local optimizer configs
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.max_duration=${TOTAL_STEPS}ba"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.optimizer.name=adopt"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.optimizer.lr=0.0021"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.optimizer.betas=[0.95,0.9999]"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.scheduler.schedulers.lr.name=constant_with_sqrt_cooldown_with_warmup"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.scheduler.schedulers.lr.t_max=${TOTAL_STEPS}ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.scheduler.schedulers.lr.t_warmup=${WARMUP_STEPS}ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ++llm_config.scheduler.schedulers.lr.t_cooldown=${COOLDOWN_STEPS}ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.scheduler.schedulers.lr.alpha_f"

#! Batch size configs (train and eval)
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.global_train_batch_size=$BATCH_SIZE"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.device_train_microbatch_size=auto"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.device_train_microbatch_size=128"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.device_train_microbatch_size=8"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.device_eval_batch_size=$BATCH_SIZE"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ++llm_config.device_eval_microbatch_size=auto"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ++llm_config.device_eval_microbatch_size=1"

#! Precision and attention implementation configs
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.precision=amp_bf16"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.precision=amp_fp16"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.model.attn_config.attn_impl=flash"
#! Selecting the sharding strategy (Default: FSDP)
if [ "$SHARDING" = "ddp" ]; then
	EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.fsdp_config"
	# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ++llm_config.callbacks.noise_scale_monitor={}"
fi

#! Additional configs
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS fl.use_unigram_metrics=true"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.eval_subset_num_batches=1"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ++llm_config.tp_config=null"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.save_interval=${SAVE_INTERVAL}ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.eval_interval=${EVAL_INTERVAL}ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.eval_first=false"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.console_log_interval=10ba"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.loggers.wandb"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.loggers.tensorboard"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.callbacks.activation_monitor_full_model.interval=1ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.callbacks.activation_monitor_full_model"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.callbacks.optimizer_monitor.interval=1ba"
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.callbacks.optimizer_monitor"
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS ~llm_config.callbacks.speed_monitor"

#! Evaluation gauntlet
EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS icl_tasks_config=empty eval_gauntlet_config=empty" # Empty gauntlet

#! Load pre-trained models
#! Note that, in case both are set, the NDArrays checkpoint will be loaded last
# CHECKPOINT_PATH=""
#! Load a model from a checkpoint of type .pt (can reside in the S3 bucket)
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS llm_config.load_path=$CHECKPOINT_PATH"
#! Load a model from a checkpoint of type NDArrays (can reside in the S3 bucket)
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS pretrained_model_path=$CHECKPOINT_PATH"

#! Loading a different weight matrix for the token embedding module of type NDArrays (can reside in the S3 bucket)
#! NOTE: This is crucial to set when analyzing DEPT-SPEC checkpoints since the token embedding matrix is not shared across the clients
# WTE_PARAMETERS_PATH=""
# EXTERNAL_CONFIGS="$EXTERNAL_CONFIGS wte_parameters_path=$WTE_PARAMETERS_PATH"

# External configs will be appended at the end of the command, thus rule over those in the `base_centralised_training.sh` script
export EXTERNAL_CONFIGS

PYTORCH_MASTER_ADDRESS=127.0.0.1 PYTORCH_NODE_RANK=0 PYTORCH_WORLD_SIZE=2 PYTORCH_BASE_RANK=0 bash $PROJECT_PATH/scripts/base_centralised_training.sh -p "$PROJECT_PATH"

# PYTORCH_MASTER_ADDRESS=128.232.115.48 PYTORCH_NODE_RANK=1 PYTORCH_WORLD_SIZE=4 PYTORCH_BASE_RANK=2 PYTORCH_MASTER_PORT=11680 bash $PROJECT_PATH/scripts/base_centralised_training.sh -p "$PROJECT_PATH"
