127 lines
4.0 KiB
Bash
Executable File
127 lines
4.0 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
cd "${ROOT_DIR}"
|
|
|
|
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
|
|
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
|
|
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
|
|
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
|
|
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
|
|
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
|
|
|
|
if [[ -f .venv/bin/activate ]]; then
|
|
source .venv/bin/activate
|
|
elif [[ "${DRY_RUN:-0}" != "1" ]]; then
|
|
echo "Missing .venv. Run ./scripts/setup_env.sh first." >&2
|
|
exit 2
|
|
fi
|
|
mkdir -p outputs runs logs
|
|
|
|
TRAIN_JSONL="${TRAIN_JSONL:-data/processed/training_probe/train.jsonl}"
|
|
VAL_JSONL="${VAL_JSONL:-data/processed/training_probe/validation.jsonl}"
|
|
MAX_LENGTH="${MAX_LENGTH:-262144}"
|
|
SAVE_STEPS="${SAVE_STEPS:-1000}"
|
|
EVAL_STEPS="${EVAL_STEPS:-1000}"
|
|
LOGGING_STEPS="${LOGGING_STEPS:-1}"
|
|
NUM_EPOCHS="${NUM_EPOCHS:-1}"
|
|
WARMUP_RATIO="${WARMUP_RATIO:-0.1}"
|
|
LR_SCHEDULER_TYPE="${LR_SCHEDULER_TYPE:-cosine}"
|
|
LORA_RANK="${LORA_RANK:-32}"
|
|
DEFAULT_PER_DEVICE_BATCH_SIZE="${DEFAULT_PER_DEVICE_BATCH_SIZE:-1}"
|
|
DEFAULT_GRAD_ACCUM_STEPS="${DEFAULT_GRAD_ACCUM_STEPS:-1}"
|
|
DEFAULT_EVAL_BATCH_SIZE="${DEFAULT_EVAL_BATCH_SIZE:-1}"
|
|
|
|
env_key() {
|
|
printf '%s' "$1" | tr '[:lower:]-' '[:upper:]_' | sed 's/[^A-Z0-9_]/_/g'
|
|
}
|
|
|
|
env_or_default() {
|
|
local name="$1"
|
|
local fallback="$2"
|
|
if [[ -n "${!name:-}" ]]; then
|
|
printf '%s' "${!name}"
|
|
else
|
|
printf '%s' "${fallback}"
|
|
fi
|
|
}
|
|
|
|
require_file() {
|
|
if [[ ! -f "$1" ]]; then
|
|
echo "Missing required file: $1" >&2
|
|
exit 2
|
|
fi
|
|
}
|
|
|
|
run_swift_train() {
|
|
local model_path="$1"
|
|
local train_type="$2"
|
|
local run_name="$3"
|
|
local output_dir="outputs/${run_name}"
|
|
local tb_dir="runs/${run_name}"
|
|
local log_file="logs/${run_name}.log"
|
|
local run_key
|
|
run_key="$(env_key "${run_name}")"
|
|
|
|
local default_lr
|
|
if [[ "${train_type}" == "lora" ]]; then
|
|
default_lr="${LORA_LEARNING_RATE:-5e-5}"
|
|
else
|
|
default_lr="${FULL_LEARNING_RATE:-1e-5}"
|
|
fi
|
|
local learning_rate
|
|
learning_rate="$(env_or_default "${run_key}_LEARNING_RATE" "${LEARNING_RATE:-${default_lr}}")"
|
|
|
|
local type_key
|
|
type_key="$(env_key "${train_type}")"
|
|
local type_bsz_var="${type_key}_PER_DEVICE_BATCH_SIZE"
|
|
local type_accum_var="${type_key}_GRAD_ACCUM_STEPS"
|
|
local per_device_batch_size
|
|
local grad_accum_steps
|
|
local eval_batch_size
|
|
per_device_batch_size="$(env_or_default "${run_key}_PER_DEVICE_BATCH_SIZE" "${PER_DEVICE_BATCH_SIZE:-${!type_bsz_var:-${DEFAULT_PER_DEVICE_BATCH_SIZE}}}")"
|
|
grad_accum_steps="$(env_or_default "${run_key}_GRAD_ACCUM_STEPS" "${GRAD_ACCUM_STEPS:-${!type_accum_var:-${DEFAULT_GRAD_ACCUM_STEPS}}}")"
|
|
eval_batch_size="$(env_or_default "${run_key}_EVAL_BATCH_SIZE" "${EVAL_PER_DEVICE_BATCH_SIZE:-${DEFAULT_EVAL_BATCH_SIZE}}")"
|
|
|
|
require_file "${TRAIN_JSONL}"
|
|
require_file "${VAL_JSONL}"
|
|
mkdir -p "${output_dir}" "${tb_dir}" logs
|
|
|
|
local cmd=(
|
|
swift sft
|
|
--model "${model_path}"
|
|
--dataset "${TRAIN_JSONL}"
|
|
--val_dataset "${VAL_JSONL}"
|
|
--train_type "${train_type}"
|
|
--torch_dtype bfloat16
|
|
--num_train_epochs "${NUM_EPOCHS}"
|
|
--per_device_train_batch_size "${per_device_batch_size}"
|
|
--per_device_eval_batch_size "${eval_batch_size}"
|
|
--gradient_accumulation_steps "${grad_accum_steps}"
|
|
--learning_rate "${learning_rate}"
|
|
--warmup_ratio "${WARMUP_RATIO}"
|
|
--lr_scheduler_type "${LR_SCHEDULER_TYPE}"
|
|
--max_length "${MAX_LENGTH}"
|
|
--save_steps "${SAVE_STEPS}"
|
|
--eval_steps "${EVAL_STEPS}"
|
|
--logging_steps "${LOGGING_STEPS}"
|
|
--report_to tensorboard
|
|
--logging_dir "${tb_dir}"
|
|
--output_dir "${output_dir}"
|
|
--save_total_limit "${SAVE_TOTAL_LIMIT:-3}"
|
|
--dataloader_num_workers "${DATALOADER_NUM_WORKERS:-4}"
|
|
)
|
|
|
|
if [[ "${train_type}" == "lora" ]]; then
|
|
cmd+=(--lora_rank "${LORA_RANK}")
|
|
fi
|
|
|
|
printf '%q ' "${cmd[@]}" | tee "${log_file}.cmd"
|
|
echo
|
|
if [[ "${DRY_RUN:-0}" == "1" ]]; then
|
|
return 0
|
|
fi
|
|
"${cmd[@]}" 2>&1 | tee "${log_file}"
|
|
}
|