File size: 4,358 Bytes
6e8abde | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | #!/usr/bin/env bash
# Train EVIE-8B (single 4096-d head). Requires DATA_ROOT, HARDNEG_ROOT.
set -euo pipefail
ROOT="${ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
REPO="${REPO:-$(cd "$ROOT/../.." && pwd)}"
EVIE_ROOT="${EVIE_ROOT:-$REPO}"
# shellcheck source=/dev/null
source "$REPO/code/shared/lib.sh"
evie_resolve_python
evie_pythonpath
evie_workdirs
export PYTHONPATH="$PYTHONPATH:$ROOT/scripts"
cd "$ROOT"
PY="$PYTHON"
RUN_NAME="${RUN_NAME:-evie-8b}"
DATA_ROOT="${DATA_ROOT:?set DATA_ROOT}"
HARDNEG_ROOT="${HARDNEG_ROOT:?set HARDNEG_ROOT}"
BASE_MODEL="${BASE_MODEL:-Qwen/Qwen3.5}"
COL_DIM="${COL_DIM:-4096}"
SOURCES="${SOURCES:-colpali_train_set vdr-multilingual-train VisRAG-Ret-Train-Synthetic-data VisRAG-Ret-Train-In-domain-data tatdqa_train tabfquad_train_set}"
BSZ="${BSZ:-2}"
EFF_BATCH="${EFF_BATCH:-512}"
MVT="${MVT:-1024}"
EPOCHS="${EPOCHS:-1}"
SEED="${SEED:-42}"
BIDIR="${BIDIR:-on}"
GRAD_CHECKPOINTING="${GRAD_CHECKPOINTING:-off}"
LR="${LR:-4.57e-5}"
WEIGHT_DECAY="${WEIGHT_DECAY:-0.02}"
WARMUP_RATIO="${WARMUP_RATIO:-0.08}"
LORA_R="${LORA_R:-32}"
LORA_ALPHA="${LORA_ALPHA:-128}"
LORA_DROPOUT="${LORA_DROPOUT:-0.197}"
LOSS_TEMPERATURE="${LOSS_TEMPERATURE:-0.02}"
NUM_HARD_NEGS="${NUM_HARD_NEGS:-2}"
USE_HARDNEGATIVES="${USE_HARDNEGATIVES:-on}"
HARDNEG_IN_BATCH_WEIGHT="${HARDNEG_IN_BATCH_WEIGHT:-0.5}"
REPORT_TO="${REPORT_TO:-wandb,tensorboard}"
WANDB_PROJECT="${WANDB_PROJECT:-evie}"
WANDB_NAME="${WANDB_NAME:-$RUN_NAME}"
EVAL_AFTER_TRAIN="${EVAL_AFTER_TRAIN:-1}"
SKIP_IF_TRAINED="${SKIP_IF_TRAINED:-1}"
OVERWRITE="${OVERWRITE:-auto}"
TB_DIR="${TB_DIR:-$RUNS_DIR/$RUN_NAME/tensorboard}"
NNODES="${NNODES:-1}"
NODE_RANK="${NODE_RANK:-${RANK:-0}}"
MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
MASTER_PORT="${MASTER_PORT:-29500}"
unset PYTHONHOME
export PATH="$(dirname "$PY"):$PATH"
evie_batch
mkdir -p "$LOG_DIR"
LOG_FILE="$LOG_DIR/train_${RUN_NAME}_node${NODE_RANK}_$(date +%Y%m%d_%H%M%S).log"
exec > >(tee -a "$LOG_FILE") 2>&1
echo "[log] $LOG_FILE"
echo "[dist] ${NNODES}x${NPROC_PER_NODE} rank=$NODE_RANK ${MASTER_ADDR}:${MASTER_PORT}"
OUTPUT_DIR="${OUTPUT_DIR:-$RUNS_DIR/$RUN_NAME}"
WANDB_DIR="${WANDB_DIR:-$OUTPUT_DIR/wandb}"
SKIP_TRAIN=0
if [[ -f "$OUTPUT_DIR/adapter_model.safetensors" && -f "$OUTPUT_DIR/adapter_config.json" \
&& "$SKIP_IF_TRAINED" == "1" && -z "${RESUME:-}" && "${OVERWRITE:-0}" != "1" ]]; then
SKIP_TRAIN=1
echo "[train] skip: adapter already at $OUTPUT_DIR"
fi
if [[ "$SKIP_TRAIN" != "1" ]]; then
evie_nccl
export HARDNEG_SUBDIR="${HARDNEG_SUBDIR:-judged}"
evie_wandb
RESUME="${RESUME:-}"
evie_prepare_output
read -r -a SOURCE_ARGS <<< "$SOURCES"
ARGS=(
--base-model "$BASE_MODEL" --col-dim "$COL_DIM" --data-root "$DATA_ROOT"
--output-dir "$OUTPUT_DIR" --sources "${SOURCE_ARGS[@]}"
--epochs "$EPOCHS" --seed "$SEED"
--per-device-batch-size "$BSZ" --grad-accum "$GRAD_ACCUM"
--learning-rate "$LR" --weight-decay "$WEIGHT_DECAY" --warmup-ratio "$WARMUP_RATIO"
--max-visual-tokens "$MVT" --dataloader-workers "${DL_WORKERS:-8}"
--dataloader-prefetch-factor "${DL_PREFETCH:-4}"
--bidirectional-attention "$BIDIR" --grad-checkpointing "$GRAD_CHECKPOINTING"
--lora-r "$LORA_R" --lora-alpha "$LORA_ALPHA" --lora-dropout "$LORA_DROPOUT"
--loss-temperature "$LOSS_TEMPERATURE"
--num-hard-negs "$NUM_HARD_NEGS" --use-hardnegatives "$USE_HARDNEGATIVES"
--hardneg-in-batch-weight "$HARDNEG_IN_BATCH_WEIGHT"
--report-to "$REPORT_TO" --logging-dir "$TB_DIR" --run-name "$RUN_NAME"
--hardneg-root "$HARDNEG_ROOT"
)
[[ -n "$RESUME" ]] && ARGS+=(--resume-from-checkpoint "$RESUME")
[[ -n "${MAX_STEPS:-}" ]] && ARGS+=(--max-steps "$MAX_STEPS")
[[ -n "${MAX_SAMPLES:-}" ]] && ARGS+=(--max-samples-per-source "$MAX_SAMPLES")
echo "[train] $RUN_NAME seed=$SEED ${NNODES}x${NPROC_PER_NODE} global_batch=$ACTUAL_GLOBAL mvt=$MVT"
"$PY" -m torch.distributed.run \
--nnodes="$NNODES" --nproc_per_node="$NPROC_PER_NODE" --node_rank="$NODE_RANK" \
--master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" scripts/train.py "${ARGS[@]}"
fi
if [[ "$EVAL_AFTER_TRAIN" == "1" ]]; then
export MASTER_PORT="${EVAL_MASTER_PORT:-29501}"
export RUN_NAME NNODES NODE_RANK MASTER_ADDR NPROC_PER_NODE
export EVAL_BASE_MODEL="${EVAL_BASE_MODEL:-$BASE_MODEL}"
bash "$REPO/code/shared/eval_run.sh"
fi
echo "== complete: $RUN_NAME =="
|