wehub-resource-sync
50 lines
1.8 KiB
Bash
Executable File
50 lines
1.8 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Turnkey post-training pipeline: SFT -> Reward Model -> {DPO, PPO} -> GRPO -> eval table.
|
|
# Assumes the base model is already pretrained (scripts/pretrain_base.py ->
|
|
# /ephemeral/ckpts/base_pretrained.pt) and the datasets are prepared (scripts/prepare_*).
|
|
#
|
|
# Usage (from repo root):
|
|
# bash scripts/run_posttraining.sh # use both GPUs (torchrun)
|
|
# NPROC=1 bash scripts/run_posttraining.sh # single GPU
|
|
#
|
|
# Each stage writes a checkpoint to /ephemeral/ckpts and metrics JSONL to /ephemeral/logs.
|
|
set -euo pipefail
|
|
|
|
cd "$(dirname "$0")/.."
|
|
export PYTHONPATH=. HF_HOME=/ephemeral/hf_cache PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
|
PY=/ephemeral/venv/bin/python
|
|
NPROC=${NPROC:-2}
|
|
|
|
run() { # run a training script single- or multi-GPU
|
|
if [ "$NPROC" -gt 1 ]; then
|
|
/ephemeral/venv/bin/torchrun --standalone --nproc_per_node="$NPROC" "$@"
|
|
else
|
|
$PY "$@"
|
|
fi
|
|
}
|
|
|
|
echo "############ 1/5 SFT ############"
|
|
run scripts/train_sft.py
|
|
|
|
echo "############ 2/5 Reward Model ############"
|
|
run scripts/train_reward.py
|
|
|
|
echo "############ 3/5 DPO ############"
|
|
run scripts/train_dpo.py --loss_type dpo
|
|
|
|
echo "############ 4/5 PPO (GSM8K, verifier reward) ############"
|
|
run scripts/train_ppo.py --reward_source verifier
|
|
|
|
echo "############ 5/5 GRPO (arithmetic curriculum -> GSM8K) ############"
|
|
run scripts/train_grpo.py
|
|
|
|
echo "############ Eval: GSM8K accuracy across stages ############"
|
|
TABLE=/ephemeral/logs/stage_table.jsonl
|
|
rm -f "$TABLE"
|
|
for s in base_pretrained sft dpo ppo grpo; do
|
|
[ -f "/ephemeral/ckpts/$s.pt" ] && \
|
|
$PY scripts/eval_post_training.py --ckpt "/ephemeral/ckpts/$s.pt" --label "$s" --limit 200 --append "$TABLE"
|
|
done
|
|
$PY scripts/eval_post_training.py --table "$TABLE"
|
|
echo "Done. Metrics in /ephemeral/logs, checkpoints in /ephemeral/ckpts."
|