EVAFRILL-Mo/train_3b_dpo_1gpu.sh at main · pathcosmos/EVAFRILL-Mo · GitHub

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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env bash
# ==========================================================================
# EVAFRILL-Mo 3B DPO Training — H100 MIG 3g.40gb (Single GPU)
#
# Runs DPO Round 1 + Round 2 sequentially with resilient restart.
#
# Round 1: Full preference data, 3000 steps, beta=0.1, lr=5e-7
# Round 2: High-quality subset, 2000 steps, beta=0.05, lr=1e-7
# ==========================================================================

set -euo pipefail

# Activate Python env if needed
if [ -f /root/ai-env/bin/activate ]; then
    source /root/ai-env/bin/activate
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

SFT_CKPT="checkpoints/3b_sft_v2/checkpoint-best"
DPO_DATA="data/preference/combined_preference.jsonl"
TOKENIZER="tokenizer/korean_sp/tokenizer.json"

# ==========================================
# DPO Round 1
# ==========================================
echo "=========================================="
echo "DPO Round 1: Full data, 3000 steps"
echo "=========================================="

ROUND1_DIR="checkpoints/3b_dpo_r1"
ROUND1_LOG="checkpoints/3b_dpo_r1/train.log"

MAX_RETRIES=5
RETRY_DELAY=30
RETRIES=0

while [ $RETRIES -lt $MAX_RETRIES ]; do
    # Find latest checkpoint for resume
    RESUME_ARG=""
    if [ -d "$ROUND1_DIR" ]; then
        LATEST=$(ls -td "$ROUND1_DIR"/checkpoint-[0-9]* 2>/dev/null | head -1)
        if [ -n "$LATEST" ] && [ -f "$LATEST/model.pt" ]; then
            RESUME_ARG="--resume $LATEST"
            echo "Resuming from: $LATEST"
        fi
    fi

    python3 train/dpo.py \
        --sft_checkpoint "$SFT_CKPT" \
        --dpo_data "$DPO_DATA" \
        --tokenizer "$TOKENIZER" \
        --checkpoint_dir "$ROUND1_DIR" \
        --config configs/h100_mig/dpo_3b_1gpu.yaml \
        --device cuda:0 \
        --log_file "$ROUND1_LOG" \
        $RESUME_ARG \
        && break

    RETRIES=$((RETRIES + 1))
    echo "[WARN] Round 1 crashed (attempt $RETRIES/$MAX_RETRIES). Restarting in ${RETRY_DELAY}s..."
    sleep $RETRY_DELAY
done

if [ $RETRIES -ge $MAX_RETRIES ]; then
    echo "[ERROR] Round 1 failed after $MAX_RETRIES attempts"
    exit 1
fi

echo "Round 1 complete."

# ==========================================
# DPO Round 2 (more conservative)
# ==========================================
echo ""
echo "=========================================="
echo "DPO Round 2: Conservative, 2000 steps"
echo "=========================================="

# Use Round 1 merged checkpoint as base
ROUND1_MERGED="$ROUND1_DIR/checkpoint-merged"
if [ ! -d "$ROUND1_MERGED" ]; then
    echo "[ERROR] Round 1 merged checkpoint not found: $ROUND1_MERGED"
    exit 1
fi

ROUND2_DIR="checkpoints/3b_dpo_r2"
ROUND2_LOG="checkpoints/3b_dpo_r2/train.log"
RETRIES=0

while [ $RETRIES -lt $MAX_RETRIES ]; do
    RESUME_ARG=""
    if [ -d "$ROUND2_DIR" ]; then
        LATEST=$(ls -td "$ROUND2_DIR"/checkpoint-[0-9]* 2>/dev/null | head -1)
        if [ -n "$LATEST" ] && [ -f "$LATEST/model.pt" ]; then
            RESUME_ARG="--resume $LATEST"
            echo "Resuming from: $LATEST"
        fi
    fi

    python3 train/dpo.py \
        --sft_checkpoint "$ROUND1_MERGED" \
        --dpo_data "$DPO_DATA" \
        --tokenizer "$TOKENIZER" \
        --checkpoint_dir "$ROUND2_DIR" \
        --max_steps 2000 \
        --beta 0.05 \
        --lr 1e-7 \
        --warmup_steps 50 \
        --device cuda:0 \
        --log_file "$ROUND2_LOG" \
        $RESUME_ARG \
        && break

    RETRIES=$((RETRIES + 1))
    echo "[WARN] Round 2 crashed (attempt $RETRIES/$MAX_RETRIES). Restarting in ${RETRY_DELAY}s..."
    sleep $RETRY_DELAY
done

if [ $RETRIES -ge $MAX_RETRIES ]; then
    echo "[ERROR] Round 2 failed after $MAX_RETRIES attempts"
    exit 1
fi

echo "Round 2 complete."

# ==========================================
# SLERP Merge: SFT + DPO
# ==========================================
echo ""
echo "=========================================="
echo "SLERP Merge: SFT ↔ DPO"
echo "=========================================="

ROUND2_MERGED="$ROUND2_DIR/checkpoint-merged"
if [ ! -d "$ROUND2_MERGED" ]; then
    echo "[WARN] Round 2 merged checkpoint not found, using Round 1"
    ROUND2_MERGED="$ROUND1_MERGED"
fi

python3 scripts/merge_checkpoints.py \
    --ckpt_a "$SFT_CKPT" \
    --ckpt_b "$ROUND2_MERGED" \
    --output "checkpoints/3b_dpo/checkpoint-slerp" \
    --alpha 0.5

echo ""
echo "=========================================="
echo "DPO Pipeline Complete!"
echo "  Round 1: $ROUND1_DIR"
echo "  Round 2: $ROUND2_DIR"
echo "  SLERP:   checkpoints/3b_dpo/checkpoint-slerp"
echo "=========================================="